MCPcopy Create free account
hub / github.com/apache/arrow / MakeIntegerBatch

Function MakeIntegerBatch

cpp/src/arrow/acero/test_util_internal.cc:172–214  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

170namespace {
171
172Result<ExecBatch> MakeIntegerBatch(const std::vector<std::function<int64_t(int)>>& gens,
173 const std::shared_ptr<Schema>& schema,
174 int batch_start_row, int batch_size) {
175 int n_fields = schema->num_fields();
176 if (gens.size() != static_cast<size_t>(n_fields)) {
177 return Status::Invalid("mismatching generator-vector and schema size");
178 }
179 auto memory_pool = default_memory_pool();
180 std::vector<Datum> values(n_fields);
181 for (int f = 0; f < n_fields; f++) {
182 std::shared_ptr<Array> array;
183 auto type = schema->field(f)->type();
184
185#define ARROW_TEST_INT_BUILD_CASE(id) \
186 case Type::id: { \
187 using T = typename TypeIdTraits<Type::id>::Type; \
188 using CType = typename TypeTraits<T>::CType; \
189 using Builder = typename TypeTraits<T>::BuilderType; \
190 ARROW_ASSIGN_OR_RAISE(auto a_builder, MakeBuilder(type, memory_pool)); \
191 Builder& builder = *checked_cast<Builder*>(a_builder.get()); \
192 ARROW_RETURN_NOT_OK(builder.Reserve(batch_size)); \
193 for (int j = 0; j < batch_size; j++) { \
194 builder.UnsafeAppend(static_cast<CType>(gens[f](batch_start_row + j))); \
195 } \
196 ARROW_RETURN_NOT_OK(builder.Finish(&array)); \
197 break; \
198 }
199
200 switch (type->id()) {
201 ARROW_TEST_INT_BUILD_CASE(INT8)
202 ARROW_TEST_INT_BUILD_CASE(INT16)
203 ARROW_TEST_INT_BUILD_CASE(INT32)
204 ARROW_TEST_INT_BUILD_CASE(INT64)
205 default:
206 return Status::TypeError("building ", type->ToString());
207 }
208
209#undef ARROW_TEST_INT_BUILD_CASE
210
211 values[f] = Datum(array);
212 }
213 return ExecBatch(std::move(values), batch_size);
214}
215
216} // namespace
217

Callers 1

NextMethod · 0.85

Calls 11

default_memory_poolFunction · 0.85
InvalidFunction · 0.50
TypeErrorFunction · 0.50
DatumClass · 0.50
ExecBatchFunction · 0.50
num_fieldsMethod · 0.45
sizeMethod · 0.45
typeMethod · 0.45
fieldMethod · 0.45
idMethod · 0.45
ToStringMethod · 0.45

Tested by

no test coverage detected