diff --git a/pyiceberg/utils/schema_conversion.py b/pyiceberg/utils/schema_conversion.py index 7e6343141b..d1900f68f5 100644 --- a/pyiceberg/utils/schema_conversion.py +++ b/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" diff --git a/tests/avro/test_file.py b/tests/avro/test_file.py index 175089c30b..ab6feae661 100644 --- a/tests/avro/test_file.py +++ b/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, diff --git a/tests/utils/test_schema_conversion.py b/tests/utils/test_schema_conversion.py index b110f92bd2..a1d45a4d6f 100644 --- a/tests/utils/test_schema_conversion.py +++ b/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