Avro: read named type references and write unique fixed type names

pyiceberg/utils/schema_conversion.py

@@ -81,7 +81,15 @@
 AvroType = str | Any
 
 
+# Avro types that have a name and can be referred to by it later in the schema
+AVRO_NAMED_TYPES = frozenset({"record", "enum", "fixed"})
+
+
 class AvroSchemaConversion:
+    def __init__(self) -> None:
+        # Named types defined so far, so that a later reference by name resolves to them
+        self._named_types: dict[str, IcebergType] = {}
+
     def avro_to_iceberg(self, avro_schema: dict[str, Any]) -> Schema:
         """Convert an Apache Avro into an Apache Iceberg schema equivalent.
 
@@ -127,6 +135,7 @@ def avro_to_iceberg(self, avro_schema: dict[str, Any]) -> Schema:
         Returns:
             Equivalent Iceberg schema.
         """
+        self._named_types = {}
         return Schema(*[self._convert_field(field) for field in avro_schema["fields"]], schema_id=1)
 
     def iceberg_to_avro(self, schema: Schema, schema_name: str | None = None) -> AvroType:
@@ -197,26 +206,33 @@ def _convert_schema(self, avro_type: str | dict[str, Any]) -> IcebergType:
         """
         if isinstance(avro_type, str) and avro_type in PRIMITIVE_FIELD_TYPE_MAPPING:
             return PRIMITIVE_FIELD_TYPE_MAPPING[avro_type]
+        elif isinstance(avro_type, str) and avro_type in self._named_types:
+            # A reference to a named type defined earlier, e.g. "decimal_10_2"
+            return self._named_types[avro_type]
         elif isinstance(avro_type, dict):
             if "logicalType" in avro_type:
-                return self._convert_logical_type(avro_type)
+                iceberg_type = self._convert_logical_type(avro_type)
             else:
                 # Resolve potential nested types
                 while "type" in avro_type and isinstance(avro_type["type"], dict):
                     avro_type = avro_type["type"]
                 type_identifier = avro_type["type"]
                 if type_identifier == "record":
-                    return self._convert_record_type(avro_type)
+                    iceberg_type = self._convert_record_type(avro_type)
                 elif type_identifier == "array":
-                    return self._convert_array_type(avro_type)
+                    iceberg_type = self._convert_array_type(avro_type)
                 elif type_identifier == "map":
-                    return self._convert_map_type(avro_type)
+                    iceberg_type = self._convert_map_type(avro_type)
                 elif type_identifier == "fixed":
-                    return self._convert_fixed_type(avro_type)
+                    iceberg_type = self._convert_fixed_type(avro_type)
                 elif isinstance(type_identifier, str) and type_identifier in PRIMITIVE_FIELD_TYPE_MAPPING:
-                    return PRIMITIVE_FIELD_TYPE_MAPPING[type_identifier]
+                    iceberg_type = PRIMITIVE_FIELD_TYPE_MAPPING[type_identifier]
                 else:
                     raise TypeError(f"Type not recognized: {avro_type}")
+            if avro_type.get("type") in AVRO_NAMED_TYPES and isinstance(name := avro_type.get("name"), str):
+                # A repeated definition of the same name is accepted, as older writers produced it
+                self._named_types[name] = iceberg_type
+            return iceberg_type
         else:
             raise TypeError(f"Type not recognized: {avro_type}")
 
@@ -506,6 +522,16 @@ def __init__(self, schema_name: str | None) -> None:
         """
         self.schema_name = schema_name
 
+    @staticmethod
+    def _name_fixed(result: AvroType, field_id: int) -> None:
+        """Add the field-id to the name of a fixed type.
+
+        Avro allows only one definition per name, and a decimal, uuid or fixed type can be used by more
+        than one field. Field-ids are unique, so the names stay unique, the same way records are named.
+        """
+        if isinstance(result, dict) and result.get("type") == "fixed":
+            result["name"] = f"{result['name']}_{field_id}"
+
     def schema(self, schema: Schema, struct_result: AvroType) -> AvroType:
         if isinstance(struct_result, dict) and self.schema_name is not None:
             struct_result["name"] = self.schema_name
@@ -527,6 +553,7 @@ def field(self, field: NestedField, field_result: AvroType) -> AvroType:
         # Sets the schema name
         if isinstance(field_result, dict) and field_result.get("type") == "record":
             field_result["name"] = f"r{field.field_id}"
+        self._name_fixed(field_result, field.field_id)
 
         original_name = field.name
         sanitized_name = make_compatible_name(original_name)
@@ -554,9 +581,12 @@ def list(self, list_type: ListType, element_result: AvroType) -> AvroType:
         # Sets the schema name in case of a record
         if isinstance(element_result, dict) and element_result.get("type") == "record":
             element_result["name"] = f"r{self.last_list_field_id}"
