diff --git a/pyiceberg/avro/resolver.py b/pyiceberg/avro/resolver.py index 3f19ae9878..430e3bb171 100644 --- a/pyiceberg/avro/resolver.py +++ b/pyiceberg/avro/resolver.py @@ -147,7 +147,7 @@ def struct(self, struct: StructType, field_results: list[Writer]) -> Writer: return StructWriter(tuple((pos, result) for pos, result in enumerate(field_results))) def field(self, field: NestedField, field_result: Writer) -> Writer: - return field_result if field.required else OptionWriter(field_result) + return field_result if field.required or isinstance(field.field_type, UnknownType) else OptionWriter(field_result) def list(self, list_type: ListType, element_result: Writer) -> Writer: return ListWriter(element_result) @@ -302,7 +302,9 @@ def struct(self, file_schema: StructType, record_struct: IcebergType | None, fil return StructWriter(field_writers=tuple(results)) def field(self, file_field: NestedField, record_type: IcebergType | None, field_writer: Writer) -> Writer: - return field_writer if file_field.required else OptionWriter(field_writer) + return ( + field_writer if file_field.required or isinstance(file_field.field_type, UnknownType) else OptionWriter(field_writer) + ) def list(self, file_list_type: ListType, file_list: IcebergType | None, element_writer: Writer) -> Writer: return ListWriter(element_writer if file_list_type.element_required else OptionWriter(element_writer)) diff --git a/pyiceberg/utils/schema_conversion.py b/pyiceberg/utils/schema_conversion.py index 7e6343141b..f6448aa565 100644 --- a/pyiceberg/utils/schema_conversion.py +++ b/pyiceberg/utils/schema_conversion.py @@ -534,7 +534,8 @@ def field(self, field: NestedField, field_result: AvroType) -> AvroType: result = { "name": sanitized_name, FIELD_ID_PROP: field.field_id, - "type": field_result if field.required else ["null", field_result], + # Unknown maps to Avro "null", and a union cannot contain "null" twice + "type": field_result if field.required or isinstance(field.field_type, UnknownType) else ["null", field_result], } if original_name != sanitized_name: diff --git a/tests/avro/test_file.py b/tests/avro/test_file.py index 175089c30b..ad2185d172 100644 --- a/tests/avro/test_file.py +++ b/tests/avro/test_file.py @@ -54,6 +54,7 @@ TimestampType, TimestamptzType, TimeType, + UnknownType, UUIDType, ) from pyiceberg.utils.schema_conversion import AvroSchemaConversion @@ -544,3 +545,23 @@ def test_avro_output_file_tell_includes_buffered_records() -> None: out.write_block([manifest_entry(0)]) assert out.tell() > after_header + + +def test_write_optional_unknown_type() -> None: + # A partition field whose source column was dropped has the unknown type + schema = Schema( + NestedField(field_id=1, name="before", field_type=StringType(), required=False), + NestedField(field_id=2, name="dropped", field_type=UnknownType(), required=False), + NestedField(field_id=3, name="after", field_type=LongType(), required=True), + ) + + with TemporaryDirectory() as tmpdir: + tmp_avro_file = tmpdir + "/unknown.avro" + with avro.AvroOutputFile[Record](PyArrowFileIO().new_output(tmp_avro_file), schema, "unknown") as out: + out.write_block([Record("hello", None, 42)]) + + with avro.AvroFile[Record](PyArrowFileIO().new_input(tmp_avro_file), schema, {-1: Record}) as avro_reader: + assert list(avro_reader) == [Record("hello", None, 42)] + + with open(tmp_avro_file, "rb") as fo: + assert list(reader(fo=fo)) == [{"before": "hello", "dropped": None, "after": 42}] diff --git a/tests/catalog/test_catalog_behaviors.py b/tests/catalog/test_catalog_behaviors.py index 4310ed1eb6..36efcb8cfb 100644 --- a/tests/catalog/test_catalog_behaviors.py +++ b/tests/catalog/test_catalog_behaviors.py @@ -848,6 +848,28 @@ def test_write_and_evolve(catalog: Catalog, format_version: int) -> None: txn.delete("foo = 'a'") +@pytest.mark.parametrize("format_version", [1, 2]) +def test_delete_after_dropping_old_partition_source_column(catalog: Catalog, format_version: int) -> None: + identifier = f"default.delete_after_dropping_old_partition_source_v{format_version}" + catalog.create_namespace_if_not_exists("default") + + pa_table = pa.Table.from_pydict({"id": [1, 2], "category": ["a", "b"]}) + tbl = catalog.create_table(identifier=identifier, schema=pa_table.schema, properties={"format-version": str(format_version)}) + with tbl.update_spec() as spec_update: + spec_update.add_identity("category") + tbl.append(pa_table) + + with tbl.update_spec() as spec_update: + spec_update.remove_field("category") + with tbl.update_schema() as schema_update: + schema_update.delete_column("category") + + # Rewrites the manifest of the old spec, whose partition source column is gone + tbl.delete("id == 1") + + assert catalog.load_table(identifier).scan().to_arrow()["id"].to_pylist() == [2] + + # Merge manifests @pytest.mark.parametrize("format_version", [1, 2]) def test_merge_manifests_local_file_system(catalog: Catalog, arrow_table_with_null: pa.Table, format_version: int) -> None: