diff --git a/src/iceberg/avro/avro_data_util.cc b/src/iceberg/avro/avro_data_util.cc index 39f2948b9..c3843c2f6 100644 --- a/src/iceberg/avro/avro_data_util.cc +++ b/src/iceberg/avro/avro_data_util.cc @@ -748,7 +748,6 @@ Status ExtractDatumFromArray(const ::arrow::Array& array, int64_t index, return {}; } - // TODO(gangwu): support LARGE_STRING. case ::arrow::Type::STRING: { const auto& string_array = internal::checked_cast(array); @@ -756,7 +755,13 @@ Status ExtractDatumFromArray(const ::arrow::Array& array, int64_t index, return {}; } - // TODO(gangwu): support LARGE_BINARY. + case ::arrow::Type::LARGE_STRING: { + const auto& large_string_array = + internal::checked_cast(array); + datum->value() = large_string_array.GetString(index); + return {}; + } + case ::arrow::Type::BINARY: { const auto& binary_array = internal::checked_cast(array); @@ -767,6 +772,16 @@ Status ExtractDatumFromArray(const ::arrow::Array& array, int64_t index, return {}; } + case ::arrow::Type::LARGE_BINARY: { + const auto& large_binary_array = + internal::checked_cast(array); + std::string_view value = large_binary_array.GetView(index); + datum->value>().assign( + reinterpret_cast(value.data()), + reinterpret_cast(value.data()) + value.size()); + return {}; + } + case ::arrow::Type::FIXED_SIZE_BINARY: { const auto& fixed_array = internal::checked_cast(array); @@ -835,7 +850,6 @@ Status ExtractDatumFromArray(const ::arrow::Array& array, int64_t index, return {}; } - // TODO(gangwu): support LARGE_LIST. case ::arrow::Type::LIST: { const auto& list_array = internal::checked_cast(array); auto& avro_array = datum->value<::avro::GenericArray>(); @@ -855,6 +869,26 @@ Status ExtractDatumFromArray(const ::arrow::Array& array, int64_t index, return {}; } + case ::arrow::Type::LARGE_LIST: { + const auto& large_list_array = + internal::checked_cast(array); + auto& avro_array = datum->value<::avro::GenericArray>(); + auto& elements = avro_array.value(); + + auto start = large_list_array.value_offset(index); + auto end = large_list_array.value_offset(index + 1); + auto length = end - start; + + auto values = large_list_array.values(); + elements.resize(length, ::avro::GenericDatum(avro_array.schema()->leafAt(0))); + + for (int64_t i = 0; i < length; ++i) { + ICEBERG_RETURN_UNEXPECTED( + ExtractDatumFromArray(*values, start + i, &elements[i])); + } + return {}; + } + case ::arrow::Type::MAP: { const auto& map_array = internal::checked_cast(array); auto start = map_array.value_offset(index); diff --git a/src/iceberg/test/avro_data_test.cc b/src/iceberg/test/avro_data_test.cc index 9ec0221da..78ef2bc6f 100644 --- a/src/iceberg/test/avro_data_test.cc +++ b/src/iceberg/test/avro_data_test.cc @@ -21,6 +21,8 @@ #include #include +#include +#include #include #include #include @@ -1547,6 +1549,131 @@ TEST(ExtractDatumFromArrayTest, MapStringToInt) { } } +TEST(ExtractDatumFromArrayTest, LargeListOfIntegers) { + Schema iceberg_schema({ + SchemaField::MakeRequired( + 1, "numbers", + std::make_shared(SchemaField::MakeRequired(2, "element", int32()))), + }); + + ::avro::NodePtr avro_node; + ASSERT_THAT(ToAvroNodeVisitor{}.Visit(iceberg_schema, &avro_node), IsOk()); + + // Build a LargeListArray manually — JSON parsing always produces regular ListArray, + // so we use builders to get the large-offset variant. + auto int32_builder = std::make_shared<::arrow::Int32Builder>(); + ::arrow::LargeListBuilder list_builder(::arrow::default_memory_pool(), int32_builder); + + ASSERT_TRUE(list_builder.Append().ok()); + ASSERT_TRUE(int32_builder->Append(10).ok()); + ASSERT_TRUE(int32_builder->Append(11).ok()); + ASSERT_TRUE(int32_builder->Append(12).ok()); + + ASSERT_TRUE(list_builder.Append().ok()); + ASSERT_TRUE(int32_builder->Append(20).ok()); + ASSERT_TRUE(int32_builder->Append(21).ok()); + + ASSERT_TRUE(list_builder.Append().ok()); + ASSERT_TRUE(int32_builder->Append(30).ok()); + ASSERT_TRUE(int32_builder->Append(31).ok()); + ASSERT_TRUE(int32_builder->Append(32).ok()); + ASSERT_TRUE(int32_builder->Append(33).ok()); + + std::shared_ptr<::arrow::Array> large_list_array; + ASSERT_TRUE(list_builder.Finish(&large_list_array).ok()); + + auto arrow_array = ::arrow::StructArray::Make( + ::arrow::ArrayVector{large_list_array}, + std::vector>{::arrow::field( + "numbers", ::arrow::large_list(::arrow::int32()))}) + .ValueOrDie(); + + std::vector> expected = {{10, 11, 12}, {20, 21}, {30, 31, 32, 33}}; + + auto verify_record = [&](int64_t index, const std::vector& expected_numbers) { + ::avro::GenericDatum extracted_datum(avro_node); + ASSERT_THAT(ExtractDatumFromArray(*arrow_array, index, &extracted_datum), IsOk()); + const auto& record = extracted_datum.value<::avro::GenericRecord>(); + const auto& avro_arr = record.fieldAt(0).value<::avro::GenericArray>(); + const auto& elements = avro_arr.value(); + + ASSERT_EQ(elements.size(), expected_numbers.size()); + for (size_t i = 0; i < expected_numbers.size(); ++i) { + EXPECT_EQ(elements[i].value(), expected_numbers[i]); + } + }; + + for (size_t i = 0; i < expected.size(); ++i) { + verify_record(i, expected[i]); + } +} + +TEST(ExtractDatumFromArrayTest, LargeString) { + Schema iceberg_schema({ + SchemaField::MakeRequired(1, "name", string()), + }); + + ::avro::NodePtr avro_node; + ASSERT_THAT(ToAvroNodeVisitor{}.Visit(iceberg_schema, &avro_node), IsOk()); + + // Build a LargeStringArray manually + ::arrow::LargeStringBuilder builder; + ASSERT_TRUE(builder.Append("hello").ok()); + ASSERT_TRUE(builder.Append("world").ok()); + ASSERT_TRUE(builder.Append("iceberg").ok()); + std::shared_ptr<::arrow::Array> large_string_array; + ASSERT_TRUE(builder.Finish(&large_string_array).ok()); + + auto arrow_array = + ::arrow::StructArray::Make(::arrow::ArrayVector{large_string_array}, + std::vector>{ + ::arrow::field("name", ::arrow::large_utf8())}) + .ValueOrDie(); + + std::vector expected = {"hello", "world", "iceberg"}; + + for (int64_t i = 0; std::cmp_less(i, expected.size()); ++i) { + ::avro::GenericDatum extracted_datum(avro_node); + ASSERT_THAT(ExtractDatumFromArray(*arrow_array, i, &extracted_datum), IsOk()); + const auto& record = extracted_datum.value<::avro::GenericRecord>(); + EXPECT_EQ(record.fieldAt(0).value(), expected[i]); + } +} + +TEST(ExtractDatumFromArrayTest, LargeBinary) { + Schema iceberg_schema({ + SchemaField::MakeRequired(1, "data", binary()), + }); + + ::avro::NodePtr avro_node; + ASSERT_THAT(ToAvroNodeVisitor{}.Visit(iceberg_schema, &avro_node), IsOk()); + + // Build a LargeBinaryArray manually + ::arrow::LargeBinaryBuilder builder; + ASSERT_TRUE(builder.Append("abc").ok()); + ASSERT_TRUE(builder.Append("bcd").ok()); + ASSERT_TRUE(builder.Append("cde").ok()); + std::shared_ptr<::arrow::Array> large_binary_array; + ASSERT_TRUE(builder.Finish(&large_binary_array).ok()); + + auto arrow_array = + ::arrow::StructArray::Make(::arrow::ArrayVector{large_binary_array}, + std::vector>{ + ::arrow::field("data", ::arrow::large_binary())}) + .ValueOrDie(); + + for (int64_t i = 0; i < 3; ++i) { + ::avro::GenericDatum extracted_datum(avro_node); + ASSERT_THAT(ExtractDatumFromArray(*arrow_array, i, &extracted_datum), IsOk()); + const auto& record = extracted_datum.value<::avro::GenericRecord>(); + const auto& bytes = record.fieldAt(0).value>(); + ASSERT_EQ(bytes.size(), 3); + EXPECT_EQ(bytes[0], static_cast('a' + i)); + EXPECT_EQ(bytes[1], static_cast('b' + i)); + EXPECT_EQ(bytes[2], static_cast('c' + i)); + } +} + TEST(ExtractDatumFromArrayTest, ErrorHandling) { Schema iceberg_schema({SchemaField::MakeRequired(1, "a", int32())}); ::avro::NodePtr avro_node;