diff --git a/src/iceberg/avro/avro_direct_decoder.cc b/src/iceberg/avro/avro_direct_decoder.cc index 23bb8ca7f..465bc4b24 100644 --- a/src/iceberg/avro/avro_direct_decoder.cc +++ b/src/iceberg/avro/avro_direct_decoder.cc @@ -160,7 +160,8 @@ Status DecodeStructToBuilder(const ::avro::NodePtr& avro_node, ::avro::Decoder& // Build a map from Avro field index to projection index (cached per struct schema) // -1 means the field should be skipped - const FieldProjection* cache_key = projections.data(); + // Empty projections all share one data() pointer, so the Avro node is part of the key. + const DecodeContext::ProjectionCacheKey cache_key{projections.data(), avro_node.get()}; auto cache_it = ctx.avro_to_projection_cache.find(cache_key); std::vector* avro_to_projection; @@ -168,18 +169,23 @@ Status DecodeStructToBuilder(const ::avro::NodePtr& avro_node, ::avro::Decoder& // Use cached mapping avro_to_projection = &cache_it->second; } else { - // Build and cache the mapping - auto [inserted_it, inserted] = ctx.avro_to_projection_cache.emplace( - cache_key, std::vector(avro_node->leaves(), -1)); - avro_to_projection = &inserted_it->second; - + // Build and validate the mapping before caching so a failure leaves no partial entry + std::vector mapping(avro_node->leaves(), -1); for (size_t proj_idx = 0; proj_idx < projections.size(); ++proj_idx) { const auto& field_projection = projections[proj_idx]; if (field_projection.kind == FieldProjection::Kind::kProjected) { size_t avro_field_index = std::get(field_projection.from); - (*avro_to_projection)[avro_field_index] = static_cast(proj_idx); + if (avro_field_index >= mapping.size()) { + return InvalidArgument( + "Projected field index {} out of range for Avro record with {} fields", + avro_field_index, mapping.size()); + } + mapping[avro_field_index] = static_cast(proj_idx); } } + avro_to_projection = + &ctx.avro_to_projection_cache.emplace(cache_key, std::move(mapping)) + .first->second; } // Read all Avro fields in order (must maintain decoder position) diff --git a/src/iceberg/avro/avro_direct_decoder_internal.h b/src/iceberg/avro/avro_direct_decoder_internal.h index d067b550e..1e3505fcf 100644 --- a/src/iceberg/avro/avro_direct_decoder_internal.h +++ b/src/iceberg/avro/avro_direct_decoder_internal.h @@ -19,6 +19,12 @@ #pragma once +#include +#include +#include +#include +#include + #include #include #include @@ -41,9 +47,17 @@ struct DecodeContext { // Scratch buffer for binary/fixed/uuid/decimal data (reused across rows) std::vector bytes_scratch; // Cache for avro field index to projection index mapping - // Key: pointer to projections array (identifies struct schema) + // Key: {projections array, Avro record node}; empty structs share a null data() // Value: vector mapping avro field index -> projection index (-1 if not projected) - std::unordered_map> avro_to_projection_cache; + using ProjectionCacheKey = std::pair; + struct ProjectionCacheKeyHash { + size_t operator()(const ProjectionCacheKey& key) const noexcept { + return std::hash{}(key.first) * 31 + + std::hash{}(key.second); + } + }; + std::unordered_map, ProjectionCacheKeyHash> + avro_to_projection_cache; }; /// \brief Directly decode Avro data to Arrow array builders without GenericDatum diff --git a/src/iceberg/test/avro_data_test.cc b/src/iceberg/test/avro_data_test.cc index 90ed4d037..9ec0221da 100644 --- a/src/iceberg/test/avro_data_test.cc +++ b/src/iceberg/test/avro_data_test.cc @@ -39,6 +39,7 @@ #include #include "iceberg/avro/avro_data_util_internal.h" +#include "iceberg/avro/avro_direct_decoder_internal.h" #include "iceberg/avro/avro_direct_encoder_internal.h" #include "iceberg/avro/avro_schema_util_internal.h" #include "iceberg/expression/literal.h" @@ -1900,6 +1901,55 @@ const std::vector kRoundTripTestCases = { }, }; +TEST(DirectDecoderTest, ProjectionCacheKeyedPerAvroNodeForEmptyStructs) { + auto make_struct = [](int32_t first_id, int count) { + std::vector fields; + for (int i = 0; i < count; ++i) { + fields.push_back(SchemaField::MakeRequired(first_id + i, "f" + std::to_string(i), + std::make_shared())); + } + return std::make_shared(std::move(fields)); + }; + Schema file_schema({SchemaField::MakeRequired(1, "narrow", make_struct(10, 1)), + SchemaField::MakeRequired(2, "wide", make_struct(20, 6))}); + Schema read_schema({SchemaField::MakeRequired(1, "narrow", make_struct(10, 0)), + SchemaField::MakeRequired(2, "wide", make_struct(20, 0))}); + + ::avro::NodePtr avro_node; + ASSERT_THAT(ToAvroNodeVisitor{}.Visit(file_schema, &avro_node), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto projection, + Project(read_schema, avro_node, /*prune_source=*/false)); + + auto out = ::avro::memoryOutputStream(); + auto encoder = ::avro::binaryEncoder(); + encoder->init(*out); + for (int i = 0; i < 7; ++i) { + encoder->encodeInt(i); + } + encoder->flush(); + auto in = ::avro::memoryInputStream(*out); + auto decoder = ::avro::binaryDecoder(); + decoder->init(*in); + + ArrowSchema arrow_c_schema; + ASSERT_THAT(ToArrowSchema(read_schema, &arrow_c_schema), IsOk()); + auto arrow_schema = ::arrow::ImportSchema(&arrow_c_schema).ValueOrDie(); + auto builder = + ::arrow::MakeBuilder(std::make_shared<::arrow::StructType>(arrow_schema->fields())) + .ValueOrDie(); + + DecodeContext ctx; + ASSERT_THAT(DecodeAvroToBuilder(avro_node, *decoder, projection, read_schema, + /*metadata_context=*/{}, builder.get(), ctx), + IsOk()); + + std::vector mapping_sizes; + for (const auto& [key, mapping] : ctx.avro_to_projection_cache) { + mapping_sizes.push_back(mapping.size()); + } + EXPECT_THAT(mapping_sizes, ::testing::UnorderedElementsAre(1, 6, 2)); +} + INSTANTIATE_TEST_SUITE_P(AllTypes, AvroRoundTripConversionTest, ::testing::ValuesIn(kRoundTripTestCases), [](const ::testing::TestParamInfo& info) { diff --git a/src/iceberg/test/avro_test.cc b/src/iceberg/test/avro_test.cc index a631b4a10..a16116410 100644 --- a/src/iceberg/test/avro_test.cc +++ b/src/iceberg/test/avro_test.cc @@ -18,6 +18,7 @@ */ #include +#include #include #include #include @@ -399,6 +400,54 @@ TEST_F(AvroReaderTest, ReadWithBatchSize) { ASSERT_NO_FATAL_FAILURE(VerifyExhausted(*reader)); } +TEST_F(AvroReaderTest, ProjectEmptyStructsOfDifferentWidths) { + auto make_struct = [](int32_t first_id, int count) { + std::vector fields; + for (int i = 0; i < count; ++i) { + fields.push_back(SchemaField::MakeRequired(first_id + i, std::format("f{}", i), + std::make_shared())); + } + return std::make_shared(std::move(fields)); + }; + auto write_schema = std::make_shared( + std::vector{SchemaField::MakeRequired(1, "narrow", make_struct(10, 1)), + SchemaField::MakeRequired(2, "wide", make_struct(20, 6))}); + + ArrowSchema arrow_c_schema; + ASSERT_THAT(ToArrowSchema(*write_schema, &arrow_c_schema), IsOk()); + auto arrow_schema = ::arrow::ImportType(&arrow_c_schema).ValueOrDie(); + auto array = + ::arrow::json::ArrayFromJSONString( + arrow_schema, R"([[[1], [1, 2, 3, 4, 5, 6]], [[7], [7, 8, 9, 10, 11, 12]]])") + .ValueOrDie(); + struct ArrowArray arrow_array; + ASSERT_TRUE(::arrow::ExportArray(*array, &arrow_array).ok()); + + auto writer_result = WriterFactoryRegistry::Open( + FileFormatType::kAvro, + {.path = temp_avro_file_, .schema = write_schema, .io = file_io_}); + ASSERT_THAT(writer_result, IsOk()); + auto writer = std::move(writer_result.value()); + ASSERT_THAT(writer->Write(&arrow_array), IsOk()); + ASSERT_THAT(writer->Close(), IsOk()); + + auto read_schema = std::make_shared( + std::vector{SchemaField::MakeRequired(1, "narrow", make_struct(10, 0)), + SchemaField::MakeRequired(2, "wide", make_struct(20, 0))}); + ReaderProperties reader_properties; + reader_properties.Set(ReaderProperties::kAvroSkipDatum, true); + auto reader_result = ReaderFactoryRegistry::Open( + FileFormatType::kAvro, {.path = temp_avro_file_, + .io = file_io_, + .projection = read_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"([[[], []], [[], []]])")); + ASSERT_NO_FATAL_FAILURE(VerifyExhausted(*reader)); +} + TEST_F(AvroReaderTest, BufferSizeConfiguration) { // Test default buffer size ReaderProperties properties1;