+        self._name_fixed(element_result, self.last_list_field_id)
         return {"type": "array", "element-id": self.last_list_field_id, "items": element_result}
 
     def map(self, map_type: MapType, key_result: AvroType, value_result: AvroType) -> AvroType:
+        self._name_fixed(key_result, self.last_map_key_field_id)
+        self._name_fixed(value_result, self.last_map_value_field_id)
         if isinstance(key_result, StringType):
             # Avro Maps does not support other keys than a String,
             return {

tests/avro/test_file.py

@@ -455,6 +455,46 @@ def field_uuid(self) -> UUID:
         assert record[idx] == avro_entry_read_with_fastavro[idx], f"Invalid {field} read with fastavro"
 
 
+def test_repeated_named_types_round_trip() -> None:
+    """Named types used by more than one field get unique names in the file schema and read back with both readers"""
+    schema = Schema(
+        NestedField(field_id=1, name="price", field_type=DecimalType(10, 2), required=True),
+        NestedField(field_id=2, name="cost", field_type=DecimalType(10, 2), required=False),
+        NestedField(field_id=3, name="id", field_type=UUIDType(), required=True),
+        NestedField(field_id=4, name="parent_id", field_type=UUIDType(), required=False),
+        schema_id=1,
+    )
+    record = Record(
+        Decimal("123.45"),
+        Decimal("-6.78"),
+        UUID("12345678-1234-5678-1234-567812345678"),
+        UUID("87654321-4321-8765-4321-876543218765"),
+    )
+
+    with TemporaryDirectory() as tmpdir:
+        tmp_avro_file = tmpdir + "/repeated_named_types.avro"
+        with avro.AvroOutputFile[Record](PyArrowFileIO().new_output(tmp_avro_file), schema, "repeated_named_types") as out:
+            out.write_block([record])
+
+        with avro.AvroFile[Record](PyArrowFileIO().new_input(tmp_avro_file), schema, {-1: Record}) as avro_reader:
+            assert avro_reader.header.get_schema() == schema
+            avro_entry = next(iter(avro_reader))
+
+        with open(tmp_avro_file, "rb") as fo:
+            r = reader(fo=fo)
+            file_fields = r.writer_schema["fields"]
+            avro_entry_read_with_fastavro = list(next(iter(r)).values())
+
+    assert [f["name"] for f in (file_fields[0]["type"], file_fields[1]["type"][1])] == ["decimal_10_2_1", "decimal_10_2_2"]
+    assert [f["name"] for f in (file_fields[2]["type"], file_fields[3]["type"][1])] == ["uuid_fixed_3", "uuid_fixed_4"]
+    for idx in range(len(record)):
+        assert record[idx] == avro_entry[idx]
+    # fastavro returns a uuid stored as fixed as its raw bytes
+    price, cost, id_, parent_id = avro_entry_read_with_fastavro
+    assert (price, cost) == (record[0], record[1])
+    assert (UUID(bytes=id_), UUID(bytes=parent_id)) == (record[2], record[3])
+
+
 def manifest_entry(index: int) -> ManifestEntry:
     return ManifestEntry.from_args(
         status=ManifestEntryStatus.ADDED,

tests/utils/test_schema_conversion.py

@@ -389,3 +389,107 @@ def test_iceberg_to_avro_manifest(avro_schema_manifest_entry: dict[str, Any]) ->
     iceberg_schema = AvroSchemaConversion().avro_to_iceberg(avro_schema_manifest_entry)
     avro_result = AvroSchemaConversion().iceberg_to_avro(iceberg_schema, schema_name="manifest_entry")
     assert avro_schema_manifest_entry == avro_result
+
+
+DECIMAL_10_2 = {"type": "fixed", "name": "decimal_10_2", "size": 5, "logicalType": "decimal", "precision": 10, "scale": 2}
+
+
+def test_avro_to_iceberg_resolves_named_type_references() -> None:
+    """A named type used a second time is written as a reference to its name, as Java and iceberg-rust do"""
+    avro_schema = {
+        "type": "record",
+        "name": "r102",
+        "fields": [
+            {"name": "a", "type": ["null", DECIMAL_10_2], "field-id": 1000, "default": None},
+            {"name": "b", "type": ["null", "decimal_10_2"], "field-id": 1001, "default": None},
+            {"name": "c", "type": {"type": "fixed", "name": "fixed_3", "size": 3}, "field-id": 1002},
+            {"name": "d", "type": "fixed_3", "field-id": 1003},
+            {"name": "e", "type": {"type": "fixed", "name": "uuid_fixed", "size": 16, "logicalType": "uuid"}, "field-id": 1004},
+            {"name": "f", "type": ["null", "uuid_fixed"], "field-id": 1005, "default": None},
+        ],
+    }
+
+    assert AvroSchemaConversion().avro_to_iceberg(avro_schema) == Schema(
+        NestedField(field_id=1000, name="a", field_type=DecimalType(10, 2), required=False),
+        NestedField(field_id=1001, name="b", field_type=DecimalType(10, 2), required=False),
+        NestedField(field_id=1002, name="c", field_type=FixedType(3), required=True),
+        NestedField(field_id=1003, name="d", field_type=FixedType(3), required=True),
+        NestedField(field_id=1004, name="e", field_type=UUIDType(), required=True),
+        NestedField(field_id=1005, name="f", field_type=UUIDType(), required=False),
+        schema_id=1,
+    )
+
+
+def test_avro_to_iceberg_accepts_repeated_named_type_definitions() -> None:
+    """Manifests written before named types were defined once repeat the full definition"""
+    avro_schema = {
+        "type": "record",
+        "name": "r102",
+        "fields": [
+            {"name": "a", "type": ["null", DECIMAL_10_2], "field-id": 1000, "default": None},
+            {"name": "b", "type": ["null", DECIMAL_10_2], "field-id": 1001, "default": None},
+        ],
+    }
+
+    assert AvroSchemaConversion().avro_to_iceberg(avro_schema) == Schema(
+        NestedField(field_id=1000, name="a", field_type=DecimalType(10, 2), required=False),
+        NestedField(field_id=1001, name="b", field_type=DecimalType(10, 2), required=False),
+        schema_id=1,
+    )
+
+
+def test_avro_to_iceberg_unknown_named_type_reference() -> None:
+    avro_schema = {"type": "record", "name": "r102", "fields": [{"name": "a", "type": "decimal_10_2", "field-id": 1000}]}
+
+    with pytest.raises(TypeError, match="Type not recognized: decimal_10_2"):
+        AvroSchemaConversion().avro_to_iceberg(avro_schema)
+
+
+def test_iceberg_to_avro_names_fixed_types_by_field_id() -> None:
+    """A named type can be defined once per schema, so fixed types used by several fields get the field-id in their name"""
+    schema = Schema(
+        NestedField(field_id=1, name="a", field_type=DecimalType(10, 2), required=False),
+        NestedField(field_id=2, name="b", field_type=DecimalType(10, 2), required=True),
+        NestedField(field_id=3, name="c", field_type=FixedType(3), required=True),
+        NestedField(
+            field_id=4,
+            name="nested",
+            field_type=StructType(
+                NestedField(field_id=5, name="d", field_type=FixedType(3), required=False),
+                NestedField(field_id=6, name="e", field_type=UUIDType(), required=True),
+            ),
+            required=True,
+        ),
+        NestedField(field_id=7, name="f", field_type=UUIDType(), required=False),
+        NestedField(
+            field_id=8,
+            name="g",
+            field_type=ListType(element_id=9, element_type=UUIDType(), element_required=True),
+            required=True,
+        ),
+        NestedField(
+            field_id=10,
+            name="h",
+            field_type=MapType(key_id=11, key_type=UUIDType(), value_id=12, value_type=DecimalType(10, 2)),
+            required=True,
+        ),
+        schema_id=1,
+    )
+
+    avro_schema = AvroSchemaConversion().iceberg_to_avro(schema, schema_name="r102")
+    assert isinstance(avro_schema, dict)
+    fields = avro_schema["fields"]
+    nested_fields = fields[3]["type"]["fields"]
+    map_fields = fields[6]["type"]["items"]["fields"]
+
+    assert fields[0]["type"] == ["null", {**DECIMAL_10_2, "name": "decimal_10_2_1"}]
+    assert fields[1]["type"] == {**DECIMAL_10_2, "name": "decimal_10_2_2"}
+    assert fields[2]["type"] == {"type": "fixed", "size": 3, "name": "fixed_3_3"}
+    assert nested_fields[0]["type"] == ["null", {"type": "fixed", "size": 3, "name": "fixed_3_5"}]
+    assert nested_fields[1]["type"] == {"type": "fixed", "size": 16, "logicalType": "uuid", "name": "uuid_fixed_6"}
+    assert fields[4]["type"] == ["null", {"type": "fixed", "size": 16, "logicalType": "uuid", "name": "uuid_fixed_7"}]
+    assert fields[5]["type"]["items"]["name"] == "uuid_fixed_9"
+    assert map_fields[0]["type"]["name"] == "uuid_fixed_11"
+    assert map_fields[1]["type"]["name"] == "decimal_10_2_12"
+    # and the schema reads back as the same types
+    assert AvroSchemaConversion().avro_to_iceberg(avro_schema) == schema