Skip to content
Merged
Show file tree
Hide file tree
Changes from 15 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
21 changes: 16 additions & 5 deletions src/iceberg/avro/avro_reader.cc
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@
#include "iceberg/avro/avro_data_util_internal.h"
#include "iceberg/avro/avro_schema_util_internal.h"
#include "iceberg/avro/avro_stream_internal.h"
#include "iceberg/name_mapping.h"
#include "iceberg/schema_internal.h"
#include "iceberg/util/checked_cast.h"
#include "iceberg/util/macros.h"
Expand Down Expand Up @@ -96,16 +97,26 @@ class AvroReader::Impl {
// Create a base reader without setting reader schema to enable projection.
auto base_reader =
std::make_unique<::avro::DataFileReaderBase>(std::move(input_stream));
const ::avro::ValidSchema& file_schema = base_reader->dataSchema();
::avro::ValidSchema file_schema = base_reader->dataSchema();

// Validate field ids in the file schema.
HasIdVisitor has_id_visitor;
ICEBERG_RETURN_UNEXPECTED(has_id_visitor.Visit(file_schema));

if (has_id_visitor.HasNoIds()) {
// TODO(gangwu): support applying field-ids based on name mapping
return NotImplemented("Avro file schema has no field IDs");
}
if (!has_id_visitor.AllHaveIds()) {
// Apply field IDs based on name mapping if available
if (options.name_mapping) {
ICEBERG_ASSIGN_OR_RAISE(
auto new_root_node,
MakeAvroNodeWithFieldIds(file_schema.root(), *options.name_mapping));

// Update the file schema to use the new schema with field IDs
file_schema = ::avro::ValidSchema(new_root_node);
} else {
return InvalidSchema(
"Avro file schema has no field IDs and no name mapping provided");
}
} else if (!has_id_visitor.AllHaveIds()) {
return InvalidSchema("Not all fields in the Avro file schema have field IDs");
}

Expand Down
245 changes: 245 additions & 0 deletions src/iceberg/avro/avro_schema_util.cc
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,9 @@

#include "iceberg/avro/avro_register.h"
#include "iceberg/avro/avro_schema_util_internal.h"
#include "iceberg/avro/constants.h"
#include "iceberg/metadata_columns.h"
#include "iceberg/name_mapping.h"
#include "iceberg/schema.h"
#include "iceberg/schema_util_internal.h"
#include "iceberg/util/formatter.h"
Expand Down Expand Up @@ -773,4 +775,247 @@ Result<SchemaProjection> Project(const Schema& expected_schema,
return SchemaProjection{std::move(field_projection.children)};
}

namespace {

Result<::avro::NodePtr> CreateRecordNodeWithFieldIds(const ::avro::NodePtr& original_node,
const MappedField& field) {
auto new_record_node = std::make_shared<::avro::NodeRecord>();
new_record_node->setName(original_node->name());

for (size_t i = 0; i < original_node->leaves(); ++i) {
if (i >= original_node->names()) {
return InvalidSchema("Index {} is out of bounds for names (size: {})", i,
original_node->names());
}
const std::string& field_name = original_node->nameAt(i);
Comment thread
MisterRaindrop marked this conversation as resolved.
Comment thread
MisterRaindrop marked this conversation as resolved.
::avro::NodePtr field_node = original_node->leafAt(i);
Comment thread
MisterRaindrop marked this conversation as resolved.

// TODO(liuxiaoyu): Add support for case sensitivity in name matching.
// Try to find nested field by name
const MappedField* nested_field = nullptr;
if (field.nested_mapping) {
auto fields_span = field.nested_mapping->fields();
for (const auto& f : fields_span) {
if (f.names.find(field_name) != f.names.end()) {
nested_field = &f;
break;
}
}
}

if (!nested_field) {
return InvalidSchema("Field '{}' not found in nested mapping", field_name);
}

if (!nested_field->field_id.has_value()) {
return InvalidSchema("Field ID is missing for field '{}' in nested mapping",
field_name);
}

// Preserve existing custom attributes for this field
::avro::CustomAttributes attributes;
if (i < original_node->customAttributes()) {
// Copy all existing attributes from the original node
for (const auto& attr_pair : original_node->customAttributesAt(i).attributes()) {
// Copy each existing attribute to preserve original metadata
attributes.addAttribute(attr_pair.first, attr_pair.second, /*addQuote=*/false);
}
}

// Add field ID attribute to the new node (preserving existing attributes)
attributes.addAttribute(std::string(kFieldIdProp),
std::to_string(nested_field->field_id.value()),
/*addQuote=*/false);

new_record_node->addCustomAttributesForField(attributes);

// Recursively apply field IDs to nested fields
ICEBERG_ASSIGN_OR_RAISE(auto new_nested_node,
MakeAvroNodeWithFieldIds(field_node, *nested_field));
new_record_node->addName(field_name);
new_record_node->addLeaf(new_nested_node);
}

return new_record_node;
}

Result<::avro::NodePtr> CreateArrayNodeWithFieldIds(const ::avro::NodePtr& original_node,
const MappedField& field) {
if (original_node->leaves() != 1) {
return InvalidSchema("Array type must have exactly one leaf");
}

auto new_array_node = std::make_shared<::avro::NodeArray>();

// Check if this is a map represented as array
if (HasMapLogicalType(original_node)) {
ICEBERG_ASSIGN_OR_RAISE(auto new_element_node,
MakeAvroNodeWithFieldIds(original_node->leafAt(0), field));
new_array_node->addLeaf(new_element_node);
Comment thread
MisterRaindrop marked this conversation as resolved.
return new_array_node;
}

// For regular arrays, use the first field from nested mapping as element field
if (!field.nested_mapping || field.nested_mapping->fields().empty()) {
return InvalidSchema("Array type requires nested mapping with element field");
}

const auto& element_field = field.nested_mapping->fields()[0];

if (!element_field.field_id.has_value()) {
return InvalidSchema("Field ID is missing for element field in array");
}

ICEBERG_ASSIGN_OR_RAISE(
auto new_element_node,
MakeAvroNodeWithFieldIds(original_node->leafAt(0), element_field));
new_array_node->addLeaf(new_element_node);

// Add element field ID as custom attribute
::avro::CustomAttributes element_attributes;
element_attributes.addAttribute(std::string(kElementIdProp),
Comment thread
MisterRaindrop marked this conversation as resolved.
Outdated
std::to_string(*element_field.field_id),
/*addQuote=*/false);
new_array_node->addCustomAttributesForField(element_attributes);

return new_array_node;
}

Result<::avro::NodePtr> CreateMapNodeWithFieldIds(const ::avro::NodePtr& original_node,
const MappedField& field) {
if (original_node->leaves() != 2) {
return InvalidSchema("Map type must have exactly two leaves");
}

auto new_map_node = std::make_shared<::avro::NodeMap>();

// For map types, we need to extract key and value field mappings from the nested
// mapping
if (!field.nested_mapping) {
return InvalidSchema("Map type requires nested mapping for key and value fields");
}

// For map types, use the first two fields from nested mapping as key and value
if (!field.nested_mapping || field.nested_mapping->fields().size() < 2) {
return InvalidSchema("Map type requires nested mapping with key and value fields");
}

const auto& key_mapped_field = field.nested_mapping->fields()[0];
const auto& value_mapped_field = field.nested_mapping->fields()[1];

if (!key_mapped_field.field_id || !value_mapped_field.field_id) {
return InvalidSchema("Map key and value fields must have field IDs");
}

// Add key and value nodes
ICEBERG_ASSIGN_OR_RAISE(
auto new_key_node,
MakeAvroNodeWithFieldIds(original_node->leafAt(0), key_mapped_field));
ICEBERG_ASSIGN_OR_RAISE(
auto new_value_node,
MakeAvroNodeWithFieldIds(original_node->leafAt(1), value_mapped_field));
new_map_node->addLeaf(new_key_node);
new_map_node->addLeaf(new_value_node);

// Preserve existing custom attributes from the original node and add field ID
// attributes Copy existing attributes from the original node (if any)
if (original_node->customAttributes() > 0) {
const auto& original_attrs = original_node->customAttributesAt(0);
const auto& existing_attrs = original_attrs.attributes();
for (const auto& attr_pair : existing_attrs) {
// Copy each existing attribute to preserve original metadata
::avro::CustomAttributes attributes;
attributes.addAttribute(attr_pair.first, attr_pair.second, /*addQuote=*/false);
new_map_node->addCustomAttributesForField(attributes);
Comment thread
MisterRaindrop marked this conversation as resolved.
Outdated
}
}

::avro::CustomAttributes key_attributes;
key_attributes.addAttribute(std::string(kKeyIdProp),
std::to_string(*key_mapped_field.field_id),
/*addQuote=*/false);
new_map_node->addCustomAttributesForField(key_attributes);

::avro::CustomAttributes value_attributes;
value_attributes.addAttribute(std::string(kValueIdProp),
std::to_string(*value_mapped_field.field_id),
/*addQuote=*/false);
new_map_node->addCustomAttributesForField(value_attributes);

return new_map_node;
}

Result<::avro::NodePtr> CreateUnionNodeWithFieldIds(const ::avro::NodePtr& original_node,
const MappedField& field) {
if (original_node->leaves() != 2) {
return InvalidSchema("Union type must have exactly two branches");
}

const auto& branch_0 = original_node->leafAt(0);
const auto& branch_1 = original_node->leafAt(1);

bool branch_0_is_null = (branch_0->type() == ::avro::AVRO_NULL);
bool branch_1_is_null = (branch_1->type() == ::avro::AVRO_NULL);

if (branch_0_is_null && !branch_1_is_null) {
// branch_0 is null, branch_1 is not null
ICEBERG_ASSIGN_OR_RAISE(auto new_branch_1, MakeAvroNodeWithFieldIds(branch_1, field));
auto new_union_node = std::make_shared<::avro::NodeUnion>();
new_union_node->addLeaf(branch_0); // null branch
new_union_node->addLeaf(new_branch_1);
return new_union_node;
} else if (!branch_0_is_null && branch_1_is_null) {
// branch_0 is not null, branch_1 is null
ICEBERG_ASSIGN_OR_RAISE(auto new_branch_0, MakeAvroNodeWithFieldIds(branch_0, field));
auto new_union_node = std::make_shared<::avro::NodeUnion>();
new_union_node->addLeaf(new_branch_0);
new_union_node->addLeaf(branch_1); // null branch
return new_union_node;
} else if (branch_0_is_null && branch_1_is_null) {
// Both branches are null - this is invalid
return InvalidSchema("Union type cannot have two null branches");
} else {
// Neither branch is null - this is invalid
return InvalidSchema("Union type must have exactly one null branch");
}
}

} // namespace

Result<::avro::NodePtr> MakeAvroNodeWithFieldIds(const ::avro::NodePtr& original_node,
const MappedField& mapped_field) {
switch (original_node->type()) {
case ::avro::AVRO_RECORD:
return CreateRecordNodeWithFieldIds(original_node, mapped_field);
case ::avro::AVRO_ARRAY:
return CreateArrayNodeWithFieldIds(original_node, mapped_field);
case ::avro::AVRO_MAP:
return CreateMapNodeWithFieldIds(original_node, mapped_field);
case ::avro::AVRO_UNION:
return CreateUnionNodeWithFieldIds(original_node, mapped_field);
case ::avro::AVRO_BOOL:
case ::avro::AVRO_INT:
case ::avro::AVRO_LONG:
case ::avro::AVRO_FLOAT:
case ::avro::AVRO_DOUBLE:
case ::avro::AVRO_STRING:
case ::avro::AVRO_BYTES:
case ::avro::AVRO_FIXED:
// For primitive types, just return a copy
return original_node;
case ::avro::AVRO_NULL:
case ::avro::AVRO_ENUM:
default:
return InvalidSchema("Unsupported Avro type for field ID application: {}",
ToString(original_node));
}
}

Result<::avro::NodePtr> MakeAvroNodeWithFieldIds(const ::avro::NodePtr& original_node,
const NameMapping& mapping) {
MappedField mapped_field;
mapped_field.nested_mapping = std::make_shared<MappedFields>(mapping.AsMappedFields());
return MakeAvroNodeWithFieldIds(original_node, mapped_field);
}

} // namespace iceberg::avro
15 changes: 15 additions & 0 deletions src/iceberg/avro/avro_schema_util_internal.h
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@

