|
39 | 39 | # include <unistd.h> // IWYU pragma: keep |
40 | 40 | #endif |
41 | 41 |
|
| 42 | +#include "arrow/array.h" |
| 43 | +#include "arrow/array/builder_union.h" |
42 | 44 | #include "arrow/config.h" |
43 | 45 | #include "arrow/table.h" |
| 46 | +#include "arrow/testing/builder.h" |
44 | 47 | #include "arrow/testing/random.h" |
45 | 48 | #include "arrow/type.h" |
46 | 49 | #include "arrow/util/cpu_info.h" |
@@ -242,4 +245,55 @@ std::vector<int64_t> GetSupportedHardwareFlags( |
242 | 245 | return hardware_flags; |
243 | 246 | } |
244 | 247 |
|
| 248 | +Status MakeUnion(std::shared_ptr<RecordBatch>* out) { |
| 249 | + // Define schema |
| 250 | + std::vector<std::shared_ptr<Field>> union_fields( |
| 251 | + {field("u0", int32()), field("u1", uint8())}); |
| 252 | + |
| 253 | + std::vector<int8_t> type_codes = {5, 10}; |
| 254 | + auto sparse_type = sparse_union(union_fields, type_codes); |
| 255 | + auto dense_type = dense_union(union_fields, type_codes); |
| 256 | + |
| 257 | + auto f0 = field("sparse", sparse_type); |
| 258 | + auto f1 = field("dense", dense_type); |
| 259 | + auto schema = ::arrow::schema({f0, f1}); |
| 260 | + |
| 261 | + // Create data |
| 262 | + std::vector<std::shared_ptr<Array>> sparse_children(2); |
| 263 | + std::vector<std::shared_ptr<Array>> dense_children(2); |
| 264 | + |
| 265 | + const int64_t length = 7; |
| 266 | + |
| 267 | + std::shared_ptr<Buffer> type_ids_buffer; |
| 268 | + std::vector<uint8_t> type_ids = {5, 10, 5, 5, 10, 10, 5}; |
| 269 | + RETURN_NOT_OK(CopyBufferFromVector(type_ids, default_memory_pool(), &type_ids_buffer)); |
| 270 | + |
| 271 | + std::vector<int32_t> u0_values = {0, 1, 2, 3, 4, 5, 6}; |
| 272 | + ArrayFromVector<Int32Type, int32_t>(u0_values, &sparse_children[0]); |
| 273 | + |
| 274 | + std::vector<uint8_t> u1_values = {10, 11, 12, 13, 14, 15, 16}; |
| 275 | + ArrayFromVector<UInt8Type, uint8_t>(u1_values, &sparse_children[1]); |
| 276 | + |
| 277 | + // dense children |
| 278 | + u0_values = {0, 2, 3, 7}; |
| 279 | + ArrayFromVector<Int32Type, int32_t>(u0_values, &dense_children[0]); |
| 280 | + |
| 281 | + u1_values = {11, 14, 15}; |
| 282 | + ArrayFromVector<UInt8Type, uint8_t>(u1_values, &dense_children[1]); |
| 283 | + |
| 284 | + std::shared_ptr<Buffer> offsets_buffer; |
| 285 | + std::vector<int32_t> offsets = {0, 0, 1, 2, 1, 2, 3}; |
| 286 | + RETURN_NOT_OK(CopyBufferFromVector(offsets, default_memory_pool(), &offsets_buffer)); |
| 287 | + |
| 288 | + auto sparse = std::make_shared<SparseUnionArray>(sparse_type, length, sparse_children, |
| 289 | + type_ids_buffer); |
| 290 | + auto dense = std::make_shared<DenseUnionArray>(dense_type, length, dense_children, |
| 291 | + type_ids_buffer, offsets_buffer); |
| 292 | + |
| 293 | + // construct batch |
| 294 | + std::vector<std::shared_ptr<Array>> arrays = {sparse, dense}; |
| 295 | + *out = RecordBatch::Make(schema, length, arrays); |
| 296 | + return Status::OK(); |
| 297 | +} |
| 298 | + |
245 | 299 | } // namespace arrow |
0 commit comments