Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 13 additions & 7 deletions src/iceberg/avro/avro_direct_decoder.cc
Original file line number Diff line number Diff line change
Expand Up @@ -160,26 +160,32 @@ 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()};
Comment thread
gty404 marked this conversation as resolved.
auto cache_it = ctx.avro_to_projection_cache.find(cache_key);
std::vector<int>* avro_to_projection;

if (cache_it != ctx.avro_to_projection_cache.end()) {
// 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<int>(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<int> 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<size_t>(field_projection.from);
(*avro_to_projection)[avro_field_index] = static_cast<int>(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<int>(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)
Expand Down
18 changes: 16 additions & 2 deletions src/iceberg/avro/avro_direct_decoder_internal.h
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,12 @@

#pragma once

#include <cstddef>
#include <functional>
#include <unordered_map>
#include <utility>
#include <vector>

#include <arrow/array/builder_base.h>
#include <avro/Decoder.hh>
#include <avro/Node.hh>
Expand All @@ -41,9 +47,17 @@ struct DecodeContext {
// Scratch buffer for binary/fixed/uuid/decimal data (reused across rows)
std::vector<uint8_t> 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<const FieldProjection*, std::vector<int>> avro_to_projection_cache;
using ProjectionCacheKey = std::pair<const FieldProjection*, const ::avro::Node*>;
struct ProjectionCacheKeyHash {
size_t operator()(const ProjectionCacheKey& key) const noexcept {
return std::hash<const FieldProjection*>{}(key.first) * 31 +
std::hash<const ::avro::Node*>{}(key.second);
}
};
std::unordered_map<ProjectionCacheKey, std::vector<int>, ProjectionCacheKeyHash>
avro_to_projection_cache;
};

/// \brief Directly decode Avro data to Arrow array builders without GenericDatum
Expand Down
50 changes: 50 additions & 0 deletions src/iceberg/test/avro_data_test.cc
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@
#include <gtest/gtest.h>

#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"
Expand Down Expand Up @@ -1900,6 +1901,55 @@ const std::vector<RoundTripParam> kRoundTripTestCases = {
},
};

TEST(DirectDecoderTest, ProjectionCacheKeyedPerAvroNodeForEmptyStructs) {
auto make_struct = [](int32_t first_id, int count) {
std::vector<SchemaField> fields;
for (int i = 0; i < count; ++i) {
fields.push_back(SchemaField::MakeRequired(first_id + i, "f" + std::to_string(i),
std::make_shared<IntType>()));
}
return std::make_shared<StructType>(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<size_t> 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<RoundTripParam>& info) {
Expand Down
49 changes: 49 additions & 0 deletions src/iceberg/test/avro_test.cc
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
*/

#include <array>
#include <format>
#include <initializer_list>
#include <optional>
#include <sstream>
Expand Down Expand Up @@ -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<SchemaField> fields;
for (int i = 0; i < count; ++i) {
fields.push_back(SchemaField::MakeRequired(first_id + i, std::format("f{}", i),
std::make_shared<IntType>()));
}
return std::make_shared<StructType>(std::move(fields));
};
auto write_schema = std::make_shared<Schema>(
std::vector<SchemaField>{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<Schema>(
std::vector<SchemaField>{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;
Expand Down
Loading