diff --git a/pyiceberg/table/__init__.py b/pyiceberg/table/__init__.py index c0adce84dc..0faacae27d 100644 --- a/pyiceberg/table/__init__.py +++ b/pyiceberg/table/__init__.py @@ -362,7 +362,7 @@ def _set_ref_snapshot( return updates, requirements def _build_partition_predicate( - self, partition_records: set[Record], spec: PartitionSpec, schema: Schema + self, partition_records: set[Record], spec: PartitionSpec, schema: Schema, is_projected: bool = False ) -> BooleanExpression: """Build a filter predicate matching any of the input partition records. @@ -370,10 +370,13 @@ def _build_partition_predicate( partition_records: A set of partition records to match spec: An optional partition spec, if none then defaults to current schema: An optional schema, if none then defaults to current + is_projected: Reference the partition fields rather than their source columns. + Returns: A predicate matching any of the input partition records. """ - partition_fields = [schema.find_field(field.source_id).name for field in spec.fields] + partition_fields = [(field.name if is_projected else schema.find_field(field.source_id).name) for field in spec.fields] + if not partition_records or not partition_fields: return AlwaysFalse() diff --git a/pyiceberg/table/update/snapshot.py b/pyiceberg/table/update/snapshot.py index 7931edacdd..9285c98bd6 100644 --- a/pyiceberg/table/update/snapshot.py +++ b/pyiceberg/table/update/snapshot.py @@ -381,10 +381,14 @@ def _build_delete_files_partition_predicate(self) -> None: group.add(data_file.partition) for spec_id, partition_records in partition_to_overwrite.items(): - self.delete_by_predicate( + self.partition_filters[spec_id] = Or( + self.partition_filters[spec_id], self._transaction._build_partition_predicate( - partition_records=partition_records, schema=self.schema(), spec=self.spec(spec_id) - ) + partition_records=partition_records, + schema=self.schema(), + spec=self.spec(spec_id), + is_projected=True, + ), ) diff --git a/tests/table/test_upsert.py b/tests/table/test_upsert.py index 08f90c6600..78ddbc7c5c 100644 --- a/tests/table/test_upsert.py +++ b/tests/table/test_upsert.py @@ -14,6 +14,7 @@ # 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 import pyarrow as pa @@ -26,11 +27,13 @@ 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 DayTransform +from pyiceberg.types import IntegerType, NestedField, StringType, StructType, TimestampType from tests.catalog.test_base import InMemoryCatalog @@ -714,6 +717,42 @@ def test_upsert_with_nulls(catalog: Catalog) -> None: ) +def test_upsert_on_table_partitioned_by_transform(catalog: Catalog) -> None: + """Upsert has to rewrite the matched file on a table partitioned by a non-identity transform. + + The manifest pruning in the overwrite builds its predicate from the partition records of + the deleted files. Those records hold already-transformed values, so referencing the source + column would send them through the transform twice, prune away the only relevant manifest + and leave the replaced row behind as a duplicate. + """ + identifier = "default.test_upsert_on_table_partitioned_by_transform" + _drop_table(catalog, identifier) + + schema = Schema( + NestedField(1, "k", StringType(), required=False), + NestedField(2, "v", IntegerType(), required=False), + NestedField(3, "ts", TimestampType(), required=False), + ) + spec = PartitionSpec(PartitionField(source_id=3, field_id=1000, transform=DayTransform(), name="ts_day")) + table = catalog.create_table(identifier, schema, partition_spec=spec) + + arrow_schema = schema_to_pyarrow(schema) + # A timestamp whose day ordinal is far from the value it would be read as if the + # DayTransform were applied a second time. + ts = datetime(2026, 1, 6, 12) + + def rows(pairs: list[tuple[str, int]]) -> pa_table: + return pa.Table.from_pylist([{"k": k, "v": v, "ts": ts} for k, v in pairs], schema=arrow_schema) + + table.append(rows([("a", 1), ("b", 1)])) + + res = table.upsert(rows([("a", 2)]), join_cols=["k"]) + assert_upsert_result(res, expected_updated=1, expected_inserted=0) + + arrow = table.scan().to_arrow() + assert sorted(zip(arrow["k"].to_pylist(), arrow["v"].to_pylist(), strict=True)) == [("a", 2), ("b", 1)] + + def test_transaction(catalog: Catalog) -> None: """Test the upsert within a Transaction. Make sure that if something fails the entire Transaction is rolled back."""