Avro: define named types once and read references to them

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 = {"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}")
 
@@ -497,6 +513,7 @@ class ConvertSchemaToAvro(SchemaVisitorPerPrimitiveType[AvroType]):
     last_list_field_id: int
     last_map_key_field_id: int
     last_map_value_field_id: int
+    defined_names: set[str]
 
     def __init__(self, schema_name: str | None) -> None:
         """Convert an Iceberg schema to an Avro schema.
@@ -505,6 +522,18 @@ def __init__(self, schema_name: str | None) -> None:
             schema_name: The name of the root record.
         """
         self.schema_name = schema_name
+        self.defined_names: set[str] = set()
+
+    def _define_once(self, named_type: dict[str, Any]) -> AvroType:
+        """Return the definition of a named type the first time, and a reference to its name after that.
+
+        Avro allows only one definition per name; later uses must refer to it by name.
+        """
+        name = named_type["name"]
+        if name in self.defined_names:
+            return name
+        self.defined_names.add(name)
+        return named_type
 
     def schema(self, schema: Schema, struct_result: AvroType) -> AvroType:
         if isinstance(struct_result, dict) and self.schema_name is not None:
@@ -582,17 +611,19 @@ def map(self, map_type: MapType, key_result: AvroType, value_result: AvroType) -
             }
 
     def visit_fixed(self, fixed_type: FixedType) -> AvroType:
-        return {"type": "fixed", "size": len(fixed_type), "name": f"fixed_{len(fixed_type)}"}
+        return self._define_once({"type": "fixed", "size": len(fixed_type), "name": f"fixed_{len(fixed_type)}"})
 
     def visit_decimal(self, decimal_type: DecimalType) -> AvroType:
-        return {
-            "type": "fixed",
-            "size": decimal_required_bytes(decimal_type.precision),
-            "logicalType": "decimal",
-            "precision": decimal_type.precision,
-            "scale": decimal_type.scale,
-            "name": f"decimal_{decimal_type.precision}_{decimal_type.scale}",
-        }
+        return self._define_once(
+            {
+                "type": "fixed",
+                "size": decimal_required_bytes(decimal_type.precision),
+                "logicalType": "decimal",
+                "precision": decimal_type.precision,
+                "scale": decimal_type.scale,
+                "name": f"decimal_{decimal_type.precision}_{decimal_type.scale}",
+            }
+        )
 
     def visit_boolean(self, boolean_type: BooleanType) -> AvroType:
         return "boolean"
@@ -631,7 +662,7 @@ def visit_string(self, string_type: StringType) -> AvroType:
         return "string"
 
     def visit_uuid(self, uuid_type: UUIDType) -> AvroType:
-        return {"type": "fixed", "size": 16, "logicalType": "uuid", "name": "uuid_fixed"}
+        return self._define_once({"type": "fixed", "size": 16, "logicalType": "uuid", "name": "uuid_fixed"})
 
     def visit_binary(self, binary_type: BinaryType) -> AvroType:
         return "bytes"

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 more than once are defined once in the file schema and read back by 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 file_fields[1]["type"] == ["null", "decimal_10_2"]
+    assert file_fields[3]["type"] == ["null", "uuid_fixed"]
+    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,93 @@ 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_defines_named_types_once() -> None:
+    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=DecimalType(9, 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"]
+
+    assert fields[0]["type"] == ["null", DECIMAL_10_2]
+    assert fields[1]["type"] == "decimal_10_2"
+    assert fields[2]["type"] == {"type": "fixed", "size": 3, "name": "fixed_3"}
+    assert nested_fields[0]["type"] == ["null", "fixed_3"]
+    assert nested_fields[1]["type"] == {"type": "fixed", "size": 16, "logicalType": "uuid", "name": "uuid_fixed"}
+    assert fields[4]["type"] == ["null", "uuid_fixed"]
+    # a different decimal is a different name, so it gets its own definition
+    assert fields[5]["type"]["name"] == "decimal_9_2"
+    # and the references read back as the same types
+    assert AvroSchemaConversion().avro_to_iceberg(avro_schema) == schema