fix(manifest): keep the table readable after dropping an old partition column

pyiceberg/avro/file.py

@@ -145,7 +145,7 @@ class AvroFile(Generic[D]):
         "block",
     )
     input_file: InputFile
-    read_schema: Schema | None
+    read_schema: Schema | Callable[[AvroFileHeader, Schema], Schema] | None
     read_types: dict[int, Callable[..., StructProtocol]]
     read_enums: dict[int, Callable[..., Enum]]
     header: AvroFileHeader
@@ -158,7 +158,7 @@ class AvroFile(Generic[D]):
     def __init__(
         self,
         input_file: InputFile,
-        read_schema: Schema | None = None,
+        read_schema: Schema | Callable[[AvroFileHeader, Schema], Schema] | None = None,
         read_types: dict[int, Callable[..., StructProtocol]] = EMPTY_DICT,
         read_enums: dict[int, Callable[..., Enum]] = EMPTY_DICT,
     ) -> None:
@@ -178,8 +178,9 @@ def __enter__(self) -> AvroFile[D]:
             self.decoder = new_decoder(f.read())
         self.header = self._read_header()
         self.schema = self.header.get_schema()
-        if not self.read_schema:
-            self.read_schema = self.schema
+        # The read schema can depend on the file header, e.g. on the partition spec of a manifest
+        read_schema = self.read_schema(self.header, self.schema) if callable(self.read_schema) else self.read_schema
+        self.read_schema = read_schema or self.schema
 
         self.reader = resolve_reader(self.schema, self.read_schema, self.read_types, self.read_enums)
 

pyiceberg/manifest.py

@@ -16,6 +16,7 @@
 # under the License.
 from __future__ import annotations
 
+import json
 import math
 import threading
 from abc import ABC, abstractmethod
@@ -32,7 +33,7 @@
 from pydantic_core import to_json
 
 from pyiceberg.avro.codecs import AVRO_CODEC_KEY, AvroCompressionCodec
-from pyiceberg.avro.file import AvroFile, AvroOutputFile
+from pyiceberg.avro.file import AvroFile, AvroFileHeader, AvroOutputFile
 from pyiceberg.conversions import to_bytes
 from pyiceberg.encryption.io import EncryptedInputFile
 from pyiceberg.encryption.key_metadata import StandardKeyMetadata
@@ -52,6 +53,7 @@
     PrimitiveType,
     StringType,
     StructType,
+    UnknownType,
 )
 from pyiceberg.utils.config import Config
 
