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
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
278 changes: 278 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,280 @@ Result<SchemaProjection> Project(const Schema& expected_schema,
return SchemaProjection{std::move(field_projection.children)};
}

namespace {

void CopyCustomAttributes(const ::avro::CustomAttributes& source,
::avro::CustomAttributes& target) {
for (const auto& attr_pair : source.attributes()) {
target.addAttribute(attr_pair.first, attr_pair.second, /*addQuote=*/false);
}
}

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);
::avro::NodePtr field_node = original_node->leafAt(i);

// 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);

// Check and add custom attributes
if (original_node->customAttributes() > 0) {
::avro::CustomAttributes merged_attributes;
const auto& original_attrs = original_node->customAttributesAt(0);
CopyCustomAttributes(original_attrs, merged_attributes);
// Add merged attributes if we found any
if (merged_attributes.attributes().size() > 0) {
new_array_node->addCustomAttributesForField(merged_attributes);
}
}

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);

// Create merged custom attributes with element field ID
::avro::CustomAttributes merged_attributes;

// First add our element field ID (highest priority)
merged_attributes.addAttribute(std::string(kElementIdProp),
std::to_string(*element_field.field_id),
/*addQuote=*/false);

// Then merge any custom attributes from original node
if (original_node->customAttributes() > 0) {
const auto& original_attrs = original_node->customAttributesAt(0);
CopyCustomAttributes(original_attrs, merged_attributes);
}

// Add all attributes at once
new_array_node->addCustomAttributesForField(merged_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);

// Create key and value attributes
::avro::CustomAttributes key_attributes;
::avro::CustomAttributes value_attributes;

// Add required field IDs
key_attributes.addAttribute(std::string(kKeyIdProp),
std::to_string(*key_mapped_field.field_id),
/*addQuote=*/false);
value_attributes.addAttribute(std::string(kValueIdProp),
std::to_string(*value_mapped_field.field_id),
/*addQuote=*/false);

// Merge custom attributes from original node if they exist
if (original_node->customAttributes() > 0) {
// Merge attributes for key (index 0)
const auto& original_key_attrs = original_node->customAttributesAt(0);
CopyCustomAttributes(original_key_attrs, key_attributes);

// Merge attributes for value (index 1)
if (original_node->customAttributes() > 1) {
const auto& original_value_attrs = original_node->customAttributesAt(1);
CopyCustomAttributes(original_value_attrs, value_attributes);
}
}

// Add the merged attributes to the new map node
new_map_node->addCustomAttributesForField(key_attributes);
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
Loading
Loading