| 37 | TestUnionNode() : rng_(0) {} |
| 38 | |
| 39 | std::shared_ptr<Schema> GenerateRandomSchema(size_t num_inputs) { |
| 40 | static std::vector<std::shared_ptr<DataType>> some_arrow_types = { |
| 41 | arrow::null(), arrow::boolean(), arrow::int8(), arrow::int16(), |
| 42 | arrow::int32(), arrow::int64(), arrow::float16(), arrow::float32(), |
| 43 | arrow::float64(), arrow::utf8(), arrow::binary(), arrow::date32()}; |
| 44 | |
| 45 | std::vector<std::shared_ptr<Field>> fields(num_inputs); |
| 46 | std::default_random_engine gen(42); |
| 47 | std::uniform_int_distribution<int> types_dist( |
| 48 | 0, static_cast<int>(some_arrow_types.size()) - 1); |
| 49 | for (size_t i = 0; i < num_inputs; i++) { |
| 50 | int random_index = types_dist(gen); |
| 51 | auto col_type = some_arrow_types.at(random_index); |
| 52 | fields[i] = |
| 53 | field("column_" + std::to_string(i) + "_" + col_type->ToString(), col_type); |
| 54 | } |
| 55 | return schema(fields); |
| 56 | } |
| 57 | |
| 58 | void GenerateBatchesFromSchema(const std::shared_ptr<Schema>& schema, |
| 59 | size_t num_batches, BatchesWithSchema* out_batches, |