| /* |
| * 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 <optional> |
| #include <sstream> |
| #include <string> |
| #include <string_view> |
| #include <unordered_map> |
| |
| #include <arrow/array.h> |
| #include <arrow/array/array_base.h> |
| #include <arrow/c/bridge.h> |
| #include <arrow/filesystem/filesystem.h> |
| #include <arrow/json/from_string.h> |
| #include <avro/DataFile.hh> |
| #include <avro/Generic.hh> |
| #include <gtest/gtest.h> |
| |
| #include "iceberg/arrow/arrow_io_internal.h" |
| #include "iceberg/avro/avro_constants.h" |
| #include "iceberg/avro/avro_register.h" |
| #include "iceberg/avro/avro_stream_internal.h" |
| #include "iceberg/avro/avro_writer.h" |
| #include "iceberg/file_reader.h" |
| #include "iceberg/metadata_columns.h" |
| #include "iceberg/schema.h" |
| #include "iceberg/schema_internal.h" |
| #include "iceberg/test/matchers.h" |
| #include "iceberg/test/std_io.h" |
| #include "iceberg/test/temp_file_test_base.h" |
| #include "iceberg/type.h" |
| #include "iceberg/util/checked_cast.h" |
| |
| namespace iceberg::avro { |
| |
| namespace { |
| |
| ::avro::NodePtr UnwrapOptional(const ::avro::NodePtr& node) { |
| if (node->type() != ::avro::AVRO_UNION) { |
| return node; |
| } |
| |
| for (size_t i = 0; i < node->leaves(); ++i) { |
| if (node->leafAt(i)->type() != ::avro::AVRO_NULL) { |
| return node->leafAt(i); |
| } |
| } |
| return node; |
| } |
| |
| std::optional<int32_t> FieldIdAt(const ::avro::NodePtr& node, size_t index) { |
| if (index >= node->customAttributes()) { |
| return std::nullopt; |
| } |
| |
| auto field_id = node->customAttributesAt(index).getAttribute(std::string(kFieldIdProp)); |
| if (!field_id.has_value()) { |
| return std::nullopt; |
| } |
| return std::stoi(field_id.value()); |
| } |
| |
| } // namespace |
| |
| class AvroReaderTest : public TempFileTestBase { |
| protected: |
| static void SetUpTestSuite() { RegisterAll(); } |
| |
| void SetUp() override { |
| TempFileTestBase::SetUp(); |
| file_io_ = arrow::ArrowFileSystemFileIO::MakeMockFileIO(); |
| temp_avro_file_ = "avro_reader_test.avro"; |
| } |
| |
| bool skip_datum_{true}; |
| |
| void CreateSimpleAvroFile() { |
| // Create simple avro file using the writer API instead of direct Avro library |
| auto schema = std::make_shared<Schema>(std::vector<SchemaField>{ |
| SchemaField::MakeRequired(1, "id", std::make_shared<IntType>()), |
| SchemaField::MakeOptional(2, "name", std::make_shared<StringType>())}); |
| |
| ArrowSchema arrow_c_schema; |
| ASSERT_THAT(ToArrowSchema(*schema, &arrow_c_schema), IsOk()); |
| auto arrow_schema = ::arrow::ImportType(&arrow_c_schema).ValueOrDie(); |
| |
| auto array = ::arrow::json::ArrayFromJSONString( |
| ::arrow::struct_(arrow_schema->fields()), |
| R"([[1, "Alice"], [2, "Bob"], [3, "Charlie"]])") |
| .ValueOrDie(); |
| |
| struct ArrowArray arrow_array; |
| auto export_result = ::arrow::ExportArray(*array, &arrow_array); |
| ASSERT_TRUE(export_result.ok()); |
| |
| auto writer_result = |
| WriterFactoryRegistry::Open(FileFormatType::kAvro, { |
| .path = temp_avro_file_, |
| .schema = schema, |
| .io = file_io_, |
| }); |
| ASSERT_TRUE(writer_result.has_value()); |
| auto writer = std::move(writer_result.value()); |
| ASSERT_THAT(writer->Write(&arrow_array), IsOk()); |
| ASSERT_THAT(writer->Close(), IsOk()); |
| } |
| |
| void VerifyNextBatch(Reader& reader, std::string_view expected_json) { |
| // Boilerplate to get Arrow schema |
| auto schema_result = reader.Schema(); |
| ASSERT_THAT(schema_result, IsOk()); |
| auto arrow_c_schema = std::move(schema_result.value()); |
| auto import_schema_result = ::arrow::ImportType(&arrow_c_schema); |
| auto arrow_schema = import_schema_result.ValueOrDie(); |
| |
| // Boilerplate to get Arrow array |
| auto data = reader.Next(); |
| ASSERT_THAT(data, IsOk()); |
| ASSERT_TRUE(data.value().has_value()); |
| auto arrow_c_array = data.value().value(); |
| auto data_result = ::arrow::ImportArray(&arrow_c_array, arrow_schema); |
| auto arrow_array = data_result.ValueOrDie(); |
| |
| // Verify data |
| auto expected_array = |
| ::arrow::json::ArrayFromJSONString(arrow_schema, expected_json).ValueOrDie(); |
| ASSERT_TRUE(arrow_array->Equals(*expected_array)); |
| } |
| |
| void VerifyExhausted(Reader& reader) { |
| auto data = reader.Next(); |
| ASSERT_THAT(data, IsOk()); |
| ASSERT_FALSE(data.value().has_value()); |
| } |
| |
| void WriteAndVerify(std::shared_ptr<Schema> schema, |
| const std::string& expected_string) { |
| ArrowSchema arrow_c_schema; |
| ASSERT_THAT(ToArrowSchema(*schema, &arrow_c_schema), IsOk()); |
| |
| auto arrow_schema_result = ::arrow::ImportType(&arrow_c_schema); |
| ASSERT_TRUE(arrow_schema_result.ok()); |
| auto arrow_schema = arrow_schema_result.ValueOrDie(); |
| |
| auto array_result = ::arrow::json::ArrayFromJSONString(arrow_schema, expected_string); |
| ASSERT_TRUE(array_result.ok()); |
| auto array = array_result.ValueOrDie(); |
| |
| struct ArrowArray arrow_array; |
| auto export_result = ::arrow::ExportArray(*array, &arrow_array); |
| ASSERT_TRUE(export_result.ok()); |
| |
| std::unordered_map<std::string, std::string> metadata = {{"k1", "v1"}, {"k2", "v2"}}; |
| |
| auto writer_result = |
| WriterFactoryRegistry::Open(FileFormatType::kAvro, {.path = temp_avro_file_, |
| .schema = schema, |
| .io = file_io_, |
| .metadata = metadata}); |
| ASSERT_TRUE(writer_result.has_value()); |
| auto writer = std::move(writer_result.value()); |
| ASSERT_THAT(writer->Write(&arrow_array), IsOk()); |
| ASSERT_THAT(writer->Close(), IsOk()); |
| ICEBERG_UNWRAP_OR_FAIL(auto written_length, writer->length()); |
| |
| ReaderProperties reader_properties; |
| reader_properties.Set(ReaderProperties::kAvroSkipDatum, skip_datum_); |
| |
| auto reader_result = ReaderFactoryRegistry::Open( |
| FileFormatType::kAvro, {.path = temp_avro_file_, |
| .length = written_length, |
| .io = file_io_, |
| .projection = schema, |
| .properties = std::move(reader_properties)}); |
| ASSERT_THAT(reader_result, IsOk()); |
| auto reader = std::move(reader_result.value()); |
| ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader, expected_string)); |
| ASSERT_NO_FATAL_FAILURE(VerifyExhausted(*reader)); |
| |
| auto metadata_result = reader->Metadata(); |
| ASSERT_THAT(metadata_result, IsOk()); |
| auto read_metadata = std::move(metadata_result.value()); |
| for (const auto& [key, value] : metadata) { |
| auto it = read_metadata.find(key); |
| ASSERT_NE(it, read_metadata.end()); |
| ASSERT_EQ(it->second, value); |
| } |
| } |
| |
| std::shared_ptr<FileIO> file_io_; |
| std::string temp_avro_file_; |
| }; |
| |
| TEST_F(AvroReaderTest, ReadTwoFields) { |
| CreateSimpleAvroFile(); |
| auto schema = std::make_shared<Schema>(std::vector<SchemaField>{ |
| SchemaField::MakeRequired(1, "id", std::make_shared<IntType>()), |
| SchemaField::MakeOptional(2, "name", std::make_shared<StringType>())}); |
| |
| auto reader_result = ReaderFactoryRegistry::Open( |
| FileFormatType::kAvro, |
| {.path = temp_avro_file_, .io = file_io_, .projection = schema}); |
| ASSERT_THAT(reader_result, IsOk()); |
| auto reader = std::move(reader_result.value()); |
| |
| ASSERT_NO_FATAL_FAILURE( |
| VerifyNextBatch(*reader, R"([[1, "Alice"], [2, "Bob"], [3, "Charlie"]])")); |
| ASSERT_NO_FATAL_FAILURE(VerifyExhausted(*reader)); |
| } |
| |
| TEST_F(AvroReaderTest, RoundTripWithGenericFileIO) { |
| file_io_ = std::make_shared<iceberg::test::StdFileIO>(); |
| temp_avro_file_ = CreateNewTempFilePathWithSuffix(".avro"); |
| auto schema = std::make_shared<Schema>(std::vector<SchemaField>{ |
| SchemaField::MakeRequired(1, "id", std::make_shared<IntType>()), |
| SchemaField::MakeOptional(2, "name", std::make_shared<StringType>())}); |
| |
| ASSERT_NO_FATAL_FAILURE(WriteAndVerify(schema, R"([[1, "Foo"], [2, "Bar"]])")); |
| } |
| |
| TEST_F(AvroReaderTest, ReadReorderedFieldsWithNulls) { |
| CreateSimpleAvroFile(); |
| auto schema = std::make_shared<Schema>(std::vector<SchemaField>{ |
| SchemaField::MakeOptional(2, "name", std::make_shared<StringType>()), |
| SchemaField::MakeRequired(1, "id", std::make_shared<IntType>()), |
| SchemaField::MakeOptional(3, "score", std::make_shared<DoubleType>())}); |
| |
| auto reader_result = ReaderFactoryRegistry::Open( |
| FileFormatType::kAvro, |
| {.path = temp_avro_file_, .io = file_io_, .projection = schema}); |
| ASSERT_THAT(reader_result, IsOk()); |
| auto reader = std::move(reader_result.value()); |
| |
| ASSERT_NO_FATAL_FAILURE(VerifyNextBatch( |
| *reader, R"([["Alice", 1, null], ["Bob", 2, null], ["Charlie", 3, null]])")); |
| ASSERT_NO_FATAL_FAILURE(VerifyExhausted(*reader)); |
| } |
| |
| TEST_F(AvroReaderTest, ReadWithBatchSize) { |
| CreateSimpleAvroFile(); |
| auto schema = std::make_shared<Schema>(std::vector<SchemaField>{ |
| SchemaField::MakeRequired(1, "id", std::make_shared<IntType>())}); |
| |
| ReaderProperties reader_properties; |
| reader_properties.Set(ReaderProperties::kBatchSize, int64_t{2}); |
| |
| auto reader_result = ReaderFactoryRegistry::Open( |
| FileFormatType::kAvro, {.path = temp_avro_file_, |
| .io = file_io_, |
| .projection = schema, |
| .properties = std::move(reader_properties)}); |
| ASSERT_THAT(reader_result, IsOk()); |
| auto reader = std::move(reader_result.value()); |
| |
| ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader, R"([[1], [2]])")); |
| ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader, R"([[3]])")); |
| ASSERT_NO_FATAL_FAILURE(VerifyExhausted(*reader)); |
| } |
| |
| TEST_F(AvroReaderTest, BufferSizeConfiguration) { |
| // Test default buffer size |
| ReaderProperties properties1; |
| ASSERT_EQ(properties1.Get(ReaderProperties::kAvroBufferSize), 1024 * 1024); |
| |
| // Test setting custom buffer size |
| ReaderProperties properties2; |
| constexpr int64_t kCustomBufferSize = 2 * 1024 * 1024; // 2MB |
| properties2.Set(ReaderProperties::kAvroBufferSize, kCustomBufferSize); |
| ASSERT_EQ(properties2.Get(ReaderProperties::kAvroBufferSize), kCustomBufferSize); |
| |
| // Test setting via FromMap |
| std::unordered_map<std::string, std::string> config_map = { |
| {"read.avro.buffer-size", "4194304"} // 4MB |
| }; |
| auto properties3 = ReaderProperties::FromMap(config_map); |
| ASSERT_EQ(properties3.Get(ReaderProperties::kAvroBufferSize), 4194304); |
| |
| // Test that unset returns to default |
| properties2.Unset(ReaderProperties::kAvroBufferSize); |
| ASSERT_EQ(properties2.Get(ReaderProperties::kAvroBufferSize), 1024 * 1024); |
| } |
| |
| // Parameterized test fixture for testing both DirectDecoder and GenericDatum modes |
| class AvroReaderParameterizedTest : public AvroReaderTest, |
| public ::testing::WithParamInterface<bool> { |
| protected: |
| void SetUp() override { |
| AvroReaderTest::SetUp(); |
| skip_datum_ = GetParam(); |
| } |
| }; |
| |
| TEST_P(AvroReaderParameterizedTest, AvroWriterBasicType) { |
| auto schema = std::make_shared<iceberg::Schema>(std::vector<SchemaField>{ |
| SchemaField::MakeRequired(1, "name", std::make_shared<StringType>())}); |
| |
| std::string expected_string = R"([["Hello"], ["世界"], ["nanoarrow"]])"; |
| |
| WriteAndVerify(schema, expected_string); |
| } |
| |
| TEST_P(AvroReaderParameterizedTest, AvroWriterNestedType) { |
| auto schema = std::make_shared<iceberg::Schema>(std::vector<SchemaField>{ |
| SchemaField::MakeRequired(1, "id", std::make_shared<IntType>()), |
| SchemaField::MakeRequired( |
| 2, "info", |
| std::make_shared<iceberg::StructType>(std::vector<SchemaField>{ |
| SchemaField::MakeRequired(3, "name", std::make_shared<StringType>()), |
| SchemaField::MakeRequired(4, "age", std::make_shared<IntType>())}))}); |
| |
| std::string expected_string = |
| R"([[1, ["Alice", 25]], [2, ["Bob", 30]], [3, ["Ivy", 35]]])"; |
| |
| WriteAndVerify(schema, expected_string); |
| } |
| |
| TEST_P(AvroReaderParameterizedTest, AllPrimitiveTypes) { |
| auto schema = std::make_shared<iceberg::Schema>(std::vector<SchemaField>{ |
| SchemaField::MakeRequired(1, "bool_col", std::make_shared<BooleanType>()), |
| SchemaField::MakeRequired(2, "int_col", std::make_shared<IntType>()), |
| SchemaField::MakeRequired(3, "long_col", std::make_shared<LongType>()), |
| SchemaField::MakeRequired(4, "float_col", std::make_shared<FloatType>()), |
| SchemaField::MakeRequired(5, "double_col", std::make_shared<DoubleType>()), |
| SchemaField::MakeRequired(6, "string_col", std::make_shared<StringType>()), |
| SchemaField::MakeRequired(7, "binary_col", std::make_shared<BinaryType>())}); |
| |
| std::string expected_string = R"([ |
| [true, 42, 1234567890, 3.14, 2.71828, "test", "AQID"], |
| [false, -100, -9876543210, -1.5, 0.0, "hello", "BAUG"] |
| ])"; |
| |
| WriteAndVerify(schema, expected_string); |
| } |
| |
| // Skipping DecimalType test - requires specific decimal encoding in JSON |
| |
| TEST_P(AvroReaderParameterizedTest, DateTimeTypes) { |
| auto schema = std::make_shared<iceberg::Schema>(std::vector<SchemaField>{ |
| SchemaField::MakeRequired(1, "date_col", std::make_shared<DateType>()), |
| SchemaField::MakeRequired(2, "time_col", std::make_shared<TimeType>()), |
| SchemaField::MakeRequired(3, "timestamp_col", std::make_shared<TimestampType>())}); |
| |
| // Dates as days since epoch, time/timestamps as microseconds |
| std::string expected_string = R"([ |
| [18628, 43200000000, 1640995200000000], |
| [18629, 86399000000, 1641081599000000] |
| ])"; |
| |
| WriteAndVerify(schema, expected_string); |
| } |
| |
| TEST_P(AvroReaderParameterizedTest, NestedStruct) { |
| auto schema = std::make_shared<iceberg::Schema>(std::vector<SchemaField>{ |
| SchemaField::MakeRequired(1, "id", std::make_shared<IntType>()), |
| SchemaField::MakeRequired( |
| 2, "person", |
| std::make_shared<iceberg::StructType>(std::vector<SchemaField>{ |
| SchemaField::MakeRequired(3, "name", std::make_shared<StringType>()), |
| SchemaField::MakeRequired(4, "age", std::make_shared<IntType>()), |
| SchemaField::MakeOptional( |
| 5, "address", |
| std::make_shared<iceberg::StructType>(std::vector<SchemaField>{ |
| SchemaField::MakeRequired(6, "street", |
| std::make_shared<StringType>()), |
| SchemaField::MakeRequired(7, "city", |
| std::make_shared<StringType>())}))}))}); |
| |
| std::string expected_string = R"([ |
| [1, ["Alice", 30, ["123 Main St", "NYC"]]], |
| [2, ["Bob", 25, ["456 Oak Ave", "LA"]]] |
| ])"; |
| |
| WriteAndVerify(schema, expected_string); |
| } |
| |
| TEST_P(AvroReaderParameterizedTest, ListType) { |
| auto schema = std::make_shared<iceberg::Schema>(std::vector<SchemaField>{ |
| SchemaField::MakeRequired(1, "id", std::make_shared<IntType>()), |
| SchemaField::MakeRequired(2, "tags", |
| std::make_shared<ListType>(SchemaField::MakeRequired( |
| 3, "element", std::make_shared<StringType>())))}); |
| |
| std::string expected_string = R"([ |
| [1, ["tag1", "tag2", "tag3"]], |
| [2, ["foo", "bar"]], |
| [3, []] |
| ])"; |
| |
| WriteAndVerify(schema, expected_string); |
| } |
| |
| TEST_P(AvroReaderParameterizedTest, MapType) { |
| auto schema = std::make_shared<iceberg::Schema>( |
| std::vector<SchemaField>{SchemaField::MakeRequired( |
| 1, "properties", |
| std::make_shared<MapType>( |
| SchemaField::MakeRequired(2, "key", std::make_shared<StringType>()), |
| SchemaField::MakeRequired(3, "value", std::make_shared<IntType>())))}); |
| |
| std::string expected_string = R"([ |
| [[["key1", 100], ["key2", 200]]], |
| [[["a", 1], ["b", 2], ["c", 3]]] |
| ])"; |
| |
| WriteAndVerify(schema, expected_string); |
| } |
| |
| TEST_P(AvroReaderParameterizedTest, MapTypeWithNonStringKey) { |
| auto schema = std::make_shared<iceberg::Schema>( |
| std::vector<SchemaField>{SchemaField::MakeRequired( |
| 1, "int_map", |
| std::make_shared<MapType>( |
| SchemaField::MakeRequired(2, "key", std::make_shared<IntType>()), |
| SchemaField::MakeRequired(3, "value", std::make_shared<StringType>())))}); |
| |
| std::string expected_string = R"([ |
| [[[1, "one"], [2, "two"], [3, "three"]]], |
| [[[10, "ten"], [20, "twenty"]]] |
| ])"; |
| |
| WriteAndVerify(schema, expected_string); |
| } |
| |
| TEST_F(AvroReaderTest, ProjectionSubsetAndReorder) { |
| // Write file with full schema |
| auto write_schema = std::make_shared<iceberg::Schema>(std::vector<SchemaField>{ |
| SchemaField::MakeRequired(1, "id", std::make_shared<IntType>()), |
| SchemaField::MakeRequired(2, "name", std::make_shared<StringType>()), |
| SchemaField::MakeRequired(3, "age", std::make_shared<IntType>()), |
| SchemaField::MakeRequired(4, "city", std::make_shared<StringType>())}); |
| |
| std::string write_data = R"([ |
| [1, "Alice", 25, "NYC"], |
| [2, "Bob", 30, "SF"], |
| [3, "Charlie", 35, "LA"] |
| ])"; |
| |
| // Write with full schema |
| ArrowSchema arrow_c_schema; |
| ASSERT_THAT(ToArrowSchema(*write_schema, &arrow_c_schema), IsOk()); |
| auto arrow_schema_result = ::arrow::ImportType(&arrow_c_schema); |
| ASSERT_TRUE(arrow_schema_result.ok()); |
| auto arrow_schema = arrow_schema_result.ValueOrDie(); |
| |
| auto array_result = ::arrow::json::ArrayFromJSONString(arrow_schema, write_data); |
| ASSERT_TRUE(array_result.ok()); |
| auto array = array_result.ValueOrDie(); |
| |
| struct ArrowArray arrow_array; |
| auto export_result = ::arrow::ExportArray(*array, &arrow_array); |
| ASSERT_TRUE(export_result.ok()); |
| |
| std::unordered_map<std::string, std::string> metadata = {{"k1", "v1"}}; |
| auto writer_result = |
| WriterFactoryRegistry::Open(FileFormatType::kAvro, {.path = temp_avro_file_, |
| .schema = write_schema, |
| .io = file_io_, |
| .metadata = metadata}); |
| ASSERT_TRUE(writer_result.has_value()); |
| auto writer = std::move(writer_result.value()); |
| ASSERT_THAT(writer->Write(&arrow_array), IsOk()); |
| ASSERT_THAT(writer->Close(), IsOk()); |
| ICEBERG_UNWRAP_OR_FAIL(auto written_length, writer->length()); |
| |
| // Read with projected schema: subset of columns (city, id) in different order |
| auto read_schema = std::make_shared<iceberg::Schema>(std::vector<SchemaField>{ |
| SchemaField::MakeRequired(4, "city", std::make_shared<StringType>()), |
| SchemaField::MakeRequired(1, "id", std::make_shared<IntType>())}); |
| |
| auto reader_result = |
| ReaderFactoryRegistry::Open(FileFormatType::kAvro, {.path = temp_avro_file_, |
| .length = written_length, |
| .io = file_io_, |
| .projection = read_schema}); |
| ASSERT_THAT(reader_result, IsOk()); |
| auto reader = std::move(reader_result.value()); |
| |
| // Verify reordered subset |
| ASSERT_NO_FATAL_FAILURE( |
| VerifyNextBatch(*reader, R"([["NYC", 1], ["SF", 2], ["LA", 3]])")); |
| ASSERT_NO_FATAL_FAILURE(VerifyExhausted(*reader)); |
| } |
| |
| TEST_P(AvroReaderParameterizedTest, ComplexNestedTypes) { |
| auto schema = std::make_shared<iceberg::Schema>(std::vector<SchemaField>{ |
| SchemaField::MakeRequired(1, "id", std::make_shared<IntType>()), |
| SchemaField::MakeRequired(2, "nested_list", |
| std::make_shared<ListType>(SchemaField::MakeRequired( |
| 3, "element", |
| std::make_shared<ListType>(SchemaField::MakeRequired( |
| 4, "element", std::make_shared<IntType>())))))}); |
| |
| std::string expected_string = R"([ |
| [1, [[1, 2], [3, 4]]], |
| [2, [[5], [6, 7, 8]]] |
| ])"; |
| |
| WriteAndVerify(schema, expected_string); |
| } |
| |
| TEST_P(AvroReaderParameterizedTest, OptionalFieldsWithNulls) { |
| auto schema = std::make_shared<iceberg::Schema>(std::vector<SchemaField>{ |
| SchemaField::MakeRequired(1, "id", std::make_shared<IntType>()), |
| SchemaField::MakeOptional(2, "name", std::make_shared<StringType>()), |
| SchemaField::MakeOptional(3, "age", std::make_shared<IntType>())}); |
| |
| std::string expected_string = R"([ |
| [1, "Alice", 30], |
| [2, null, 25], |
| [3, "Charlie", null], |
| [4, null, null] |
| ])"; |
| |
| WriteAndVerify(schema, expected_string); |
| } |
| |
| TEST_P(AvroReaderParameterizedTest, LargeDataset) { |
| auto schema = std::make_shared<iceberg::Schema>(std::vector<SchemaField>{ |
| SchemaField::MakeRequired(1, "id", std::make_shared<LongType>()), |
| SchemaField::MakeRequired(2, "value", std::make_shared<DoubleType>())}); |
| |
| // Generate large dataset JSON |
| std::ostringstream json; |
| json << "["; |
| for (int i = 0; i < 1000; ++i) { |
| if (i > 0) json << ", "; |
| json << "[" << i << ", " << (i * 1.5) << "]"; |
| } |
| json << "]"; |
| |
| WriteAndVerify(schema, json.str()); |
| } |
| |
| TEST_P(AvroReaderParameterizedTest, EmptyCollections) { |
| auto schema = std::make_shared<iceberg::Schema>(std::vector<SchemaField>{ |
| SchemaField::MakeRequired(1, "id", std::make_shared<IntType>()), |
| SchemaField::MakeRequired(2, "list_col", |
| std::make_shared<ListType>(SchemaField::MakeRequired( |
| 3, "element", std::make_shared<IntType>())))}); |
| |
| std::string expected_string = R"([ |
| [1, []], |
| [2, [10, 20, 30]] |
| ])"; |
| |
| WriteAndVerify(schema, expected_string); |
| } |
| |
| TEST_P(AvroReaderParameterizedTest, ReadFilePathColumn) { |
| temp_avro_file_ = "avro_metadata_path.avro"; |
| CreateSimpleAvroFile(); |
| |
| // Create schema with _path metadata column |
| auto schema = std::make_shared<Schema>(std::vector<SchemaField>{ |
| SchemaField::MakeRequired(1, "id", int32()), |
| MetadataColumns::kFilePath, |
| }); |
| |
| const std::string expected_string = R"([ |
| [1, "avro_metadata_path.avro"], |
| [2, "avro_metadata_path.avro"], |
| [3, "avro_metadata_path.avro"] |
| ])"; |
| |
| ICEBERG_UNWRAP_OR_FAIL( |
| auto reader, ReaderFactoryRegistry::Open( |
| FileFormatType::kAvro, |
| {.path = temp_avro_file_, .io = file_io_, .projection = schema})); |
| |
| VerifyNextBatch(*reader, expected_string); |
| } |
| |
| TEST_P(AvroReaderParameterizedTest, ReadRowPositionColumn) { |
| temp_avro_file_ = "avro_metadata_pos.avro"; |
| CreateSimpleAvroFile(); |
| |
| // Create schema with _pos metadata column |
| auto schema = std::make_shared<Schema>(std::vector<SchemaField>{ |
| SchemaField::MakeRequired(1, "id", int32()), |
| MetadataColumns::kRowPosition, |
| }); |
| |
| const std::string expected_string = R"([[1, 0], [2, 1], [3, 2]])"; |
| |
| ICEBERG_UNWRAP_OR_FAIL( |
| auto reader, ReaderFactoryRegistry::Open( |
| FileFormatType::kAvro, |
| {.path = temp_avro_file_, .io = file_io_, .projection = schema})); |
| } |
| |
| TEST_P(AvroReaderParameterizedTest, ReadBothMetadataColumns) { |
| temp_avro_file_ = "avro_metadata_path_pos.avro"; |
| CreateSimpleAvroFile(); |
| |
| // Create schema with both _file and _pos metadata columns |
| auto schema = std::make_shared<Schema>(std::vector<SchemaField>{ |
| SchemaField::MakeRequired(1, "id", int32()), |
| MetadataColumns::kFilePath, |
| MetadataColumns::kRowPosition, |
| }); |
| |
| const std::string expected_string = R"([ |
| [1, "avro_metadata_path_pos.avro", 0], |
| [2, "avro_metadata_path_pos.avro", 1], |
| [3, "avro_metadata_path_pos.avro", 2] |
| ])"; |
| |
| ICEBERG_UNWRAP_OR_FAIL( |
| auto reader, ReaderFactoryRegistry::Open( |
| FileFormatType::kAvro, |
| {.path = temp_avro_file_, .io = file_io_, .projection = schema})); |
| |
| VerifyNextBatch(*reader, expected_string); |
| } |
| |
| TEST_P(AvroReaderParameterizedTest, ReadMetadataOnlyProjection) { |
| temp_avro_file_ = "avro_metadata_only.avro"; |
| CreateSimpleAvroFile(); |
| |
| // Create schema with only metadata columns (no data columns) |
| auto schema = std::make_shared<Schema>(std::vector<SchemaField>{ |
| MetadataColumns::kFilePath, |
| MetadataColumns::kRowPosition, |
| }); |
| |
| const std::string expected_string = R"([ |
| ["avro_metadata_only.avro", 0], |
| ["avro_metadata_only.avro", 1], |
| ["avro_metadata_only.avro", 2] |
| ])"; |
| |
| ICEBERG_UNWRAP_OR_FAIL( |
| auto reader, ReaderFactoryRegistry::Open( |
| FileFormatType::kAvro, |
| {.path = temp_avro_file_, .io = file_io_, .projection = schema})); |
| |
| VerifyNextBatch(*reader, expected_string); |
| } |
| |
| TEST_P(AvroReaderParameterizedTest, SplitWithRowPositionNotSupported) { |
| CreateSimpleAvroFile(); |
| |
| // Create schema with _pos metadata column |
| auto schema = std::make_shared<Schema>(std::vector<SchemaField>{ |
| SchemaField::MakeRequired(1, "id", int32()), |
| MetadataColumns::kRowPosition, |
| }); |
| |
| auto reader_result = ReaderFactoryRegistry::Open( |
| FileFormatType::kAvro, {.path = temp_avro_file_, |
| .split = Split{.offset = 100, .length = 200}, |
| .io = file_io_, |
| .projection = schema}); |
| |
| ASSERT_THAT(reader_result, IsError(ErrorKind::kNotSupported)); |
| EXPECT_THAT(reader_result, |
| HasErrorMessage("'_pos' metadata column with split is not supported")); |
| } |
| |
| INSTANTIATE_TEST_SUITE_P(DirectDecoderModes, AvroReaderParameterizedTest, |
| ::testing::Bool(), |
| [](const ::testing::TestParamInfo<bool>& info) { |
| return info.param ? "DirectDecoder" : "GenericDatum"; |
| }); |
| |
| // Parameterized test fixture for testing both direct encoder and GenericDatum modes |
| class AvroWriterTest : public ::testing::Test, |
| public ::testing::WithParamInterface<bool> { |
| protected: |
| static void SetUpTestSuite() { RegisterAll(); } |
| |
| void SetUp() override { |
| file_io_ = arrow::ArrowFileSystemFileIO::MakeMockFileIO(); |
| temp_avro_file_ = "avro_writer_test.avro"; |
| skip_datum_ = GetParam(); |
| } |
| |
| void WriteAvroFile( |
| std::shared_ptr<Schema> schema, const std::string& json_data, |
| const std::unordered_map<std::string, std::string>& extra_properties = {}) { |
| ArrowSchema arrow_c_schema; |
| ASSERT_THAT(ToArrowSchema(*schema, &arrow_c_schema), IsOk()); |
| |
| auto arrow_schema_result = ::arrow::ImportType(&arrow_c_schema); |
| ASSERT_TRUE(arrow_schema_result.ok()); |
| auto arrow_schema = arrow_schema_result.ValueOrDie(); |
| |
| auto array_result = ::arrow::json::ArrayFromJSONString(arrow_schema, json_data); |
| ASSERT_TRUE(array_result.ok()); |
| auto array = array_result.ValueOrDie(); |
| |
| struct ArrowArray arrow_array; |
| auto export_result = ::arrow::ExportArray(*array, &arrow_array); |
| ASSERT_TRUE(export_result.ok()); |
| |
| std::unordered_map<std::string, std::string> metadata = { |
| {"writer_test", "direct_encoder"}}; |
| |
| WriterProperties writer_properties; |
| writer_properties.Set(WriterProperties::kAvroSkipDatum, skip_datum_); |
| for (const auto& [key, value] : extra_properties) { |
| writer_properties.mutable_configs().emplace(key, value); |
| } |
| |
| auto writer_result = WriterFactoryRegistry::Open( |
| FileFormatType::kAvro, {.path = temp_avro_file_, |
| .schema = schema, |
| .io = file_io_, |
| .metadata = metadata, |
| .properties = std::move(writer_properties)}); |
| ASSERT_TRUE(writer_result.has_value()); |
| writer_ = std::move(writer_result.value()); |
| ASSERT_THAT(writer_->Write(&arrow_array), IsOk()); |
| ASSERT_THAT(writer_->Close(), IsOk()); |
| write_schema_ = schema; |
| } |
| |
| void VerifyNextBatch(Reader& reader, std::string_view expected_json) { |
| // Boilerplate to get Arrow schema |
| auto schema_result = reader.Schema(); |
| ASSERT_THAT(schema_result, IsOk()); |
| auto arrow_c_schema = std::move(schema_result.value()); |
| auto import_schema_result = ::arrow::ImportType(&arrow_c_schema); |
| auto arrow_schema = import_schema_result.ValueOrDie(); |
| |
| // Boilerplate to get Arrow array |
| auto data = reader.Next(); |
| ASSERT_THAT(data, IsOk()); |
| ASSERT_TRUE(data.value().has_value()); |
| auto arrow_c_array = data.value().value(); |
| auto data_result = ::arrow::ImportArray(&arrow_c_array, arrow_schema); |
| auto arrow_array = data_result.ValueOrDie(); |
| |
| // Verify data |
| auto expected_array = |
| ::arrow::json::ArrayFromJSONString(arrow_schema, expected_json).ValueOrDie(); |
| ASSERT_TRUE(arrow_array->Equals(*expected_array)); |
| } |
| |
| void VerifyExhausted(Reader& reader) { |
| auto data = reader.Next(); |
| ASSERT_THAT(data, IsOk()); |
| ASSERT_FALSE(data.value().has_value()); |
| } |
| |
| void VerifyWrittenData(const std::string& expected_json) { |
| ICEBERG_UNWRAP_OR_FAIL(auto written_length, writer_->length()); |
| |
| auto reader_result = |
| ReaderFactoryRegistry::Open(FileFormatType::kAvro, {.path = temp_avro_file_, |
| .length = written_length, |
| .io = file_io_, |
| .projection = write_schema_}); |
| ASSERT_THAT(reader_result, IsOk()); |
| auto reader = std::move(reader_result.value()); |
| ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader, expected_json)); |
| ASSERT_NO_FATAL_FAILURE(VerifyExhausted(*reader)); |
| } |
| |
| ::avro::ValidSchema PhysicalAvroSchema() { |
| auto& mock_io = internal::checked_cast<arrow::ArrowFileSystemFileIO&>(*file_io_); |
| auto input = mock_io.fs()->OpenInputFile(temp_avro_file_).ValueOrDie(); |
| auto input_stream = std::make_unique<AvroInputStream>(std::move(input), 1024 * 1024); |
| ::avro::DataFileReader<::avro::GenericDatum> avro_reader(std::move(input_stream)); |
| return avro_reader.dataSchema(); |
| } |
| |
| std::shared_ptr<FileIO> file_io_; |
| std::string temp_avro_file_; |
| bool skip_datum_{true}; |
| std::shared_ptr<Schema> write_schema_; |
| std::unique_ptr<Writer> writer_; |
| }; |
| |
| TEST_P(AvroWriterTest, WritePrimitiveTypes) { |
| auto schema = std::make_shared<iceberg::Schema>(std::vector<SchemaField>{ |
| SchemaField::MakeRequired(1, "bool_col", std::make_shared<BooleanType>()), |
| SchemaField::MakeRequired(2, "int_col", std::make_shared<IntType>()), |
| SchemaField::MakeRequired(3, "long_col", std::make_shared<LongType>()), |
| SchemaField::MakeRequired(4, "float_col", std::make_shared<FloatType>()), |
| SchemaField::MakeRequired(5, "double_col", std::make_shared<DoubleType>()), |
| SchemaField::MakeRequired(6, "string_col", std::make_shared<StringType>())}); |
| |
| std::string test_data = R"([ |
| [true, 42, 1234567890, 3.14, 2.71828, "hello"], |
| [false, -100, -9876543210, -1.5, 0.0, "world"] |
| ])"; |
| |
| WriteAvroFile(schema, test_data); |
| VerifyWrittenData(test_data); |
| } |
| |
| TEST_P(AvroWriterTest, WriteTemporalTypes) { |
| auto schema = std::make_shared<iceberg::Schema>(std::vector<SchemaField>{ |
| SchemaField::MakeRequired(1, "date_col", std::make_shared<DateType>()), |
| SchemaField::MakeRequired(2, "time_col", std::make_shared<TimeType>()), |
| SchemaField::MakeRequired(3, "timestamp_col", std::make_shared<TimestampType>())}); |
| |
| std::string test_data = R"([ |
| [18628, 43200000000, 1640995200000000], |
| [18629, 86399000000, 1641081599000000] |
| ])"; |
| |
| WriteAvroFile(schema, test_data); |
| VerifyWrittenData(test_data); |
| } |
| |
| TEST_P(AvroWriterTest, WriteNestedStruct) { |
| auto schema = std::make_shared<iceberg::Schema>(std::vector<SchemaField>{ |
| SchemaField::MakeRequired(1, "id", std::make_shared<IntType>()), |
| SchemaField::MakeRequired( |
| 2, "person", |
| std::make_shared<iceberg::StructType>(std::vector<SchemaField>{ |
| SchemaField::MakeRequired(3, "name", std::make_shared<StringType>()), |
| SchemaField::MakeRequired(4, "age", std::make_shared<IntType>())}))}); |
| |
| std::string test_data = R"([ |
| [1, ["Alice", 30]], |
| [2, ["Bob", 25]] |
| ])"; |
| |
| WriteAvroFile(schema, test_data); |
| VerifyWrittenData(test_data); |
| } |
| |
| TEST_P(AvroWriterTest, WriteListType) { |
| auto schema = std::make_shared<iceberg::Schema>(std::vector<SchemaField>{ |
| SchemaField::MakeRequired(1, "id", std::make_shared<IntType>()), |
| SchemaField::MakeRequired(2, "tags", |
| std::make_shared<ListType>(SchemaField::MakeRequired( |
| 3, "element", std::make_shared<StringType>())))}); |
| |
| std::string test_data = R"([ |
| [1, ["tag1", "tag2", "tag3"]], |
| [2, ["foo", "bar"]], |
| [3, []] |
| ])"; |
| |
| WriteAvroFile(schema, test_data); |
| VerifyWrittenData(test_data); |
| } |
| |
| TEST_P(AvroWriterTest, WriteMapTypeWithStringKey) { |
| auto schema = std::make_shared<iceberg::Schema>( |
| std::vector<SchemaField>{SchemaField::MakeRequired( |
| 1, "properties", |
| std::make_shared<MapType>( |
| SchemaField::MakeRequired(2, "key", std::make_shared<StringType>()), |
| SchemaField::MakeRequired(3, "value", std::make_shared<IntType>())))}); |
| |
| std::string test_data = R"([ |
| [[["key1", 100], ["key2", 200]]], |
| [[["a", 1], ["b", 2], ["c", 3]]] |
| ])"; |
| |
| WriteAvroFile(schema, test_data); |
| VerifyWrittenData(test_data); |
| } |
| |
| TEST_P(AvroWriterTest, WriteMapTypeWithNonStringKey) { |
| auto schema = std::make_shared<iceberg::Schema>( |
| std::vector<SchemaField>{SchemaField::MakeRequired( |
| 1, "int_map", |
| std::make_shared<MapType>( |
| SchemaField::MakeRequired(2, "key", std::make_shared<IntType>()), |
| SchemaField::MakeRequired(3, "value", std::make_shared<StringType>())))}); |
| |
| std::string test_data = R"([ |
| [[[1, "one"], [2, "two"], [3, "three"]]], |
| [[[10, "ten"], [20, "twenty"]]] |
| ])"; |
| |
| WriteAvroFile(schema, test_data); |
| VerifyWrittenData(test_data); |
| } |
| |
| TEST_P(AvroWriterTest, WriteEmptyMaps) { |
| auto schema = std::make_shared<iceberg::Schema>(std::vector<SchemaField>{ |
| SchemaField::MakeRequired( |
| 1, "string_map", |
| std::make_shared<MapType>( |
| SchemaField::MakeRequired(2, "key", std::make_shared<StringType>()), |
| SchemaField::MakeRequired(3, "value", std::make_shared<IntType>()))), |
| SchemaField::MakeRequired( |
| 4, "int_map", |
| std::make_shared<MapType>( |
| SchemaField::MakeRequired(5, "key", std::make_shared<IntType>()), |
| SchemaField::MakeRequired(6, "value", std::make_shared<StringType>())))}); |
| |
| // Test empty maps for both string and non-string keys |
| std::string test_data = R"([ |
| [[], []], |
| [[["a", 1]], []] |
| ])"; |
| |
| // Just verify writing succeeds (empty maps are handled correctly by the encoder) |
| WriteAvroFile(schema, test_data); |
| VerifyWrittenData(test_data); |
| } |
| |
| TEST_P(AvroWriterTest, WriteOptionalFields) { |
| auto schema = std::make_shared<iceberg::Schema>(std::vector<SchemaField>{ |
| SchemaField::MakeRequired(1, "id", std::make_shared<IntType>()), |
| SchemaField::MakeOptional(2, "name", std::make_shared<StringType>()), |
| SchemaField::MakeOptional(3, "age", std::make_shared<IntType>())}); |
| |
| std::string test_data = R"([ |
| [1, "Alice", 30], |
| [2, null, 25], |
| [3, "Charlie", null], |
| [4, null, null] |
| ])"; |
| |
| WriteAvroFile(schema, test_data); |
| VerifyWrittenData(test_data); |
| } |
| |
| TEST_P(AvroWriterTest, WritesUnknownFieldsAsAvroNull) { |
| auto schema = std::make_shared<iceberg::Schema>(std::vector<SchemaField>{ |
| SchemaField::MakeOptional(1, "id", int32()), |
| SchemaField::MakeOptional(2, "mystery", unknown()), |
| SchemaField::MakeOptional(3, "profile", |
| std::make_shared<StructType>(std::vector<SchemaField>{ |
| SchemaField::MakeOptional(4, "name", string()), |
| SchemaField::MakeOptional(5, "secret", unknown()), |
| })), |
| }); |
| |
| std::string test_data = R"([ |
| [1, null, {"name": "Person0", "secret": null}], |
| [2, null, {"name": "Person1", "secret": null}] |
| ])"; |
| |
| WriteAvroFile(schema, test_data); |
| |
| auto avro_schema = PhysicalAvroSchema(); |
| auto root = avro_schema.root(); |
| ASSERT_EQ(root->type(), ::avro::AVRO_RECORD); |
| // Unknown fields are written as AVRO_NULL, not pruned. |
| ASSERT_EQ(root->leaves(), 3); |
| EXPECT_EQ(root->nameAt(0), "id"); |
| EXPECT_EQ(FieldIdAt(root, 0), std::make_optional(1)); |
| EXPECT_EQ(root->nameAt(1), "mystery"); |
| EXPECT_EQ(root->leafAt(1)->type(), ::avro::AVRO_NULL); |
| EXPECT_EQ(FieldIdAt(root, 1), std::make_optional(2)); |
| EXPECT_EQ(root->nameAt(2), "profile"); |
| EXPECT_EQ(FieldIdAt(root, 2), std::make_optional(3)); |
| |
| auto profile = UnwrapOptional(root->leafAt(2)); |
| ASSERT_EQ(profile->type(), ::avro::AVRO_RECORD); |
| ASSERT_EQ(profile->leaves(), 2); |
| EXPECT_EQ(profile->nameAt(0), "name"); |
| EXPECT_EQ(FieldIdAt(profile, 0), std::make_optional(4)); |
| EXPECT_EQ(profile->nameAt(1), "secret"); |
| EXPECT_EQ(profile->leafAt(1)->type(), ::avro::AVRO_NULL); |
| EXPECT_EQ(FieldIdAt(profile, 1), std::make_optional(5)); |
| |
| VerifyWrittenData(test_data); |
| } |
| |
| TEST_P(AvroWriterTest, WritesUnknownListElementsAndMapValues) { |
| auto schema = std::make_shared<iceberg::Schema>(std::vector<SchemaField>{ |
| SchemaField::MakeRequired(1, "id", int32()), |
| SchemaField::MakeRequired(2, "mysteries", |
| std::make_shared<ListType>(SchemaField::MakeOptional( |
| 3, ListType::kElementName, unknown()))), |
| SchemaField::MakeRequired( |
| 4, "properties", |
| std::make_shared<MapType>( |
| SchemaField::MakeRequired(5, MapType::kKeyName, string()), |
| SchemaField::MakeOptional(6, MapType::kValueName, unknown()))), |
| }); |
| |
| std::string test_data = R"([ |
| [1, [null, null], [["a", null], ["b", null]]], |
| [2, [], []], |
| [3, [null], [["c", null]]] |
| ])"; |
| |
| WriteAvroFile(schema, test_data); |
| |
| auto avro_schema = PhysicalAvroSchema(); |
| auto root = avro_schema.root(); |
| ASSERT_EQ(root->type(), ::avro::AVRO_RECORD); |
| ASSERT_EQ(root->leaves(), 3); |
| |
| auto mysteries = root->leafAt(1); |
| ASSERT_EQ(mysteries->type(), ::avro::AVRO_ARRAY); |
| ASSERT_EQ(mysteries->leaves(), 1); |
| EXPECT_EQ(mysteries->leafAt(0)->type(), ::avro::AVRO_NULL); |
| |
| auto properties = root->leafAt(2); |
| ASSERT_EQ(properties->type(), ::avro::AVRO_MAP); |
| ASSERT_EQ(properties->leaves(), 2); |
| EXPECT_EQ(properties->leafAt(1)->type(), ::avro::AVRO_NULL); |
| |
| VerifyWrittenData(test_data); |
| } |
| |
| TEST_P(AvroWriterTest, WritesUnknownFieldsNestedInsideListOrMapStructs) { |
| auto schema = std::make_shared<iceberg::Schema>(std::vector<SchemaField>{ |
| SchemaField::MakeOptional(1, "id", int32()), |
| SchemaField::MakeOptional(2, "events", |
| std::make_shared<ListType>(SchemaField::MakeOptional( |
| 3, ListType::kElementName, |
| std::make_shared<StructType>(std::vector<SchemaField>{ |
| SchemaField::MakeOptional(4, "name", string()), |
| SchemaField::MakeOptional(5, "secret", unknown()), |
| })))), |
| SchemaField::MakeOptional( |
| 6, "properties", |
| std::make_shared<MapType>( |
| SchemaField::MakeRequired(7, MapType::kKeyName, iceberg::string()), |
| SchemaField::MakeOptional( |
| 8, MapType::kValueName, |
| std::make_shared<StructType>(std::vector<SchemaField>{ |
| SchemaField::MakeOptional(9, "label", string()), |
| SchemaField::MakeOptional(10, "secret", unknown()), |
| })))), |
| }); |
| |
| std::string test_data = R"([ |
| [1, [{"name": "open", "secret": null}, {"name": "close", "secret": null}], [["a", {"label": "A", "secret": null}]]], |
| [2, [], []] |
| ])"; |
| |
| WriteAvroFile(schema, test_data); |
| |
| auto avro_schema = PhysicalAvroSchema(); |
| auto root = avro_schema.root(); |
| ASSERT_EQ(root->type(), ::avro::AVRO_RECORD); |
| ASSERT_EQ(root->leaves(), 3); |
| |
| auto events = UnwrapOptional(root->leafAt(1)); |
| ASSERT_EQ(events->type(), ::avro::AVRO_ARRAY); |
| auto event = UnwrapOptional(events->leafAt(0)); |
| ASSERT_EQ(event->type(), ::avro::AVRO_RECORD); |
| ASSERT_EQ(event->leaves(), 2); |
| EXPECT_EQ(event->nameAt(0), "name"); |
| EXPECT_EQ(FieldIdAt(event, 0), std::make_optional(4)); |
| EXPECT_EQ(event->nameAt(1), "secret"); |
| EXPECT_EQ(event->leafAt(1)->type(), ::avro::AVRO_NULL); |
| EXPECT_EQ(FieldIdAt(event, 1), std::make_optional(5)); |
| |
| auto properties = UnwrapOptional(root->leafAt(2)); |
| ASSERT_EQ(properties->type(), ::avro::AVRO_MAP); |
| ASSERT_EQ(properties->leaves(), 2); |
| auto value = UnwrapOptional(properties->leafAt(1)); |
| ASSERT_EQ(value->type(), ::avro::AVRO_RECORD); |
| ASSERT_EQ(value->leaves(), 2); |
| EXPECT_EQ(value->nameAt(0), "label"); |
| EXPECT_EQ(FieldIdAt(value, 0), std::make_optional(9)); |
| EXPECT_EQ(value->nameAt(1), "secret"); |
| EXPECT_EQ(value->leafAt(1)->type(), ::avro::AVRO_NULL); |
| EXPECT_EQ(FieldIdAt(value, 1), std::make_optional(10)); |
| |
| VerifyWrittenData(test_data); |
| } |
| |
| TEST_P(AvroWriterTest, WriteLargeDataset) { |
| auto schema = std::make_shared<iceberg::Schema>(std::vector<SchemaField>{ |
| SchemaField::MakeRequired(1, "id", std::make_shared<LongType>()), |
| SchemaField::MakeRequired(2, "value", std::make_shared<DoubleType>())}); |
| |
| // Generate large dataset JSON |
| std::ostringstream json; |
| json << "["; |
| for (int i = 0; i < 1000; ++i) { |
| if (i > 0) json << ", "; |
| json << "[" << i << ", " << (i * 1.5) << "]"; |
| } |
| json << "]"; |
| |
| WriteAvroFile(schema, json.str()); |
| VerifyWrittenData(json.str()); |
| } |
| |
| TEST_P(AvroWriterTest, MultipleAvroBlocks) { |
| auto schema = std::make_shared<Schema>( |
| std::vector<SchemaField>{SchemaField::MakeRequired(1, "id", int32()), |
| SchemaField::MakeRequired(2, "name", string())}); |
| |
| const std::string json_data = R"([ |
| [1, "Alice_with_a_very_long_name_to_exceed_sync_interval"], |
| [2, "Bob_with_another_very_long_name_to_exceed_sync_interval"], |
| [3, "Charlie_with_yet_another_very_long_name_to_exceed_sync"], |
| [4, "David_with_a_super_long_name_that_will_exceed_interval"], |
| [5, "Eve_with_an_extremely_long_name_to_force_new_block_here"] |
| ])"; |
| |
| const std::vector<std::pair</*sync_interval*/ std::string, /*num_blocks*/ size_t>> |
| test_cases = {{"32", 5}, {"65536", 1}}; |
| |
| for (const auto& [interval, num_blocks] : test_cases) { |
| WriteAvroFile(schema, json_data, |
| {{WriterProperties::kAvroSyncInterval.key(), interval}}); |
| VerifyWrittenData(json_data); |
| |
| // Use raw avro-cpp reader to count blocks by tracking previousSync() changes |
| auto mock_io = internal::checked_pointer_cast<arrow::ArrowFileSystemFileIO>(file_io_); |
| auto input = mock_io->fs()->OpenInputFile(temp_avro_file_).ValueOrDie(); |
| auto input_stream = std::make_unique<AvroInputStream>(std::move(input), 1024 * 1024); |
| ::avro::DataFileReader<::avro::GenericDatum> avro_reader(std::move(input_stream)); |
| ::avro::GenericDatum datum(avro_reader.dataSchema()); |
| |
| size_t block_count = 0; |
| int64_t last_sync = -1; |
| |
| while (avro_reader.read(datum)) { |
| if (int64_t current_sync = avro_reader.previousSync(); current_sync != last_sync) { |
| block_count++; |
| last_sync = current_sync; |
| } |
| } |
| |
| ASSERT_EQ(block_count, num_blocks); |
| } |
| } |
| |
| TEST_P(AvroWriterTest, Metrics) { |
| auto schema = std::make_shared<iceberg::Schema>(std::vector<SchemaField>{ |
| SchemaField::MakeRequired(1, "id", std::make_shared<IntType>()), |
| SchemaField::MakeOptional(2, "name", std::make_shared<StringType>())}); |
| |
| std::string test_data = R"([[1, "Alice"], [2, "Bob"], [3, "Charlie"]])"; |
| |
| // Write data but don't close yet |
| ArrowSchema arrow_c_schema; |
| ASSERT_THAT(ToArrowSchema(*schema, &arrow_c_schema), IsOk()); |
| auto arrow_schema = ::arrow::ImportType(&arrow_c_schema).ValueOrDie(); |
| auto array = ::arrow::json::ArrayFromJSONString(arrow_schema, test_data).ValueOrDie(); |
| struct ArrowArray arrow_array; |
| ASSERT_TRUE(::arrow::ExportArray(*array, &arrow_array).ok()); |
| |
| ICEBERG_UNWRAP_OR_FAIL( |
| writer_, |
| WriterFactoryRegistry::Open( |
| FileFormatType::kAvro, |
| {.path = temp_avro_file_, .schema = schema, .io = file_io_, .properties = {}})); |
| ASSERT_THAT(writer_->Write(&arrow_array), IsOk()); |
| |
| // Metrics should fail before close |
| ASSERT_THAT(writer_->metrics(), IsError(ErrorKind::kInvalid)); |
| |
| // After close, metrics should succeed |
| ASSERT_THAT(writer_->Close(), IsOk()); |
| ICEBERG_UNWRAP_OR_FAIL(auto metrics, writer_->metrics()); |
| ASSERT_TRUE(metrics.row_count.has_value()); |
| EXPECT_EQ(metrics.row_count.value(), 3); |
| EXPECT_TRUE(metrics.column_sizes.empty()); |
| EXPECT_TRUE(metrics.value_counts.empty()); |
| EXPECT_TRUE(metrics.null_value_counts.empty()); |
| EXPECT_TRUE(metrics.nan_value_counts.empty()); |
| EXPECT_TRUE(metrics.lower_bounds.empty()); |
| EXPECT_TRUE(metrics.upper_bounds.empty()); |
| } |
| |
| // Instantiate parameterized tests for both direct encoder and GenericDatum paths |
| INSTANTIATE_TEST_SUITE_P(DirectEncoderModes, AvroWriterTest, |
| ::testing::Values(true, false), |
| [](const ::testing::TestParamInfo<bool>& info) { |
| return info.param ? "DirectEncoder" : "GenericDatum"; |
| }); |
| |
| } // namespace iceberg::avro |
