feat: avro support applying field-ids based on name mapping (#127)

see issue #128

---------

Co-authored-by: Gang Wu <[email protected]>

Author: MisterRaindrop

src/iceberg/avro/avro_reader.cc

@@ -36,6 +36,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"
@@ -92,16 +93,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");
     }
 

src/iceberg/avro/avro_schema_util.cc

@@ -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"
@@ -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

src/iceberg/avro/avro_schema_util_internal.h

@@ -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"
@@ -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

src/iceberg/avro/constants.h

@@ -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