diff --git a/pyiceberg/utils/schema_conversion.py b/pyiceberg/utils/schema_conversion.py index 7e6343141b..4040523052 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 = 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 { diff --git a/tests/avro/test_file.py b/tests/avro/test_file.py index 175089c30b..3a5f994ed1 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 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, diff --git a/tests/utils/test_schema_conversion.py b/tests/utils/test_schema_conversion.py index b110f92bd2..fd7b6f3647 100644 --- a/tests/utils/test_schema_conversion.py +++ b/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