|
| 1 | +/* |
| 2 | + * Licensed to the Apache Software Foundation (ASF) under one |
| 3 | + * or more contributor license agreements. See the NOTICE file |
| 4 | + * distributed with this work for additional information |
| 5 | + * regarding copyright ownership. The ASF licenses this file |
| 6 | + * to you under the Apache License, Version 2.0 (the |
| 7 | + * "License"); you may not use this file except in compliance |
| 8 | + * with the License. You may obtain a copy of the License at |
| 9 | + * |
| 10 | + * http://www.apache.org/licenses/LICENSE-2.0 |
| 11 | + * |
| 12 | + * Unless required by applicable law or agreed to in writing, |
| 13 | + * software distributed under the License is distributed on an |
| 14 | + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| 15 | + * KIND, either express or implied. See the License for the |
| 16 | + * specific language governing permissions and limitations |
| 17 | + * under the License. |
| 18 | + */ |
| 19 | + |
| 20 | +#include "iceberg/arrow_c_data_internal.h" |
| 21 | + |
| 22 | +#include <sparrow/builder/builder.hpp> |
| 23 | + |
| 24 | +namespace iceberg::internal { |
| 25 | + |
| 26 | +std::pair<ArrowSchema, ArrowArray> CreateExampleArrowSchemaAndArrayByNanoarrow() { |
| 27 | + ArrowSchema out_schema; |
| 28 | + |
| 29 | + // Initializes the root struct schema |
| 30 | + NANOARROW_THROW_NOT_OK(ArrowSchemaInitFromType(&out_schema, NANOARROW_TYPE_STRUCT)); |
| 31 | + NANOARROW_THROW_NOT_OK(ArrowSchemaAllocateChildren(&out_schema, 2)); |
| 32 | + |
| 33 | + // Set up the non-nullable int64 field |
| 34 | + struct ArrowSchema* int64_field = out_schema.children[0]; |
| 35 | + ArrowSchemaInit(int64_field); |
| 36 | + NANOARROW_THROW_NOT_OK(ArrowSchemaInitFromType(int64_field, NANOARROW_TYPE_INT64)); |
| 37 | + NANOARROW_THROW_NOT_OK(ArrowSchemaSetName(int64_field, "id")); |
| 38 | + int64_field->flags &= ~ARROW_FLAG_NULLABLE; |
| 39 | + |
| 40 | + // Set up the nullable string field |
| 41 | + struct ArrowSchema* string_field = out_schema.children[1]; |
| 42 | + ArrowSchemaInit(string_field); |
| 43 | + NANOARROW_THROW_NOT_OK(ArrowSchemaInitFromType(string_field, NANOARROW_TYPE_STRING)); |
| 44 | + NANOARROW_THROW_NOT_OK(ArrowSchemaSetName(string_field, "name")); |
| 45 | + string_field->flags |= ARROW_FLAG_NULLABLE; |
| 46 | + |
| 47 | + constexpr int64_t kNumValues = 3; |
| 48 | + std::array<int64_t, kNumValues> int64_values = {1, 2, 3}; |
| 49 | + std::array<std::string, kNumValues> string_values = {"a", "b", "c"}; |
| 50 | + |
| 51 | + ArrowArray out_array; |
| 52 | + NANOARROW_THROW_NOT_OK(ArrowArrayInitFromSchema(&out_array, &out_schema, nullptr)); |
| 53 | + ArrowArray* int64_array = out_array.children[0]; |
| 54 | + ArrowArray* string_array = out_array.children[1]; |
| 55 | + |
| 56 | + NANOARROW_THROW_NOT_OK(ArrowArrayStartAppending(int64_array)); |
| 57 | + NANOARROW_THROW_NOT_OK(ArrowArrayStartAppending(string_array)); |
| 58 | + |
| 59 | + for (int64_t i = 0; i < kNumValues; i++) { |
| 60 | + NANOARROW_THROW_NOT_OK(ArrowArrayAppendInt(int64_array, int64_values[i])); |
| 61 | + NANOARROW_THROW_NOT_OK( |
| 62 | + ArrowArrayAppendString(string_array, ArrowCharView(string_values[i].c_str()))); |
| 63 | + } |
| 64 | + |
| 65 | + NANOARROW_THROW_NOT_OK(ArrowArrayFinishBuildingDefault(int64_array, nullptr)); |
| 66 | + NANOARROW_THROW_NOT_OK(ArrowArrayFinishBuildingDefault(string_array, nullptr)); |
| 67 | + |
| 68 | + out_array.length = kNumValues; |
| 69 | + out_array.null_count = 0; |
| 70 | + |
| 71 | + return {out_schema, out_array}; |
| 72 | +} |
| 73 | + |
| 74 | +std::pair<ArrowSchema, ArrowArray> CreateExampleArrowSchemaAndArrayBySparrow() { |
| 75 | + using struct_type = std::tuple<int64_t, sparrow::nullable<std::string>>; |
| 76 | + std::vector<struct_type> values = { |
| 77 | + {1, "a"}, |
| 78 | + {2, "b"}, |
| 79 | + {3, "c"}, |
| 80 | + }; |
| 81 | + auto sparrow_array = sparrow::build(values); |
| 82 | + |
| 83 | + // Demonstrate the use of arrow_proxy to modify the schema |
| 84 | + auto [_, non_owning_schema] = sparrow::get_arrow_structures(sparrow_array); |
| 85 | + non_owning_schema->children[0]->name = "id"; |
| 86 | + non_owning_schema->children[0]->flags &= ~ARROW_FLAG_NULLABLE; |
| 87 | + non_owning_schema->children[1]->name = "name"; |
| 88 | + non_owning_schema->children[1]->flags |= ARROW_FLAG_NULLABLE; |
| 89 | + |
| 90 | + auto [array, schema] = sparrow::extract_arrow_structures(std::move(sparrow_array)); |
| 91 | + return {schema, array}; |
| 92 | +} |
| 93 | + |
| 94 | +} // namespace iceberg::internal |
0 commit comments