#include <avro/Node.hh>

#include "iceberg/name_mapping.h"
#include "iceberg/result.h"
#include "iceberg/schema_util.h"
#include "iceberg/type.h"
Expand Down Expand Up @@ -148,4 +149,18 @@ std::string ToString(const ::avro::LogicalType::Type& logical_type);
/// \return True if the node has a map logical type, false otherwise.
bool HasMapLogicalType(const ::avro::NodePtr& node);

/// \brief Create a new Avro node with field IDs from name mapping.
/// \param original_node The original Avro node to copy.
/// \param mapped_field The mapped field to apply field IDs from.
/// \return A new Avro node with field IDs applied, or an error.
Result<::avro::NodePtr> MakeAvroNodeWithFieldIds(const ::avro::NodePtr& original_node,
const MappedField& mapped_field);

/// \brief Create a new Avro node with field IDs from name mapping.
/// \param original_node The original Avro node to copy.
/// \param mapping The name mapping to apply field IDs from.
/// \return A new Avro node with field IDs applied, or an error.
Result<::avro::NodePtr> MakeAvroNodeWithFieldIds(const ::avro::NodePtr& original_node,
const NameMapping& mapping);

} // namespace iceberg::avro
34 changes: 34 additions & 0 deletions src/iceberg/avro/constants.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,34 @@
/*
* 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.
*/

