diff --git a/pyiceberg/table/update/snapshot.py b/pyiceberg/table/update/snapshot.py index 7931edacdd..f98e54f148 100644 --- a/pyiceberg/table/update/snapshot.py +++ b/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 = [] diff --git a/tests/integration/test_deletes.py b/tests/integration/test_deletes.py index c49689b716..733c7fb36a 100644 --- a/tests/integration/test_deletes.py +++ b/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] diff --git a/tests/table/test_upsert.py b/tests/table/test_upsert.py index 08f90c6600..2193b90964 100644 --- a/tests/table/test_upsert.py +++ b/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)])