Fix issue #3758: manifest pruning for transformed partitions in upsert/delete

pyiceberg/table/update/snapshot.py

@@ -26,7 +26,7 @@
 from typing import TYPE_CHECKING, Generic
 
 from pyiceberg.avro.codecs import AvroCompressionCodec
-from pyiceberg.expressions import AlwaysFalse, BooleanExpression, Or
+from pyiceberg.expressions import AlwaysFalse, And, BooleanExpression, EqualTo, IsNull, Or, Reference
 from pyiceberg.expressions.visitors import (
     ROWS_MIGHT_NOT_MATCH,
     ROWS_MUST_MATCH,
@@ -91,6 +91,30 @@ def _new_manifest_list_file_name(snapshot_id: int, attempt: int, commit_uuid: uu
     return f"snap-{snapshot_id}-{attempt}-{commit_uuid}.avro"
 
 
+def _partition_records_filter(spec: PartitionSpec, partition_records: set[Record]) -> BooleanExpression:
+    """Build a filter over the partition fields matching any of the given partition records.
+
+    The returned expression references partition field names and transformed values, so it is
+    already in the domain a manifest evaluator binds against. It must not be projected through
+    `inclusive_projection`, which expects a predicate on the source columns instead.
+    """
+    partition_names = [field.name for field in spec.fields]
+    if not partition_records or not partition_names:
+        return AlwaysFalse()
+
+    per_record_exprs: list[BooleanExpression] = []
+    for partition_record in partition_records:
+        predicates: list[BooleanExpression] = [
+            EqualTo(Reference(partition_name), partition_record[pos])
+            if partition_record[pos] is not None
+            else IsNull(Reference(partition_name))
+            for pos, partition_name in enumerate(partition_names)
+        ]
+        per_record_exprs.append(And(*predicates) if len(predicates) > 1 else predicates[0])
+
+    return Or(*per_record_exprs) if len(per_record_exprs) > 1 else per_record_exprs[0]
+
+
 class _SnapshotProducer(UpdateTableMetadata[U], Generic[U]):
     commit_uuid: uuid.UUID
     _io: FileIO
@@ -212,9 +236,6 @@ def _write_delete_manifest() -> list[ManifestFile]:
             else:
                 return []
 
-        # Updates self._predicate with computed partition predicate for manifest pruning
-        self._build_delete_files_partition_predicate()
-
         executor = ExecutorFactory.get_or_create()
 
         added_manifests = executor.submit(_write_added_manifest)
@@ -373,20 +394,6 @@ def delete_by_predicate(self, predicate: BooleanExpression, case_sensitive: bool
         self._predicate = Or(self._predicate, predicate)
         self._case_sensitive = case_sensitive
 
-    def _build_delete_files_partition_predicate(self) -> None:
-        """Build BooleanExpression based on deleted data files partitions."""
-        partition_to_overwrite: dict[int, set[Record]] = {}
-        for data_file in self._deleted_data_files:
-            group = partition_to_overwrite.setdefault(data_file.spec_id, set())
-            group.add(data_file.partition)
-
-        for spec_id, partition_records in partition_to_overwrite.items():
-            self.delete_by_predicate(
-                self._transaction._build_partition_predicate(
-                    partition_records=partition_records, schema=self.schema(), spec=self.spec(spec_id)
-                )
-            )
-
 
 class _DeleteFiles(_SnapshotProducer["_DeleteFiles"]):
     """Will delete manifest entries from the current snapshot based on the predicate.
@@ -588,6 +595,32 @@ class _OverwriteFiles(_SnapshotProducer["_OverwriteFiles"]):
     Data and delete files were added and removed in a logical overwrite operation.
     """
 
+    @cached_property
+    def _deleted_files_partition_filters(self) -> dict[int, BooleanExpression]:
+        """Per-spec filters matching the partitions of the data files being replaced.
+
+        A data file records its partition values already transformed, so these reference the
+        partition fields. Deriving them from a source-column predicate instead would mean
+        projecting it onto the spec, which applies the transform a second time.
+        """
+        partition_to_overwrite: dict[int, set[Record]] = defaultdict(set)
+        for data_file in self._deleted_data_files:
+            partition_to_overwrite[data_file.spec_id].add(data_file.partition)
+
+        return {
+            spec_id: _partition_records_filter(self.spec(spec_id), partition_records)
+            for spec_id, partition_records in partition_to_overwrite.items()
+        }
+
+    def _build_manifest_evaluator(self, spec_id: int) -> Callable[[ManifestFile], bool]:
+        """Prune manifests that cannot hold any of the data files being replaced.
+
+        An overwrite never carries a row-level predicate, so the partitions of those files are
+        the only thing to match on.
+        """
+        partition_filter = self._deleted_files_partition_filters.get(spec_id, AlwaysFalse())
+        return manifest_evaluator(self.spec(spec_id), self.schema(), partition_filter, self._case_sensitive)
+
     def _existing_manifests(self) -> list[ManifestFile]:
         """Determine if there are any existing manifest files."""
         existing_files = []

tests/integration/test_deletes.py

@@ -25,12 +25,13 @@
 from pyiceberg.catalog.rest import RestCatalog
 from pyiceberg.exceptions import NoSuchTableError
 from pyiceberg.expressions import AlwaysTrue, EqualTo, LessThanOrEqual
+from pyiceberg.io.pyarrow import schema_to_pyarrow
 from pyiceberg.manifest import ManifestEntryStatus
 from pyiceberg.partitioning import PartitionField, PartitionSpec
 from pyiceberg.schema import Schema
 from pyiceberg.table import Table
 from pyiceberg.table.snapshots import Operation, Summary
-from pyiceberg.transforms import IdentityTransform
+from pyiceberg.transforms import DayTransform, IdentityTransform
 from pyiceberg.types import FloatType, IntegerType, LongType, NestedField, StringType, TimestampType
 
 
@@ -1024,3 +1025,40 @@ def test_manifest_entry_snapshot_id_after_partial_deletes(session_catalog: RestC
             f"DELETED entry snapshot_id should be {after_delete_snapshot.snapshot_id} "
             f"(the deleting snapshot), but was {entry.snapshot_id}"
         )
+
+
+@pytest.mark.integration
+def test_delete_partial_rewrite_of_transformed_partition(session_catalog: RestCatalog) -> None:
+    """Delete part of a data file in a partition whose values are transformed.
+
+    The file is rewritten rather than dropped, so the manifest referencing it has to be
+    rewritten too. Pruning used to compare the source column against the transformed
+    partition value, which never selected that manifest, and it was carried over whole
+    beside the rewritten file — see https://github.com/apache/iceberg-python/issues/3758.
+    """
+    identifier = "default.test_delete_partial_rewrite_of_transformed_partition"
+
+    try:
+        session_catalog.drop_table(identifier)
+    except NoSuchTableError:
+        pass
+
+    schema = Schema(
+        NestedField(1, "idx", IntegerType()),
+        NestedField(2, "event_ts", TimestampType()),
+    )
+    tbl = session_catalog.create_table(
+        identifier,
+        schema=schema,
+        partition_spec=PartitionSpec(PartitionField(source_id=2, field_id=1000, transform=DayTransform(), name="ts_day")),
+    )
+
+    event_ts = datetime(2026, 1, 6, 12)
+    rows = [{"idx": 1, "event_ts": event_ts}, {"idx": 2, "event_ts": event_ts}]
+    tbl.append(pa.Table.from_pylist(rows, schema=schema_to_pyarrow(schema)))
+    assert len(tbl.inspect.files()) == 1, "both rows must share a data file to exercise a partial rewrite"
+
+    tbl.delete(EqualTo("idx", 1))
+
+    assert [snapshot.summary.operation for snapshot in tbl.snapshots()] == [Operation.APPEND, Operation.OVERWRITE]
+    assert tbl.scan().to_arrow()["idx"].to_pylist() == [2]

tests/table/test_upsert.py

@@ -14,7 +14,9 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+from datetime import datetime
 from pathlib import PosixPath
+from typing import Any
 
 import pyarrow as pa
 import pytest
@@ -26,11 +28,22 @@
 from pyiceberg.expressions import AlwaysTrue, And, EqualTo, Reference
 from pyiceberg.expressions.literals import LongLiteral
 from pyiceberg.io.pyarrow import schema_to_pyarrow
+from pyiceberg.partitioning import PartitionField, PartitionSpec
 from pyiceberg.schema import Schema
 from pyiceberg.table import Table, UpsertResult
 from pyiceberg.table.snapshots import Operation
 from pyiceberg.table.upsert_util import create_match_filter
-from pyiceberg.types import IntegerType, NestedField, StringType, StructType
+from pyiceberg.transforms import (
+    BucketTransform,
+    DayTransform,
+    HourTransform,
+    IdentityTransform,
+    MonthTransform,
+    Transform,
+    TruncateTransform,
+    YearTransform,
+)
+from pyiceberg.types import IntegerType, NestedField, StringType, StructType, TimestampType
 from tests.catalog.test_base import InMemoryCatalog
 
 
@@ -888,3 +901,59 @@ def test_upsert_snapshot_properties(catalog: Catalog) -> None:
     for snapshot in snapshots[initial_snapshot_count:]:
         assert snapshot.summary is not None
         assert snapshot.summary.additional_properties.get("test_prop") == "test_value"
+
+
+@pytest.mark.parametrize(
+    "transform, source_id, kept_key, updated_key",
+    [
+        # Identity and truncate already behaved: re-applying either to a value it has
+        # produced is a no-op, so they are controls rather than regression cases.
+        (IdentityTransform(), 3, "a", "b"),
+        (TruncateTransform(1), 1, "aa", "ab"),
+        (YearTransform(), 3, "a", "b"),
+        (MonthTransform(), 3, "a", "b"),
+        (DayTransform(), 3, "a", "b"),
+        (HourTransform(), 3, "a", "b"),
+        # keys chosen to collide, so that one data file holds both rows
+        (BucketTransform(4), 1, "k0", "k1"),
+    ],
+)
+def test_upsert_partial_rewrite_of_partitioned_file(
+    catalog: Catalog, transform: Transform[Any, Any], source_id: int, kept_key: str, updated_key: str
+) -> None:
+    """Upsert a row out of a data file that also holds a row it must leave alone.
+
+    Both rows land in the same partition, so the file is rewritten rather than dropped. The
+    manifest holding it has to be rewritten too — see https://github.com/apache/iceberg-python/issues/3758,
+    where the manifest was instead carried over whole and the superseded rows stayed visible.
+    """
+    identifier = "default.test_upsert_partial_rewrite_of_partitioned_file"
+    _drop_table(catalog, identifier)
+
+    schema = Schema(
+        NestedField(1, "key", StringType(), required=False),
+        NestedField(2, "value", IntegerType(), required=False),
+        NestedField(3, "event_ts", TimestampType(), required=False),
+    )
+    tbl = catalog.create_table(
+        identifier,
+        schema=schema,
+        partition_spec=PartitionSpec(PartitionField(source_id=source_id, field_id=1000, transform=transform, name="part")),
+    )
+
+    arrow_schema = schema_to_pyarrow(schema)
+    event_ts = datetime(2026, 1, 6, 12)
+
+    def rows(*pairs: tuple[str, int]) -> pa_table:
+        return pa.Table.from_pylist(
+            [{"key": key, "value": value, "event_ts": event_ts} for key, value in pairs], schema=arrow_schema
+        )
+
+    tbl.append(rows((updated_key, 1), (kept_key, 1)))
+    assert len(tbl.inspect.files()) == 1, "both rows must share a data file to exercise a partial rewrite"
+
+    assert_upsert_result(tbl.upsert(rows((updated_key, 2)), join_cols=["key"]), expected_updated=1, expected_inserted=0)
+
+    result = tbl.scan().to_arrow()
+    actual = zip(result["key"].to_pylist(), result["value"].to_pylist(), strict=True)
+    assert sorted(actual) == sorted([(updated_key, 2), (kept_key, 1)])