@@ -913,7 +915,7 @@ def fetch_manifest_entry(
         input_file = io.new_input(self.manifest_path)
         with AvroFile[ManifestEntry](
             input_file,
-            MANIFEST_ENTRY_SCHEMAS[DEFAULT_READ_VERSION],
+            _manifest_entry_read_schema,
             read_types={-1: ManifestEntry, 2: DataFile},
             read_enums={0: ManifestEntryStatus, 101: FileFormat, 134: DataFileContent},
         ) as reader:
@@ -1052,6 +1054,44 @@ def read_manifest_list(input_file: InputFile) -> Iterator[ManifestFile]:
         yield from reader
 
 
+def _manifest_entry_read_schema(header: AvroFileHeader, file_schema: Schema) -> Schema:
+    """Read the partition tuple with all fields of the partition spec in the manifest header.
+
+    Fields of the unknown type are omitted when writing a manifest. Missing fields are added back as
+    optional fields that read as null, which keeps the partition fields after them in place.
+    """
+    read_schema = MANIFEST_ENTRY_SCHEMAS[DEFAULT_READ_VERSION]
+    spec_json = header.meta.get("partition-spec")
+    if not spec_json:
+        return read_schema
+
+    data_file = next((field.field_type for field in file_schema.fields if field.field_id == 2), None)
+    if not isinstance(data_file, StructType):
+        return read_schema
+    file_partition = next((field.field_type for field in data_file.fields if field.field_id == 102), None)
+    if not isinstance(file_partition, StructType):
+        return read_schema
+
+    file_fields = {field.field_id: field for field in file_partition.fields}
+    try:
+        spec_fields = [(field["field-id"], field["name"]) for field in json.loads(spec_json)]
+    except (ValueError, TypeError, KeyError):
+        # Fall back to the fields in the file, e.g. for old specs without field IDs or written as an object
+        return read_schema
+    if all(field_id in file_fields for field_id, _ in spec_fields):
+        return read_schema
+
+    partition_type = StructType(
+        *[
+            file_fields.get(field_id) or NestedField(field_id=field_id, name=name, field_type=UnknownType(), required=False)
+            for field_id, name in spec_fields
+        ]
+    )
+    return manifest_entry_schema_with_data_file(
+        DEFAULT_READ_VERSION, data_file_with_partition(partition_type, DEFAULT_READ_VERSION)
+    )
+
+
 def _inherit_from_manifest(entry: ManifestEntry, manifest: ManifestFile) -> ManifestEntry:
     """
     Inherits properties from manifest file.
@@ -1172,16 +1212,20 @@ def _meta(self) -> dict[str, str]:
             AVRO_CODEC_KEY: self._compression,
         }
 
-    def _with_partition(self, format_version: TableVersion) -> Schema:
-        data_file_type = data_file_with_partition(
-            format_version=format_version, partition_type=self._spec.partition_type(self._schema)
-        )
+    def _with_partition(self, format_version: TableVersion, skip_unknown: bool = False) -> Schema:
+        partition_type = self._spec.partition_type(self._schema)
+        if skip_unknown:
+            # The record keeps all fields; the Avro file omits unknown ones, which read back as null
+            partition_type = StructType(
+                *[field for field in partition_type.fields if not isinstance(field.field_type, UnknownType)]
+            )
+        data_file_type = data_file_with_partition(format_version=format_version, partition_type=partition_type)
         return manifest_entry_schema_with_data_file(format_version=format_version, data_file=data_file_type)
 
     def new_writer(self) -> AvroOutputFile[ManifestEntry]:
         return AvroOutputFile[ManifestEntry](
             output_file=self._output_file,
-            file_schema=self._with_partition(self.version),
+            file_schema=self._with_partition(self.version, skip_unknown=True),
             record_schema=self._with_partition(DEFAULT_READ_VERSION),
             schema_name="manifest_entry",
             metadata=self._meta,

tests/catalog/test_catalog_behaviors.py

@@ -48,7 +48,7 @@
 from pyiceberg.table.sorting import NullOrder, SortDirection, SortField, SortOrder
 from pyiceberg.table.update import AddSchemaUpdate, SetCurrentSchemaUpdate
 from pyiceberg.transforms import IdentityTransform
-from pyiceberg.typedef import Identifier
+from pyiceberg.typedef import Identifier, Record
 from pyiceberg.types import BooleanType, DoubleType, IntegerType, LongType, NestedField, StringType
 
 
@@ -863,6 +863,63 @@ 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")
+
+    tbl = catalog.load_table(identifier)
+    assert tbl.scan().to_arrow()["id"].to_pylist() == [2]
+    # The skipped field reads back as null instead of disappearing from the partition tuple
+    snapshot = tbl.current_snapshot()
+    assert snapshot is not None
+    partitions = [
+        entry.data_file.partition for manifest in snapshot.manifests(tbl.io) for entry in manifest.fetch_manifest_entry(tbl.io)
+    ]
+    assert partitions == [Record(None)]
+
+
+@pytest.mark.parametrize("format_version", [1, 2])
+def test_filter_after_dropping_earlier_partition_source_column(catalog: Catalog, format_version: int) -> None:
+    identifier = f"default.filter_after_dropping_earlier_partition_source_v{format_version}"
+    catalog.create_namespace_if_not_exists("default")
+
+    pa_table = pa.Table.from_pydict({"id": [1, 2, 3], "category": ["a", "b", "a"], "region": ["x", "y", "y"]})
+    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")
+        spec_update.add_identity("region")
+    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")
+
+    # The dropped field comes before region in the old spec, so region must not shift when the manifest is read
+    tbl.delete("id == 1")
+    assert sorted(catalog.load_table(identifier).scan(row_filter="region == 'y'").to_arrow()["id"].to_pylist()) == [2, 3]
+
+    # Reads the rewritten manifest and writes it again
+    tbl = catalog.load_table(identifier)
+    tbl.delete("id == 2")
+    assert catalog.load_table(identifier).scan(row_filter="region == 'y'").to_arrow()["id"].to_pylist() == [3]
+
+
 # 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:

tests/utils/test_manifest.py

@@ -46,15 +46,18 @@
     _inherit_from_manifest,
     _manifests,
     clear_manifest_cache,
+    data_file_with_partition,
+    manifest_entry_schema_with_data_file,
     read_manifest_list,
     write_manifest,
     write_manifest_list,
 )
-from pyiceberg.partitioning import UNPARTITIONED_PARTITION_SPEC, PartitionSpec
+from pyiceberg.partitioning import UNPARTITIONED_PARTITION_SPEC, PartitionField, PartitionSpec
 from pyiceberg.schema import Schema
 from pyiceberg.table.snapshots import Operation, Snapshot, Summary
+from pyiceberg.transforms import IdentityTransform
 from pyiceberg.typedef import Record, TableVersion
-from pyiceberg.types import IntegerType, ListType, LongType, NestedField, StructType
+from pyiceberg.types import IntegerType, ListType, LongType, NestedField, StringType, StructType, UnknownType
 from tests.encryption.memory_kms import MemoryKeyManagementClient
 
 
@@ -1519,3 +1522,84 @@ def test_negative_manifest_cache_size_raises_value_error(monkeypatch: pytest.Mon
     finally:
         monkeypatch.delenv("PYICEBERG_MANIFEST_CACHE_SIZE", raising=False)
         importlib.reload(manifest_module)
+
+
+def _entry_with_partition(partition: Record) -> ManifestEntry:
+    return ManifestEntry.from_args(
+        status=ManifestEntryStatus.ADDED,
+        snapshot_id=1,
+        data_file=DataFile.from_args(
+            content=DataFileContent.DATA,
+            file_path="s3://bucket/data/file.parquet",
+            file_format=FileFormat.PARQUET,
+            partition=partition,
+            record_count=100,
+            file_size_in_bytes=1024,
+        ),
+    )
+
+
+# The source column of "category" (field 1) has been dropped, so its partition type is unknown
+DROPPED_SOURCE_SCHEMA = Schema(NestedField(2, "region", StringType(), False))
+DROPPED_SOURCE_SPEC = PartitionSpec(
+    PartitionField(source_id=1, field_id=1000, transform=IdentityTransform(), name="category"),
+    PartitionField(source_id=2, field_id=1001, transform=IdentityTransform(), name="region"),
+)
+
+
+def test_write_manifest_skips_unknown_partition_field(tmp_path: Path) -> None:
+    io = load_file_io()
+    for format_version in (1, 2):
+        path = str(tmp_path / f"manifest-v{format_version}.avro")
+        with write_manifest(
+            format_version=format_version,
+            spec=DROPPED_SOURCE_SPEC,
+            schema=DROPPED_SOURCE_SCHEMA,
+            output_file=io.new_output(path),
+            snapshot_id=1,
+            avro_compression="null",
+        ) as writer:
+            writer.add_entry(_entry_with_partition(Record(None, "y")))
+            manifest = writer.to_manifest_file()
+
+        with open(path, "rb") as f:
+            data_file_schema = next(field for field in fastavro.reader(f).writer_schema["fields"] if field["name"] == "data_file")
+        partition_schema = next(field for field in data_file_schema["type"]["fields"] if field["name"] == "partition")
+        assert [field["name"] for field in partition_schema["type"]["fields"]] == ["region"]
+
+        assert manifest.fetch_manifest_entry(io)[0].data_file.partition == Record(None, "y")
+
+
+def test_read_manifest_with_unknown_partition_field_written_as_null(tmp_path: Path) -> None:
+    # Java writes an unknown partition field as an Avro "null" field instead of omitting it
+    # required=True makes the Avro schema a plain "null" instead of a ["null", ...] union
+    io = load_file_io()
+    path = str(tmp_path / "manifest.avro")
+    partition_type = StructType(
+        NestedField(1000, "category", UnknownType(), required=True),
+        NestedField(1001, "region", StringType(), required=False),
+    )
+    with AvroOutputFile[ManifestEntry](
+        output_file=io.new_output(path),
+        file_schema=manifest_entry_schema_with_data_file(2, data_file_with_partition(partition_type, 2)),
+        schema_name="manifest_entry",
+        metadata={
+            "schema": DROPPED_SOURCE_SCHEMA.model_dump_json(),
+            "partition-spec": '[{"source-id":1,"field-id":1000,"transform":"identity","name":"category"},'
+            '{"source-id":2,"field-id":1001,"transform":"identity","name":"region"}]',
+            "partition-spec-id": "0",
+            "format-version": "2",
+        },
+    ) as out:
+        out.write_block([_entry_with_partition(Record(None, "y"))])
+
+    manifest = ManifestFile.from_args(
+        manifest_path=path,
+        manifest_length=1,
+        partition_spec_id=0,
+        content=ManifestContent.DATA,
+        sequence_number=1,
+        min_sequence_number=1,
+        added_snapshot_id=1,
+    )
+    assert manifest.fetch_manifest_entry(io)[0].data_file.partition == Record(None, "y")