1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
|
#pragma clang system_header
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
#pragma once
#include "contrib/libs/apache/arrow_next/cpp/src/arrow/record_batch.h"
#include "contrib/libs/apache/arrow_next/cpp/src/arrow/result.h"
#include "contrib/libs/apache/arrow_next/cpp/src/arrow/status.h"
#include "contrib/libs/apache/arrow_next/cpp/src/arrow/table_builder.h"
#include "contrib/libs/apache/arrow_next/cpp/src/arrow/util/iterator.h"
#include <type_traits>
namespace arrow20::util {
namespace detail {
// Default identity function row accessor. Used to for the common case where the value
// of each row iterated over is it's self also directly iterable.
[[nodiscard]] constexpr inline auto MakeDefaultRowAccessor() {
return [](auto& x) -> Result<decltype(std::ref(x))> { return std::ref(x); };
}
// Meta-function to check if a type `T` is a range (iterable using `std::begin()` /
// `std::end()`). `is_range<T>::value` will be false if `T` is not a valid range.
template <typename T, typename = void>
struct is_range : std::false_type {};
template <typename T>
struct is_range<T, std::void_t<decltype(std::begin(std::declval<T>())),
decltype(std::end(std::declval<T>()))>> : std::true_type {
};
} // namespace detail
/// Delete overload for `const Range&& rows` because the data's lifetime must exceed
/// the lifetime of the function call. `data` will be read when client uses the
/// `RecordBatchReader`
template <class Range, class DataPointConvertor,
class RowAccessor = decltype(detail::MakeDefaultRowAccessor())>
[[nodiscard]] typename std::enable_if_t<detail::is_range<Range>::value,
Result<std::shared_ptr<RecordBatchReader>>>
/* Result<std::shared_ptr<RecordBatchReader>>> */ RowsToBatches(
const std::shared_ptr<Schema>& schema, const Range&& rows,
DataPointConvertor&& data_point_convertor,
RowAccessor&& row_accessor = detail::MakeDefaultRowAccessor(),
MemoryPool* pool = default_memory_pool(),
const std::size_t batch_size = 1024) = delete;
/// \brief Utility function for converting any row-based structure into an
/// `arrow20::RecordBatchReader` (this can be easily converted to an `arrow20::Table` using
/// `arrow20::RecordBatchReader::ToTable()`).
///
/// Examples of supported types:
/// - `std::vector<std::vector<std::variant<int, bsl::string>>>`
/// - `std::vector<MyRowStruct>`
/// If `rows` (client’s row-based structure) is not a valid C++ range, the client will
/// need to either make it iterable, or make an adapter/wrapper that is a valid C++
/// range.
/// The client must provide a `DataPointConvertor` callable type that will convert the
/// structure’s data points into the corresponding arrow types.
/// Complex nested rows can be supported by providing a custom `row_accessor` instead
/// of the default.
/// Example usage:
/// \code{.cpp}
/// auto IntConvertor = [](ArrayBuilder& array_builder, int value) {
/// return static_cast<Int64Builder&>(array_builder).Append(value);
/// };
/// std::vector<std::vector<int>> data = {{1, 2, 4}, {5, 6, 7}};
/// auto batches = RowsToBatches(kTestSchema, data, IntConvertor);
/// \endcode
/// \param[in] schema - The schema to be used in the `RecordBatchReader`
/// \param[in] rows - Iterable row-based structure that will be converted to arrow
/// batches
/// \param[in] data_point_convertor - Client provided callable type that will convert
/// the structure’s data points into the corresponding arrow types. The convertor must
/// return an error `Status` if an error happens during conversion.
/// \param[in] row_accessor - In the common case where the value of each row iterated
/// over is it's self also directly iterable, the client can just use the default.
/// The provided callable must take the values of the `rows` range and return a
/// `std::reference_wrapper<Range>` to the data points in a given row. The data points
/// must be in order of their corresponding fields in the schema.
/// see: /ref `MakeDefaultRowAccessor`
/// \param[in] pool - The MemoryPool to use for allocations.
/// \param[in] batch_size - Number of rows to insert into each RecordBatch.
/// \return `Result<std::shared_ptr<RecordBatchReader>>>` result will be a
/// `std::shared_ptr<RecordBatchReader>>` if not errors occurred, else an error status.
template <class Range, class DataPointConvertor,
class RowAccessor = decltype(detail::MakeDefaultRowAccessor())>
[[nodiscard]] typename std::enable_if_t<detail::is_range<Range>::value,
Result<std::shared_ptr<RecordBatchReader>>>
/* Result<std::shared_ptr<RecordBatchReader>>> */ RowsToBatches(
const std::shared_ptr<Schema>& schema, const Range& rows,
DataPointConvertor&& data_point_convertor,
RowAccessor&& row_accessor = detail::MakeDefaultRowAccessor(),
MemoryPool* pool = default_memory_pool(), const std::size_t batch_size = 1024) {
auto make_next_batch =
[pool = pool, batch_size = batch_size, rows_ittr = std::begin(rows),
rows_ittr_end = std::end(rows), schema = schema,
row_accessor = std::forward<RowAccessor>(row_accessor),
data_point_convertor = std::forward<DataPointConvertor>(
data_point_convertor)]() mutable -> Result<std::shared_ptr<RecordBatch>> {
if (rows_ittr == rows_ittr_end) return NULLPTR;
ARROW_ASSIGN_OR_RAISE(auto record_batch_builder,
RecordBatchBuilder::Make(schema, pool, batch_size));
for (size_t i = 0; i < batch_size && (rows_ittr != rows_ittr_end);
i++, std::advance(rows_ittr, 1)) {
int col_index = 0;
ARROW_ASSIGN_OR_RAISE(const auto row, row_accessor(*rows_ittr));
// If the accessor returns a `std::reference_wrapper` unwrap if
const auto& row_unwrapped = [&]() {
if constexpr (detail::is_range<decltype(row)>::value)
return row;
else
return row.get();
}();
for (auto& data_point : row_unwrapped) {
ArrayBuilder* array_builder = record_batch_builder->GetField(col_index);
ARROW_RETURN_IF(array_builder == NULLPTR,
Status::Invalid("array_builder == NULLPTR"));
ARROW_RETURN_NOT_OK(data_point_convertor(*array_builder, data_point));
col_index++;
}
}
ARROW_ASSIGN_OR_RAISE(auto result, record_batch_builder->Flush());
return result;
};
return RecordBatchReader::MakeFromIterator(MakeFunctionIterator(make_next_batch),
schema);
}
} // namespace arrow20::util
|