#pragma once

#include <string_view>

namespace iceberg::avro {

// Avro logical type constants
constexpr std::string_view kMapLogicalType = "map";

// Name mapping field constants
constexpr std::string_view kElement = "element";
constexpr std::string_view kKey = "key";
constexpr std::string_view kValue = "value";

} // namespace iceberg::avro
3 changes: 3 additions & 0 deletions src/iceberg/file_reader.h
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,9 @@ struct ICEBERG_EXPORT ReaderOptions {
/// \brief The filter to apply to the data. Reader implementations may ignore this if
/// the file format does not support filtering.
std::shared_ptr<class Expression> filter;
/// \brief Name mapping for schema evolution compatibility. Used when reading files
/// that may have different field names than the current schema.
std::shared_ptr<class NameMapping> name_mapping;
/// \brief Format-specific or implementation-specific properties.
std::unordered_map<std::string, std::string> properties;
};
Expand Down
7 changes: 4 additions & 3 deletions src/iceberg/name_mapping.cc
Original file line number Diff line number Diff line change
Expand Up @@ -152,22 +152,23 @@ const std::unordered_map<int32_t, MappedFieldConstRef>& MappedFields::LazyIdToFi
NameMapping::NameMapping(std::unique_ptr<MappedFields> mapping)
: mapping_(std::move(mapping)) {}

std::optional<MappedFieldConstRef> NameMapping::Find(int32_t id) {
std::optional<MappedFieldConstRef> NameMapping::Find(int32_t id) const {
const auto& fields_by_id = LazyFieldsById();
if (auto iter = fields_by_id.find(id); iter != fields_by_id.cend()) {
return iter->second;
}
return std::nullopt;
}

std::optional<MappedFieldConstRef> NameMapping::Find(std::span<const std::string> names) {
std::optional<MappedFieldConstRef> NameMapping::Find(
std::span<const std::string> names) const {
if (names.empty()) {
return std::nullopt;
}
return Find(JoinByDot(names));
}

std::optional<MappedFieldConstRef> NameMapping::Find(const std::string& name) {
std::optional<MappedFieldConstRef> NameMapping::Find(const std::string& name) const {
const auto& fields_by_name = LazyFieldsByName();
if (auto iter = fields_by_name.find(name); iter != fields_by_name.cend()) {
return iter->second;
Expand Down
Loading
Loading