| /* | |
| * 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. | |
| */ | |
| #include <functional> | |
| #include <memory> | |
| #include <mutex> | |
| #include <optional> | |
| #include <string> | |
| #include <string_view> | |
| #include <thread> | |
| #include <unordered_map> | |
| #include <utility> | |
| #include <vector> | |
| #include <arrow/api.h> | |
| #include <arrow/c/bridge.h> | |
| #include <arrow/extension/uuid.h> | |
| #include <arrow/result.h> | |
| #include <arrow/type_fwd.h> | |
| #include <arrow/util/config.h> | |
| #include <arrow/util/key_value_metadata.h> | |
| #include <arrow/util/thread_pool.h> | |
| #include <gtest/gtest.h> | |
| #include "iceberg/arrow/arrow_status_internal.h" | |
| #include "iceberg/arrow/metadata_column_util_internal.h" | |
| #include "iceberg/constants.h" | |
| #include "iceberg/result.h" | |
| #include "iceberg/schema.h" | |
| #include "iceberg/schema_internal.h" | |
| #include "iceberg/test/matchers.h" | |
| #include "iceberg/util/executor.h" | |
| #include "iceberg/util/task_group.h" | |
| namespace iceberg { | |
| TEST(MetadataColumnUtilTest, RowIdsWithoutFirstRowId) { | |
| ICEBERG_UNWRAP_OR_FAIL( | |
| auto row_ids, | |
| arrow::MakeRowIdArray(std::nullopt, nullptr, 3, ::arrow::default_memory_pool())); | |
| ASSERT_TRUE(row_ids->ValidateFull().ok()); | |
| EXPECT_EQ(row_ids->length(), 3); | |
| EXPECT_EQ(row_ids->null_count(), row_ids->length()); | |
| } | |
| struct ToArrowSchemaParam { | |
| std::shared_ptr<Type> iceberg_type; | |
| bool optional = true; | |
| std::shared_ptr<::arrow::DataType> arrow_type; | |
| }; | |
| class ToArrowSchemaTest : public ::testing::TestWithParam<ToArrowSchemaParam> {}; | |
| TEST_P(ToArrowSchemaTest, PrimitiveType) { | |
| constexpr std::string_view kFieldName = "foo"; | |
| constexpr int32_t kFieldId = 1024; | |
| const auto& param = GetParam(); | |
| Schema schema( | |
| {param.optional ? SchemaField::MakeOptional(kFieldId, std::string(kFieldName), | |
| param.iceberg_type) | |
| : SchemaField::MakeRequired(kFieldId, std::string(kFieldName), | |
| param.iceberg_type)}, | |
| /*schema_id=*/0); | |
| ArrowSchema arrow_schema; | |
| ASSERT_THAT(ToArrowSchema(schema, &arrow_schema), IsOk()); | |
| auto imported_schema = ::arrow::ImportSchema(&arrow_schema).ValueOrDie(); | |
| ASSERT_EQ(imported_schema->num_fields(), 1); | |
| auto field = imported_schema->field(0); | |
| ASSERT_EQ(field->name(), kFieldName); | |
| ASSERT_EQ(field->nullable(), param.optional); | |
| ASSERT_TRUE(field->type()->Equals(param.arrow_type)); | |
| auto metadata = field->metadata(); | |
| ASSERT_TRUE(metadata->Contains(kParquetFieldIdKey)); | |
| ASSERT_EQ(metadata->Get(kParquetFieldIdKey), std::to_string(kFieldId)); | |
| } | |
| INSTANTIATE_TEST_SUITE_P( | |
| SchemaConversion, ToArrowSchemaTest, | |
| ::testing::Values( | |
| ToArrowSchemaParam{.iceberg_type = iceberg::boolean(), | |
| .optional = false, | |
| .arrow_type = ::arrow::boolean()}, | |
| ToArrowSchemaParam{.iceberg_type = iceberg::int32(), | |
| .optional = true, | |
| .arrow_type = ::arrow::int32()}, | |
| ToArrowSchemaParam{.iceberg_type = iceberg::int64(), | |
| .arrow_type = ::arrow::int64()}, | |
| ToArrowSchemaParam{.iceberg_type = iceberg::float32(), | |
| .arrow_type = ::arrow::float32()}, | |
| ToArrowSchemaParam{.iceberg_type = iceberg::float64(), | |
| .arrow_type = ::arrow::float64()}, | |
| ToArrowSchemaParam{.iceberg_type = iceberg::decimal(10, 2), | |
| .arrow_type = ::arrow::decimal128(10, 2)}, | |
| ToArrowSchemaParam{.iceberg_type = iceberg::date(), | |
| .arrow_type = ::arrow::date32()}, | |
| ToArrowSchemaParam{.iceberg_type = iceberg::time(), | |
| .arrow_type = ::arrow::time64(::arrow::TimeUnit::MICRO)}, | |
| ToArrowSchemaParam{.iceberg_type = iceberg::timestamp(), | |
| .arrow_type = ::arrow::timestamp(::arrow::TimeUnit::MICRO)}, | |
| ToArrowSchemaParam{ | |
| .iceberg_type = iceberg::timestamp_tz(), | |
| .arrow_type = ::arrow::timestamp(::arrow::TimeUnit::MICRO, "UTC")}, | |
| ToArrowSchemaParam{.iceberg_type = iceberg::timestamp_ns(), | |
| .arrow_type = ::arrow::timestamp(::arrow::TimeUnit::NANO)}, | |
| ToArrowSchemaParam{ | |
| .iceberg_type = iceberg::timestamptz_ns(), | |
| .arrow_type = ::arrow::timestamp(::arrow::TimeUnit::NANO, "UTC")}, | |
| ToArrowSchemaParam{.iceberg_type = iceberg::string(), | |
| .arrow_type = ::arrow::utf8()}, | |
| ToArrowSchemaParam{.iceberg_type = iceberg::binary(), | |
| .arrow_type = ::arrow::binary()}, | |
| ToArrowSchemaParam{.iceberg_type = iceberg::geometry(), | |
| .arrow_type = ::arrow::binary()}, | |
| ToArrowSchemaParam{.iceberg_type = iceberg::geography(), | |
| .arrow_type = ::arrow::binary()}, | |
| ToArrowSchemaParam{.iceberg_type = iceberg::uuid(), | |
| .arrow_type = ::arrow::extension::uuid()}, | |
| ToArrowSchemaParam{.iceberg_type = iceberg::fixed(20), | |
| .arrow_type = ::arrow::fixed_size_binary(20)}, | |
| ToArrowSchemaParam{.iceberg_type = iceberg::unknown(), | |
| .arrow_type = ::arrow::null()})); | |
| TEST(ToArrowSchemaTest, UnsupportedV3Types) { | |
| const std::vector<std::shared_ptr<Type>> unsupported_types = {iceberg::variant()}; | |
| for (const auto& unsupported_type : unsupported_types) { | |
| Schema schema( | |
| {SchemaField::MakeOptional(/*field_id=*/1, "unsupported", unsupported_type)}, | |
| /*schema_id=*/0); | |
| ArrowSchema arrow_schema{}; | |
| ASSERT_THAT(ToArrowSchema(schema, &arrow_schema), | |
| HasErrorMessage("is not supported by Arrow conversion")); | |
| EXPECT_EQ(arrow_schema.release, nullptr); | |
| } | |
| } | |
| TEST(ToArrowSchemaTest, ReleasesSchemaOnConversionFailure) { | |
| // fixed(0) passes CheckArrowCompatible but nanoarrow rejects its non-positive width. | |
| Schema schema({SchemaField::MakeOptional(/*field_id=*/1, "invalid_fixed", fixed(0))}, | |
| /*schema_id=*/0); | |
| ArrowSchema arrow_schema{}; | |
| EXPECT_THAT(ToArrowSchema(schema, &arrow_schema), IsError(ErrorKind::kInvalidSchema)); | |
| EXPECT_EQ(arrow_schema.release, nullptr); | |
| } | |
| namespace { | |
| void CheckArrowField(const ::arrow::Field& field, ::arrow::Type::type type_id, | |
| std::string_view name, bool nullable, int32_t field_id) { | |
| ASSERT_EQ(field.name(), name); | |
| ASSERT_EQ(field.nullable(), nullable); | |
| ASSERT_EQ(field.type()->id(), type_id); | |
| auto metadata = field.metadata(); | |
| ASSERT_TRUE(metadata != nullptr); | |
| ASSERT_TRUE(metadata->Contains(kParquetFieldIdKey)); | |
| ASSERT_EQ(metadata->Get(kParquetFieldIdKey), std::to_string(field_id)); | |
| } | |
| } // namespace | |
| TEST(ToArrowSchemaTest, StructType) { | |
| constexpr int32_t kStructFieldId = 1; | |
| constexpr int32_t kIntFieldId = 2; | |
| constexpr int32_t kStrFieldId = 3; | |
| constexpr std::string_view kStructFieldName = "struct_field"; | |
| constexpr std::string_view kIntFieldName = "int_field"; | |
| constexpr std::string_view kStrFieldName = "str_field"; | |
| auto struct_type = std::make_shared<StructType>(std::vector<SchemaField>{ | |
| SchemaField::MakeRequired(kIntFieldId, std::string(kIntFieldName), | |
| iceberg::int32()), | |
| SchemaField::MakeOptional(kStrFieldId, std::string(kStrFieldName), | |
| iceberg::string())}); | |
| Schema schema({SchemaField::MakeRequired(kStructFieldId, std::string(kStructFieldName), | |
| struct_type)}, | |
| /*schema_id=*/0); | |
| ArrowSchema arrow_schema; | |
| ASSERT_THAT(ToArrowSchema(schema, &arrow_schema), IsOk()); | |
| auto imported_schema = ::arrow::ImportSchema(&arrow_schema).ValueOrDie(); | |
| ASSERT_EQ(imported_schema->num_fields(), 1); | |
| auto field = imported_schema->field(0); | |
| ASSERT_NO_FATAL_FAILURE(CheckArrowField(*field, ::arrow::Type::STRUCT, kStructFieldName, | |
| /*nullable=*/false, kStructFieldId)); | |
| auto struct_field = std::static_pointer_cast<::arrow::StructType>(field->type()); | |
| ASSERT_EQ(struct_field->num_fields(), 2); | |
| ASSERT_NO_FATAL_FAILURE(CheckArrowField(*struct_field->field(0), ::arrow::Type::INT32, | |
| kIntFieldName, /*nullable=*/false, | |
| kIntFieldId)); | |
| ASSERT_NO_FATAL_FAILURE(CheckArrowField(*struct_field->field(1), ::arrow::Type::STRING, | |
| kStrFieldName, /*nullable=*/true, kStrFieldId)); | |
| } | |
| TEST(ToArrowSchemaTest, ListType) { | |
| constexpr std::string_view kListFieldName = "list_field"; | |
| constexpr std::string_view kElemFieldName = "element"; | |
| constexpr int32_t kListFieldId = 1; | |
| constexpr int32_t kElemFieldId = 2; | |
| auto list_type = std::make_shared<ListType>(SchemaField::MakeOptional( | |
| kElemFieldId, std::string(kElemFieldName), iceberg::int64())); | |
| Schema schema( | |
| {SchemaField::MakeRequired(kListFieldId, std::string(kListFieldName), list_type)}, | |
| /*schema_id=*/0); | |
| ArrowSchema arrow_schema; | |
| ASSERT_THAT(ToArrowSchema(schema, &arrow_schema), IsOk()); | |
| auto imported_schema = ::arrow::ImportSchema(&arrow_schema).ValueOrDie(); | |
| ASSERT_EQ(imported_schema->num_fields(), 1); | |
| auto field = imported_schema->field(0); | |
| ASSERT_NO_FATAL_FAILURE(CheckArrowField(*field, ::arrow::Type::LIST, kListFieldName, | |
| /*nullable=*/false, kListFieldId)); | |
| auto list_field = std::static_pointer_cast<::arrow::ListType>(field->type()); | |
| ASSERT_NO_FATAL_FAILURE(CheckArrowField(*list_field->value_field(), | |
| ::arrow::Type::INT64, kElemFieldName, | |
| /*nullable=*/true, kElemFieldId)); | |
| } | |
| TEST(ToArrowSchemaTest, MapType) { | |
| constexpr std::string_view kMapFieldName = "map_field"; | |
| constexpr std::string_view kKeyFieldName = "key"; | |
| constexpr std::string_view kValueFieldName = "value"; | |
| constexpr int32_t kFieldId = 1; | |
| constexpr int32_t kKeyFieldId = 2; | |
| constexpr int32_t kValueFieldId = 3; | |
| auto map_type = std::make_shared<MapType>( | |
| SchemaField::MakeRequired(kKeyFieldId, std::string(kKeyFieldName), | |
| iceberg::string()), | |
| SchemaField::MakeOptional(kValueFieldId, std::string(kValueFieldName), | |
| iceberg::int32())); | |
| Schema schema( | |
| {SchemaField::MakeRequired(kFieldId, std::string(kMapFieldName), map_type)}, | |
| /*schema_id=*/0); | |
| ArrowSchema arrow_schema; | |
| ASSERT_THAT(ToArrowSchema(schema, &arrow_schema), IsOk()); | |
| auto imported_schema = ::arrow::ImportSchema(&arrow_schema).ValueOrDie(); | |
| ASSERT_EQ(imported_schema->num_fields(), 1); | |
| auto field = imported_schema->field(0); | |
| ASSERT_NO_FATAL_FAILURE(CheckArrowField(*field, ::arrow::Type::MAP, kMapFieldName, | |
| /*nullable=*/false, kFieldId)); | |
| auto map_field = std::static_pointer_cast<::arrow::MapType>(field->type()); | |
| auto key_field = map_field->key_field(); | |
| ASSERT_NO_FATAL_FAILURE(CheckArrowField(*key_field, ::arrow::Type::STRING, | |
| kKeyFieldName, | |
| /*nullable=*/false, kKeyFieldId)); | |
| auto value_field = map_field->item_field(); | |
| ASSERT_NO_FATAL_FAILURE(CheckArrowField(*value_field, ::arrow::Type::INT32, | |
| kValueFieldName, | |
| /*nullable=*/true, kValueFieldId)); | |
| } | |
| TEST(ToArrowSchemaTest, NestedUnknownFieldsRoundTrip) { | |
| Schema schema( | |
| { | |
| SchemaField::MakeOptional( | |
| /*field_id=*/1, "profile", | |
| std::make_shared<StructType>(std::vector<SchemaField>{ | |
| SchemaField::MakeOptional(/*field_id=*/2, "mystery", | |
| iceberg::unknown()), | |
| })), | |
| SchemaField::MakeOptional( | |
| /*field_id=*/3, "mysteries", | |
| std::make_shared<ListType>(SchemaField::MakeOptional( | |
| /*field_id=*/4, "element", iceberg::unknown()))), | |
| SchemaField::MakeOptional( | |
| /*field_id=*/5, "properties", | |
| std::make_shared<MapType>( | |
| SchemaField::MakeRequired(/*field_id=*/6, "key", iceberg::string()), | |
| SchemaField::MakeOptional(/*field_id=*/7, "value", | |
| iceberg::unknown()))), | |
| }, | |
| /*schema_id=*/0); | |
| ArrowSchema arrow_c_schema; | |
| ASSERT_THAT(ToArrowSchema(schema, &arrow_c_schema), IsOk()); | |
| auto imported_schema = ::arrow::ImportSchema(&arrow_c_schema).ValueOrDie(); | |
| ASSERT_EQ(imported_schema->num_fields(), 3); | |
| auto profile_type = | |
| std::static_pointer_cast<::arrow::StructType>(imported_schema->field(0)->type()); | |
| ASSERT_EQ(profile_type->num_fields(), 1); | |
| ASSERT_NO_FATAL_FAILURE(CheckArrowField(*profile_type->field(0), ::arrow::Type::NA, | |
| "mystery", /*nullable=*/true, | |
| /*field_id=*/2)); | |
| auto mysteries_type = | |
| std::static_pointer_cast<::arrow::ListType>(imported_schema->field(1)->type()); | |
| ASSERT_NO_FATAL_FAILURE(CheckArrowField(*mysteries_type->value_field(), | |
| ::arrow::Type::NA, "element", | |
| /*nullable=*/true, /*field_id=*/4)); | |
| auto properties_type = | |
| std::static_pointer_cast<::arrow::MapType>(imported_schema->field(2)->type()); | |
| ASSERT_NO_FATAL_FAILURE(CheckArrowField(*properties_type->key_field(), | |
| ::arrow::Type::STRING, "key", | |
| /*nullable=*/false, /*field_id=*/6)); | |
| ASSERT_NO_FATAL_FAILURE(CheckArrowField(*properties_type->item_field(), | |
| ::arrow::Type::NA, "value", | |
| /*nullable=*/true, /*field_id=*/7)); | |
| ArrowSchema exported_schema; | |
| ASSERT_TRUE(::arrow::ExportSchema(*imported_schema, &exported_schema).ok()); | |
| auto schema_result = FromArrowSchema(exported_schema, /*schema_id=*/0); | |
| ASSERT_THAT(schema_result, IsOk()); | |
| ArrowSchemaRelease(&exported_schema); | |
| const auto& round_tripped_schema = *schema_result.value(); | |
| ASSERT_EQ(round_tripped_schema.fields().size(), 3); | |
| const auto* profile = | |
| dynamic_cast<const StructType*>(round_tripped_schema.fields()[0].type().get()); | |
| ASSERT_NE(profile, nullptr); | |
| ASSERT_EQ(profile->fields()[0].type()->type_id(), TypeId::kUnknown); | |
| const auto* mysteries = | |
| dynamic_cast<const ListType*>(round_tripped_schema.fields()[1].type().get()); | |
| ASSERT_NE(mysteries, nullptr); | |
| ASSERT_EQ(mysteries->fields()[0].type()->type_id(), TypeId::kUnknown); | |
| const auto* properties = | |
| dynamic_cast<const MapType*>(round_tripped_schema.fields()[2].type().get()); | |
| ASSERT_NE(properties, nullptr); | |
| ASSERT_EQ(properties->value().type()->type_id(), TypeId::kUnknown); | |
| } | |
| struct FromArrowSchemaParam { | |
| std::shared_ptr<::arrow::DataType> arrow_type; | |
| bool optional = true; | |
| std::shared_ptr<Type> iceberg_type; | |
| }; | |
| class FromArrowSchemaTest : public ::testing::TestWithParam<FromArrowSchemaParam> {}; | |
| TEST_P(FromArrowSchemaTest, PrimitiveType) { | |
| constexpr std::string_view kFieldName = "foo"; | |
| constexpr int32_t kFieldId = 1024; | |
| const auto& param = GetParam(); | |
| auto metadata = | |
| ::arrow::key_value_metadata(std::unordered_map<std::string, std::string>{ | |
| {std::string(kParquetFieldIdKey), std::to_string(kFieldId)}}); | |
| auto arrow_schema = ::arrow::schema({::arrow::field( | |
| std::string(kFieldName), param.arrow_type, param.optional, std::move(metadata))}); | |
| ArrowSchema exported_schema; | |
| ASSERT_TRUE(::arrow::ExportSchema(*arrow_schema, &exported_schema).ok()); | |
| auto type_result = FromArrowSchema(exported_schema, /*schema_id=*/1); | |
| ASSERT_THAT(type_result, IsOk()); | |
| ArrowSchemaRelease(&exported_schema); | |
| const auto& schema = type_result.value(); | |
| ASSERT_EQ(schema->schema_id(), 1); | |
| ASSERT_EQ(schema->fields().size(), 1); | |
| const auto& field = schema->fields()[0]; | |
| ASSERT_EQ(field.name(), kFieldName); | |
| ASSERT_EQ(field.field_id(), kFieldId); | |
| ASSERT_EQ(field.optional(), param.optional); | |
| ASSERT_EQ(*field.type(), *param.iceberg_type); | |
| } | |
| TEST(FromArrowSchemaTest, RejectMalformedFieldIdMetadata) { | |
| for (const auto& field_id : {"1x", "2147483648", ""}) { | |
| auto metadata = | |
| ::arrow::key_value_metadata(std::unordered_map<std::string, std::string>{ | |
| {std::string(kParquetFieldIdKey), field_id}}); | |
| auto arrow_schema = ::arrow::schema({::arrow::field( | |
| "foo", ::arrow::int32(), /*nullable=*/true, std::move(metadata))}); | |
| ArrowSchema exported_schema; | |
| ASSERT_TRUE(::arrow::ExportSchema(*arrow_schema, &exported_schema).ok()); | |
| auto result = FromArrowSchema(exported_schema, /*schema_id=*/1); | |
| ArrowSchemaRelease(&exported_schema); | |
| EXPECT_THAT(result, IsError(ErrorKind::kInvalidSchema)); | |
| EXPECT_THAT(result, | |
| HasErrorMessage(std::format("Invalid Arrow field ID: '{}'", field_id))); | |
| } | |
| } | |
| INSTANTIATE_TEST_SUITE_P( | |
| SchemaConversion, FromArrowSchemaTest, | |
| ::testing::Values( | |
| FromArrowSchemaParam{.arrow_type = ::arrow::boolean(), | |
| .optional = false, | |
| .iceberg_type = iceberg::boolean()}, | |
| FromArrowSchemaParam{.arrow_type = ::arrow::int32(), | |
| .optional = true, | |
| .iceberg_type = iceberg::int32()}, | |
| FromArrowSchemaParam{.arrow_type = ::arrow::int64(), | |
| .iceberg_type = iceberg::int64()}, | |
| FromArrowSchemaParam{.arrow_type = ::arrow::float32(), | |
| .iceberg_type = iceberg::float32()}, | |
| FromArrowSchemaParam{.arrow_type = ::arrow::float64(), | |
| .iceberg_type = iceberg::float64()}, | |
| FromArrowSchemaParam{.arrow_type = ::arrow::decimal128(10, 2), | |
| .iceberg_type = iceberg::decimal(10, 2)}, | |
| FromArrowSchemaParam{.arrow_type = ::arrow::date32(), | |
| .iceberg_type = iceberg::date()}, | |
| FromArrowSchemaParam{.arrow_type = ::arrow::time64(::arrow::TimeUnit::MICRO), | |
| .iceberg_type = iceberg::time()}, | |
| FromArrowSchemaParam{.arrow_type = ::arrow::timestamp(::arrow::TimeUnit::MICRO), | |
| .iceberg_type = iceberg::timestamp()}, | |
| FromArrowSchemaParam{ | |
| .arrow_type = ::arrow::timestamp(::arrow::TimeUnit::MICRO, "UTC"), | |
| .iceberg_type = std::make_shared<TimestampTzType>()}, | |
| FromArrowSchemaParam{.arrow_type = ::arrow::timestamp(::arrow::TimeUnit::NANO), | |
| .iceberg_type = iceberg::timestamp_ns()}, | |
| FromArrowSchemaParam{ | |
| .arrow_type = ::arrow::timestamp(::arrow::TimeUnit::NANO, "UTC"), | |
| .iceberg_type = iceberg::timestamptz_ns()}, | |
| FromArrowSchemaParam{.arrow_type = ::arrow::utf8(), | |
| .iceberg_type = iceberg::string()}, | |
| FromArrowSchemaParam{.arrow_type = ::arrow::binary(), | |
| .iceberg_type = iceberg::binary()}, | |
| FromArrowSchemaParam{.arrow_type = ::arrow::extension::uuid(), | |
| .iceberg_type = iceberg::uuid()}, | |
| FromArrowSchemaParam{.arrow_type = ::arrow::fixed_size_binary(20), | |
| .iceberg_type = iceberg::fixed(20)}, | |
| FromArrowSchemaParam{.arrow_type = ::arrow::null(), | |
| .iceberg_type = iceberg::unknown()})); | |
| TEST(FromArrowSchemaTest, RejectRequiredNullFieldAsUnknown) { | |
| auto metadata = | |
| ::arrow::key_value_metadata(std::unordered_map<std::string, std::string>{ | |
| {std::string(kParquetFieldIdKey), "1"}}); | |
| auto arrow_schema = ::arrow::schema({::arrow::field( | |
| "mystery", ::arrow::null(), /*nullable=*/false, std::move(metadata))}); | |
| ArrowSchema exported_schema; | |
| ASSERT_TRUE(::arrow::ExportSchema(*arrow_schema, &exported_schema).ok()); | |
| auto schema_result = FromArrowSchema(exported_schema, /*schema_id=*/0); | |
| ArrowSchemaRelease(&exported_schema); | |
| ASSERT_THAT(schema_result, IsError(ErrorKind::kInvalidSchema)); | |
| ASSERT_THAT(schema_result, | |
| HasErrorMessage("Arrow null field 'mystery' must be nullable")); | |
| } | |
| TEST(FromArrowSchemaTest, RejectRequiredNullListElementAsUnknown) { | |
| auto list_metadata = | |
| ::arrow::key_value_metadata(std::unordered_map<std::string, std::string>{ | |
| {std::string(kParquetFieldIdKey), "1"}}); | |
| auto element_metadata = | |
| ::arrow::key_value_metadata(std::unordered_map<std::string, std::string>{ | |
| {std::string(kParquetFieldIdKey), "2"}}); | |
| auto element_field = ::arrow::field("element", ::arrow::null(), /*nullable=*/false, | |
| std::move(element_metadata)); | |
| auto arrow_schema = | |
| ::arrow::schema({::arrow::field("mysteries", ::arrow::list(element_field), | |
| /*nullable=*/true, std::move(list_metadata))}); | |
| ArrowSchema exported_schema; | |
| ASSERT_TRUE(::arrow::ExportSchema(*arrow_schema, &exported_schema).ok()); | |
| auto schema_result = FromArrowSchema(exported_schema, /*schema_id=*/0); | |
| ArrowSchemaRelease(&exported_schema); | |
| ASSERT_THAT(schema_result, IsError(ErrorKind::kInvalidSchema)); | |
| ASSERT_THAT(schema_result, | |
| HasErrorMessage("Arrow null field 'element' must be nullable")); | |
| } | |
| TEST(FromArrowSchemaTest, StructType) { | |
| constexpr int32_t kStructFieldId = 1; | |
| constexpr int32_t kIntFieldId = 2; | |
| constexpr int32_t kStrFieldId = 3; | |
| constexpr std::string_view kStructFieldName = "struct_field"; | |
| constexpr std::string_view kIntFieldName = "int_field"; | |
| constexpr std::string_view kStrFieldName = "str_field"; | |
| auto int_field = ::arrow::field( | |
| std::string(kIntFieldName), ::arrow::int32(), /*nullable=*/false, | |
| ::arrow::key_value_metadata(std::unordered_map<std::string, std::string>{ | |
| {std::string(kParquetFieldIdKey), std::to_string(kIntFieldId)}})); | |
| auto str_field = ::arrow::field( | |
| std::string(kStrFieldName), ::arrow::utf8(), /*nullable=*/true, | |
| ::arrow::key_value_metadata(std::unordered_map<std::string, std::string>{ | |
| {std::string(kParquetFieldIdKey), std::to_string(kStrFieldId)}})); | |
| auto struct_type = ::arrow::struct_({int_field, str_field}); | |
| auto struct_field = ::arrow::field( | |
| std::string(kStructFieldName), struct_type, /*nullable=*/false, | |
| ::arrow::key_value_metadata(std::unordered_map<std::string, std::string>{ | |
| {std::string(kParquetFieldIdKey), std::to_string(kStructFieldId)}})); | |
| auto arrow_schema = ::arrow::schema({struct_field}); | |
| ArrowSchema exported_schema; | |
| ASSERT_TRUE(::arrow::ExportSchema(*arrow_schema, &exported_schema).ok()); | |
| auto schema_result = FromArrowSchema(exported_schema, /*schema_id=*/0); | |
| ASSERT_THAT(schema_result, IsOk()); | |
| ArrowSchemaRelease(&exported_schema); | |
| const auto& iceberg_schema = schema_result.value(); | |
| ASSERT_EQ(iceberg_schema->schema_id(), 0); | |
| ASSERT_EQ(iceberg_schema->fields().size(), 1); | |
| const auto& field = iceberg_schema->fields()[0]; | |
| ASSERT_EQ(field.name(), kStructFieldName); | |
| ASSERT_EQ(field.field_id(), kStructFieldId); | |
| ASSERT_FALSE(field.optional()); | |
| ASSERT_EQ(field.type()->type_id(), TypeId::kStruct); | |
| auto* struct_field_type = dynamic_cast<const StructType*>(field.type().get()); | |
| ASSERT_NE(struct_field_type, nullptr); | |
| ASSERT_EQ(struct_field_type->fields().size(), 2); | |
| const auto& int_iceberg_field = struct_field_type->fields()[0]; | |
| ASSERT_EQ(int_iceberg_field.name(), kIntFieldName); | |
| ASSERT_EQ(int_iceberg_field.field_id(), kIntFieldId); | |
| ASSERT_FALSE(int_iceberg_field.optional()); | |
| ASSERT_EQ(int_iceberg_field.type()->type_id(), TypeId::kInt); | |
| const auto& str_iceberg_field = struct_field_type->fields()[1]; | |
| ASSERT_EQ(str_iceberg_field.name(), kStrFieldName); | |
| ASSERT_EQ(str_iceberg_field.field_id(), kStrFieldId); | |
| ASSERT_TRUE(str_iceberg_field.optional()); | |
| ASSERT_EQ(str_iceberg_field.type()->type_id(), TypeId::kString); | |
| } | |
| TEST(FromArrowSchemaTest, ListType) { | |
| constexpr std::string_view kListFieldName = "list_field"; | |
| constexpr std::string_view kElemFieldName = "element"; | |
| constexpr int32_t kListFieldId = 1; | |
| constexpr int32_t kElemFieldId = 2; | |
| auto element_field = ::arrow::field( | |
| std::string(kElemFieldName), ::arrow::int64(), /*nullable=*/true, | |
| ::arrow::key_value_metadata(std::unordered_map<std::string, std::string>{ | |
| {std::string(kParquetFieldIdKey), std::to_string(kElemFieldId)}})); | |
| auto list_type = ::arrow::list(element_field); | |
| auto list_field = ::arrow::field( | |
| std::string(kListFieldName), list_type, /*nullable=*/false, | |
| ::arrow::key_value_metadata(std::unordered_map<std::string, std::string>{ | |
| {std::string(kParquetFieldIdKey), std::to_string(kListFieldId)}})); | |
| auto arrow_schema = ::arrow::schema({list_field}); | |
| ArrowSchema exported_schema; | |
| ASSERT_TRUE(::arrow::ExportSchema(*arrow_schema, &exported_schema).ok()); | |
| auto schema_result = FromArrowSchema(exported_schema, /*schema_id=*/0); | |
| ASSERT_THAT(schema_result, IsOk()); | |
| ArrowSchemaRelease(&exported_schema); | |
| const auto& iceberg_schema = schema_result.value(); | |
| ASSERT_EQ(iceberg_schema->schema_id(), 0); | |
| ASSERT_EQ(iceberg_schema->fields().size(), 1); | |
| const auto& field = iceberg_schema->fields()[0]; | |
| ASSERT_EQ(field.name(), kListFieldName); | |
| ASSERT_EQ(field.field_id(), kListFieldId); | |
| ASSERT_FALSE(field.optional()); | |
| ASSERT_EQ(field.type()->type_id(), TypeId::kList); | |
| auto* list_field_type = dynamic_cast<const ListType*>(field.type().get()); | |
| ASSERT_NE(list_field_type, nullptr); | |
| const auto& element = list_field_type->fields()[0]; | |
| ASSERT_EQ(element.name(), kElemFieldName); | |
| ASSERT_EQ(element.field_id(), kElemFieldId); | |
| ASSERT_TRUE(element.optional()); | |
| ASSERT_EQ(element.type()->type_id(), TypeId::kLong); | |
| } | |
| TEST(FromArrowSchemaTest, MapType) { | |
| constexpr std::string_view kMapFieldName = "map_field"; | |
| constexpr std::string_view kKeyFieldName = "key"; | |
| constexpr std::string_view kValueFieldName = "value"; | |
| constexpr int32_t kFieldId = 1; | |
| constexpr int32_t kKeyFieldId = 2; | |
| constexpr int32_t kValueFieldId = 3; | |
| auto key_field = ::arrow::field( | |
| std::string(kKeyFieldName), ::arrow::utf8(), /*nullable=*/false, | |
| ::arrow::key_value_metadata(std::unordered_map<std::string, std::string>{ | |
| {std::string(kParquetFieldIdKey), std::to_string(kKeyFieldId)}})); | |
| auto value_field = ::arrow::field( | |
| std::string(kValueFieldName), ::arrow::int32(), /*nullable=*/true, | |
| ::arrow::key_value_metadata(std::unordered_map<std::string, std::string>{ | |
| {std::string(kParquetFieldIdKey), std::to_string(kValueFieldId)}})); | |
| auto map_type = std::make_shared<::arrow::MapType>(key_field, value_field); | |
| auto map_field = ::arrow::field( | |
| std::string(kMapFieldName), map_type, /*nullable=*/true, | |
| ::arrow::key_value_metadata(std::unordered_map<std::string, std::string>{ | |
| {std::string(kParquetFieldIdKey), std::to_string(kFieldId)}})); | |
| auto arrow_schema = ::arrow::schema({map_field}); | |
| ArrowSchema exported_schema; | |
| ASSERT_TRUE(::arrow::ExportSchema(*arrow_schema, &exported_schema).ok()); | |
| auto schema_result = FromArrowSchema(exported_schema, /*schema_id=*/0); | |
| ASSERT_THAT(schema_result, IsOk()); | |
| ArrowSchemaRelease(&exported_schema); | |
| const auto& iceberg_schema = schema_result.value(); | |
| ASSERT_EQ(iceberg_schema->schema_id(), 0); | |
| ASSERT_EQ(iceberg_schema->fields().size(), 1); | |
| const auto& field = iceberg_schema->fields()[0]; | |
| ASSERT_EQ(field.name(), kMapFieldName); | |
| ASSERT_EQ(field.field_id(), kFieldId); | |
| ASSERT_TRUE(field.optional()); | |
| ASSERT_EQ(field.type()->type_id(), TypeId::kMap); | |
| auto* map_field_type = dynamic_cast<const MapType*>(field.type().get()); | |
| ASSERT_NE(map_field_type, nullptr); | |
| const auto& key = map_field_type->key(); | |
| ASSERT_EQ(key.name(), kKeyFieldName); | |
| ASSERT_EQ(key.field_id(), kKeyFieldId); | |
| ASSERT_FALSE(key.optional()); | |
| ASSERT_EQ(key.type()->type_id(), TypeId::kString); | |
| const auto& value = map_field_type->value(); | |
| ASSERT_EQ(value.name(), kValueFieldName); | |
| ASSERT_EQ(value.field_id(), kValueFieldId); | |
| ASSERT_TRUE(value.optional()); | |
| ASSERT_EQ(value.type()->type_id(), TypeId::kInt); | |
| } | |
| TEST(ArrowExecutorAdapterTest, RunsTaskGroupOnThreadPool) { | |
| #ifndef ARROW_ENABLE_THREADING | |
| GTEST_SKIP() << "Test requires ARROW_ENABLE_THREADING=ON"; | |
| #endif | |
| class ArrowExecutorAdapter final : public Executor { | |
| public: | |
| explicit ArrowExecutorAdapter(::arrow::internal::Executor& executor) | |
| : executor_(executor) {} | |
| Status Submit(ExecutorTask task) override { | |
| ICEBERG_ARROW_RETURN_NOT_OK(executor_.Spawn(std::move(task))); | |
| return {}; | |
| } | |
| private: | |
| ::arrow::internal::Executor& executor_; | |
| }; | |
| auto thread_pool = ::arrow::internal::ThreadPool::Make(2).ValueOrDie(); | |
| ArrowExecutorAdapter executor(*thread_pool); | |
| std::mutex mutex; | |
| std::vector<std::thread::id> thread_ids; | |
| auto status = TaskGroup() | |
| .SetExecutor(std::ref(executor)) | |
| .Submit([&]() -> Status { | |
| std::lock_guard lock(mutex); | |
| thread_ids.push_back(std::this_thread::get_id()); | |
| return {}; | |
| }) | |
| .Submit([&]() -> Status { | |
| std::lock_guard lock(mutex); | |
| thread_ids.push_back(std::this_thread::get_id()); | |
| return {}; | |
| }) | |
| .Run(); | |
| EXPECT_THAT(status, IsOk()); | |
| EXPECT_EQ(thread_ids.size(), 2); | |
| EXPECT_NE(thread_ids[0], std::this_thread::get_id()); | |
| EXPECT_NE(thread_ids[1], std::this_thread::get_id()); | |
| EXPECT_TRUE(thread_pool->Shutdown().ok()); | |
| } | |
| } // namespace iceberg |