From ff4a96622b3d50a0eb4671133ab4da3884d2c070 Mon Sep 17 00:00:00 2001 From: Nico Ritschel Date: Mon, 17 Aug 2026 20:16:49 -0700 Subject: [PATCH 1/6] Plan equality deletes in table scans --- pyiceberg/table/__init__.py | 6 +- pyiceberg/table/delete_file_index.py | 131 ++++++++++++++++-- pyiceberg/table/update/validate.py | 4 +- tests/table/test_delete_file_index.py | 192 +++++++++++++++++++++++++- 4 files changed, 311 insertions(+), 22 deletions(-) diff --git a/pyiceberg/table/__init__.py b/pyiceberg/table/__init__.py index 3dffc2270c..b545be702e 100644 --- a/pyiceberg/table/__init__.py +++ b/pyiceberg/table/__init__.py @@ -2804,7 +2804,7 @@ def plan_files( List of FileScanTasks that contain both data and delete files. """ data_entries: list[ManifestEntry] = [] - delete_index = DeleteFileIndex() + delete_index = DeleteFileIndex(self.table_metadata.schema()) residual_evaluators: dict[int, Callable[[DataFile], ResidualEvaluator]] = KeyDefaultDict(self._build_residual_evaluator) @@ -2815,10 +2815,8 @@ def plan_files( data_file = manifest_entry.data_file if data_file.content == DataFileContent.DATA: data_entries.append(manifest_entry) - elif data_file.content == DataFileContent.POSITION_DELETES: + elif data_file.content in (DataFileContent.POSITION_DELETES, DataFileContent.EQUALITY_DELETES): delete_index.add_delete_file(manifest_entry, partition_key=data_file.partition) - elif data_file.content == DataFileContent.EQUALITY_DELETES: - raise ValueError("PyIceberg does not yet support equality deletes: https://github.com/apache/iceberg/issues/6568") else: raise ValueError(f"Unknown DataFileContent ({data_file.content}): {manifest_entry}") diff --git a/pyiceberg/table/delete_file_index.py b/pyiceberg/table/delete_file_index.py index 3f513aabe5..88a8695c14 100644 --- a/pyiceberg/table/delete_file_index.py +++ b/pyiceberg/table/delete_file_index.py @@ -17,11 +17,17 @@ from __future__ import annotations from bisect import bisect_left +from typing import TYPE_CHECKING +from pyiceberg.conversions import from_bytes from pyiceberg.expressions import EqualTo from pyiceberg.expressions.visitors import _InclusiveMetricsEvaluator -from pyiceberg.manifest import INITIAL_SEQUENCE_NUMBER, POSITIONAL_DELETE_SCHEMA, DataFile, ManifestEntry +from pyiceberg.manifest import INITIAL_SEQUENCE_NUMBER, POSITIONAL_DELETE_SCHEMA, DataFile, DataFileContent, ManifestEntry from pyiceberg.typedef import Record +from pyiceberg.types import NestedField + +if TYPE_CHECKING: + from pyiceberg.schema import Schema PATH_FIELD_ID = 2147483546 @@ -59,6 +65,16 @@ def referenced_delete_files(self) -> list[DataFile]: return [data_file for data_file, _ in self._files] +class EqualityDeletes(PositionDeletes): + """Collect equality delete files indexed by the newest data sequence they may affect.""" + + def add(self, delete_file: DataFile, seq_num: int) -> None: + # Equality deletes apply only to data whose sequence is strictly less than + # the delete sequence. Indexing at seq - 1 lets the shared >= lookup encode + # that rule without special cases at lookup time. + super().add(delete_file, seq_num - 1) + + def _has_path_bounds(delete_file: DataFile) -> bool: lower = delete_file.lower_bounds upper = delete_file.upper_bounds @@ -76,6 +92,76 @@ def _applies_to_data_file(delete_file: DataFile, data_file: DataFile) -> bool: return evaluator.eval(delete_file) +def _is_all_null(data_file: DataFile, field_id: int) -> bool: + null_counts = data_file.null_value_counts + value_counts = data_file.value_counts + if not null_counts or not value_counts: + return False + null_count = null_counts.get(field_id) + value_count = value_counts.get(field_id) + return null_count is not None and value_count is not None and null_count == value_count + + +def _has_no_nulls(data_file: DataFile, field_id: int) -> bool: + null_counts = data_file.null_value_counts + return bool(null_counts) and null_counts.get(field_id) == 0 + + +def _contains_null(data_file: DataFile, field: NestedField) -> bool: + if field.required: + return False + null_counts = data_file.null_value_counts + if not null_counts: + return True + null_count = null_counts.get(field.field_id) + return null_count is None or null_count > 0 + + +def _equality_delete_applies_to_data_file(delete_file: DataFile, data_file: DataFile, schema: Schema) -> bool: + """Conservatively prune equality deletes whose metrics cannot match a data file.""" + if not delete_file.equality_ids: + return True + + for field_id in delete_file.equality_ids: + try: + field = schema.find_field(field_id) + except ValueError: + # A dropped field can still exist in older data and delete files. + return True + if not field.field_type.is_primitive: + continue + + if _contains_null(data_file, field) and _contains_null(delete_file, field): + continue + if _is_all_null(data_file, field_id) and _has_no_nulls(delete_file, field_id): + return False + if _is_all_null(delete_file, field_id) and _has_no_nulls(data_file, field_id): + return False + + delete_lower = delete_file.lower_bounds + delete_upper = delete_file.upper_bounds + data_lower = data_file.lower_bounds + data_upper = data_file.upper_bounds + if ( + delete_lower + and delete_upper + and data_lower + and data_upper + and field_id in delete_lower + and field_id in delete_upper + and field_id in data_lower + and field_id in data_upper + ): + field_type = field.field_type + if ( + from_bytes(field_type, delete_upper[field_id]) < from_bytes(field_type, data_lower[field_id]) + or from_bytes(field_type, delete_lower[field_id]) > from_bytes(field_type, data_upper[field_id]) + ): + return False + + return True + + def _referenced_data_file_path(delete_file: DataFile) -> str | None: """Return the path, if the path bounds evaluate to the same location.""" lower_bounds = delete_file.lower_bounds @@ -103,27 +189,30 @@ def _partition_key(spec_id: int, partition: Record | None) -> tuple[int, Record] class DeleteFileIndex: - """Indexes position delete files by partition and by exact data file path.""" + """Index position and equality delete files by their Iceberg applicability rules.""" - def __init__(self) -> None: + def __init__(self, schema: Schema | None = None) -> None: + self._schema = schema self._by_partition: dict[tuple[int, Record], PositionDeletes] = {} self._by_path: dict[str, PositionDeletes] = {} + self._equality_deletes: dict[tuple[int, Record] | None, EqualityDeletes] = {} def is_empty(self) -> bool: - return not self._by_partition and not self._by_path + return not self._by_partition and not self._by_path and not self._equality_deletes def add_delete_file(self, manifest_entry: ManifestEntry, partition_key: Record | None = None) -> None: delete_file = manifest_entry.data_file seq = manifest_entry.sequence_number or INITIAL_SEQUENCE_NUMBER - target_path = _referenced_data_file_path(delete_file) - - if target_path: - deletes = self._by_path.setdefault(target_path, PositionDeletes()) - deletes.add(delete_file, seq) + if delete_file.content == DataFileContent.EQUALITY_DELETES: + # An unpartitioned equality delete is global, including for partitioned + # data files. Partitioned deletes remain scoped to their exact spec/key. + key = _partition_key(delete_file.spec_id or 0, partition_key) if partition_key else None + self._equality_deletes.setdefault(key, EqualityDeletes()).add(delete_file, seq) + elif target_path := _referenced_data_file_path(delete_file): + self._by_path.setdefault(target_path, PositionDeletes()).add(delete_file, seq) else: key = _partition_key(delete_file.spec_id or 0, partition_key) - deletes = self._by_partition.setdefault(key, PositionDeletes()) - deletes.add(delete_file, seq) + self._by_partition.setdefault(key, PositionDeletes()).add(delete_file, seq) def for_data_file(self, seq_num: int, data_file: DataFile, partition_key: Record | None = None) -> set[DataFile]: if self.is_empty(): @@ -133,16 +222,25 @@ def for_data_file(self, seq_num: int, data_file: DataFile, partition_key: Record spec_id = data_file.spec_id or 0 key = _partition_key(spec_id, partition_key) - partition_deletes = self._by_partition.get(key) - if partition_deletes: + if partition_deletes := self._by_partition.get(key): for delete_file in partition_deletes.filter_by_seq(seq_num): if _applies_to_data_file(delete_file, data_file): deletes.add(delete_file) - path_deletes = self._by_path.get(data_file.file_path) - if path_deletes: + if path_deletes := self._by_path.get(data_file.file_path): deletes.update(path_deletes.filter_by_seq(seq_num)) + candidates: list[DataFile] = [] + if partition_equality_deletes := self._equality_deletes.get(key): + candidates.extend(partition_equality_deletes.filter_by_seq(seq_num)) + if global_equality_deletes := self._equality_deletes.get(None): + candidates.extend(global_equality_deletes.filter_by_seq(seq_num)) + + for delete_file in candidates: + if self._schema and not _equality_delete_applies_to_data_file(delete_file, data_file, self._schema): + continue + deletes.add(delete_file) + return deletes def referenced_delete_files(self) -> list[DataFile]: @@ -154,4 +252,7 @@ def referenced_delete_files(self) -> list[DataFile]: for deletes in self._by_path.values(): data_files.extend(deletes.referenced_delete_files()) + for deletes in self._equality_deletes.values(): + data_files.extend(deletes.referenced_delete_files()) + return data_files diff --git a/pyiceberg/table/update/validate.py b/pyiceberg/table/update/validate.py index df8506aab4..f80887b362 100644 --- a/pyiceberg/table/update/validate.py +++ b/pyiceberg/table/update/validate.py @@ -248,13 +248,13 @@ def _added_delete_files( DeleteFileIndex """ if table.format_version < 2: - return DeleteFileIndex() + return DeleteFileIndex(table.schema()) manifests, snapshot_ids = _validation_history( table, parent_snapshot, starting_snapshot, VALIDATE_ADDED_DELETE_FILES_OPERATIONS, ManifestContent.DELETES ) - dfi = DeleteFileIndex() + dfi = DeleteFileIndex(table.schema()) for manifest in manifests: for entry in manifest.fetch_manifest_entry(table.io, discard_deleted=True): diff --git a/tests/table/test_delete_file_index.py b/tests/table/test_delete_file_index.py index 09dd9ac81b..7ee3c34552 100644 --- a/tests/table/test_delete_file_index.py +++ b/tests/table/test_delete_file_index.py @@ -16,12 +16,22 @@ # under the License. import pytest +from pyiceberg.conversions import to_bytes from pyiceberg.manifest import DataFile, DataFileContent, FileFormat, ManifestEntry, ManifestEntryStatus +from pyiceberg.schema import Schema from pyiceberg.table.delete_file_index import PATH_FIELD_ID, DeleteFileIndex, PositionDeletes from pyiceberg.typedef import Record +from pyiceberg.types import IntegerType, LongType, NestedField, StringType -def _create_data_file(file_path: str = "s3://bucket/data.parquet", spec_id: int = 0) -> DataFile: +def _create_data_file( + file_path: str = "s3://bucket/data.parquet", + spec_id: int = 0, + lower_bounds: dict[int, bytes] | None = None, + upper_bounds: dict[int, bytes] | None = None, + null_value_counts: dict[int, int] | None = None, + value_counts: dict[int, int] | None = None, +) -> DataFile: data_file = DataFile.from_args( content=DataFileContent.DATA, file_path=file_path, @@ -29,6 +39,10 @@ def _create_data_file(file_path: str = "s3://bucket/data.parquet", spec_id: int partition=Record(), record_count=100, file_size_in_bytes=1000, + lower_bounds=lower_bounds, + upper_bounds=upper_bounds, + null_value_counts=null_value_counts, + value_counts=value_counts, ) data_file._spec_id = spec_id return data_file @@ -51,6 +65,33 @@ def _create_positional_delete( return ManifestEntry.from_args(status=ManifestEntryStatus.ADDED, sequence_number=sequence_number, data_file=delete_file) +def _create_equality_delete( + sequence_number: int = 1, + spec_id: int = 0, + partition: Record | None = None, + equality_ids: list[int] | None = None, + lower_bounds: dict[int, bytes] | None = None, + upper_bounds: dict[int, bytes] | None = None, + null_value_counts: dict[int, int] | None = None, + value_counts: dict[int, int] | None = None, +) -> ManifestEntry: + delete_file = DataFile.from_args( + content=DataFileContent.EQUALITY_DELETES, + file_path=f"s3://bucket/eq-delete-{sequence_number}.parquet", + file_format=FileFormat.PARQUET, + partition=partition or Record(), + record_count=10, + file_size_in_bytes=100, + equality_ids=equality_ids or [1, 2], + lower_bounds=lower_bounds, + upper_bounds=upper_bounds, + null_value_counts=null_value_counts, + value_counts=value_counts, + ) + delete_file._spec_id = spec_id + return ManifestEntry.from_args(status=ManifestEntryStatus.ADDED, sequence_number=sequence_number, data_file=delete_file) + + def _create_partition_delete(sequence_number: int = 1, spec_id: int = 0, partition: Record | None = None) -> ManifestEntry: delete_file = DataFile.from_args( content=DataFileContent.POSITION_DELETES, @@ -103,6 +144,20 @@ def test_sequence_number_filtering() -> None: assert len(index.for_data_file(7, data_file)) == 0 +def test_equality_delete_sequence_number_is_strictly_greater() -> None: + index = DeleteFileIndex() + + index.add_delete_file(_create_equality_delete(sequence_number=2)) + index.add_delete_file(_create_equality_delete(sequence_number=4)) + + data_file = _create_data_file() + + assert len(index.for_data_file(1, data_file)) == 2 + assert len(index.for_data_file(2, data_file)) == 1 + assert len(index.for_data_file(3, data_file)) == 1 + assert len(index.for_data_file(4, data_file)) == 0 + + def test_path_specific_deletes() -> None: index = DeleteFileIndex() @@ -187,3 +242,138 @@ def test_record_equality_for_partition_lookup() -> None: assert len(index.for_data_file(1, data_file, partition_b)) == 1 assert len(index.for_data_file(1, data_file, partition_c)) == 0 + + +def test_equality_delete_sequence_number_filtering() -> None: + index = DeleteFileIndex() + equality_delete = _create_equality_delete(sequence_number=2) + index.add_delete_file(equality_delete) + + data_file = _create_data_file() + assert equality_delete.data_file in index.for_data_file(1, data_file) + assert equality_delete.data_file not in index.for_data_file(2, data_file) + assert equality_delete.data_file not in index.for_data_file(3, data_file) + + +def test_equality_and_position_delete_sequence_semantics_differ() -> None: + data_file = _create_data_file() + position_delete = _create_positional_delete(sequence_number=10) + equality_delete = _create_equality_delete(sequence_number=10) + index = DeleteFileIndex() + index.add_delete_file(position_delete) + index.add_delete_file(equality_delete) + + assert index.for_data_file(10, data_file) == {position_delete.data_file} + assert index.for_data_file(9, data_file) == {position_delete.data_file, equality_delete.data_file} + assert index.for_data_file(11, data_file) == set() + + +def test_global_equality_deletes_apply_to_partitioned_data() -> None: + index = DeleteFileIndex() + global_delete = _create_equality_delete(sequence_number=10) + partition_delete = _create_equality_delete(sequence_number=20) + partition_a = Record(1) + partition_b = Record(2) + index.add_delete_file(global_delete) + index.add_delete_file(partition_delete, partition_a) + + data_file = _create_data_file() + assert index.for_data_file(1, data_file, partition_a) == {global_delete.data_file, partition_delete.data_file} + assert index.for_data_file(1, data_file, partition_b) == {global_delete.data_file} + + +def test_equality_delete_metrics_filtering() -> None: + index = DeleteFileIndex(Schema(NestedField(1, "id", IntegerType(), required=True))) + equality_delete = _create_equality_delete( + sequence_number=100, + equality_ids=[1], + lower_bounds={1: to_bytes(IntegerType(), 10)}, + upper_bounds={1: to_bytes(IntegerType(), 20)}, + ) + index.add_delete_file(equality_delete) + + before = _create_data_file( + lower_bounds={1: to_bytes(IntegerType(), 0)}, upper_bounds={1: to_bytes(IntegerType(), 5)} + ) + overlap = _create_data_file( + lower_bounds={1: to_bytes(IntegerType(), 15)}, upper_bounds={1: to_bytes(IntegerType(), 25)} + ) + after = _create_data_file( + lower_bounds={1: to_bytes(IntegerType(), 25)}, upper_bounds={1: to_bytes(IntegerType(), 30)} + ) + assert index.for_data_file(1, before) == set() + assert index.for_data_file(1, overlap) == {equality_delete.data_file} + assert index.for_data_file(1, after) == set() + + +@pytest.mark.parametrize( + ("delete_nulls", "data_nulls"), + [((10, 10), (0, 100)), ((0, 10), (100, 100))], +) +def test_equality_delete_prunes_disjoint_null_populations( + delete_nulls: tuple[int, int], data_nulls: tuple[int, int] +) -> None: + index = DeleteFileIndex(Schema(NestedField(1, "id", IntegerType(), required=False))) + equality_delete = _create_equality_delete( + sequence_number=10, + equality_ids=[1], + null_value_counts={1: delete_nulls[0]}, + value_counts={1: delete_nulls[1]}, + ) + index.add_delete_file(equality_delete) + data_file = _create_data_file(null_value_counts={1: data_nulls[0]}, value_counts={1: data_nulls[1]}) + assert index.for_data_file(1, data_file) == set() + + +def test_equality_delete_metrics_after_int_to_long_promotion() -> None: + index = DeleteFileIndex(Schema(NestedField(1, "id", LongType(), required=True))) + equality_delete = _create_equality_delete( + sequence_number=100, + equality_ids=[1], + lower_bounds={1: to_bytes(IntegerType(), 10)}, + upper_bounds={1: to_bytes(IntegerType(), 20)}, + ) + index.add_delete_file(equality_delete) + before = _create_data_file( + lower_bounds={1: to_bytes(IntegerType(), 0)}, upper_bounds={1: to_bytes(IntegerType(), 5)} + ) + overlap = _create_data_file( + lower_bounds={1: to_bytes(IntegerType(), 15)}, upper_bounds={1: to_bytes(IntegerType(), 25)} + ) + assert index.for_data_file(1, before) == set() + assert index.for_data_file(1, overlap) == {equality_delete.data_file} + + +def test_equality_delete_dropped_field_is_not_pruned() -> None: + index = DeleteFileIndex(Schema(NestedField(2, "other", StringType(), required=True))) + equality_delete = _create_equality_delete( + sequence_number=10, + equality_ids=[1], + lower_bounds={1: to_bytes(IntegerType(), 10)}, + upper_bounds={1: to_bytes(IntegerType(), 20)}, + ) + index.add_delete_file(equality_delete) + data_file = _create_data_file( + lower_bounds={1: to_bytes(IntegerType(), 15)}, upper_bounds={1: to_bytes(IntegerType(), 25)} + ) + assert index.for_data_file(1, data_file) == {equality_delete.data_file} + + +def test_equality_delete_is_not_pruned_when_both_files_contain_nulls() -> None: + index = DeleteFileIndex(Schema(NestedField(1, "id", IntegerType(), required=False))) + equality_delete = _create_equality_delete( + sequence_number=100, + equality_ids=[1], + lower_bounds={1: to_bytes(IntegerType(), 10)}, + upper_bounds={1: to_bytes(IntegerType(), 20)}, + null_value_counts={1: 1}, + value_counts={1: 10}, + ) + index.add_delete_file(equality_delete) + data_file = _create_data_file( + lower_bounds={1: to_bytes(IntegerType(), 0)}, + upper_bounds={1: to_bytes(IntegerType(), 5)}, + null_value_counts={1: 1}, + value_counts={1: 100}, + ) + assert index.for_data_file(1, data_file) == {equality_delete.data_file} From d342bdc3a119e3a230dffb11b3b114b6ae9479a7 Mon Sep 17 00:00:00 2001 From: Nico Ritschel Date: Mon, 17 Aug 2026 20:16:49 -0700 Subject: [PATCH 2/6] Apply equality deletes in Arrow scans --- pyiceberg/io/pyarrow.py | 231 +++++++++++++++++++-- pyiceberg/table/__init__.py | 11 +- tests/catalog/test_scan_planning_models.py | 12 +- tests/io/test_pyarrow.py | 231 +++++++++++++++++++++ 4 files changed, 457 insertions(+), 28 deletions(-) diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py index c36f1639d9..f7f4d10898 100644 --- a/pyiceberg/io/pyarrow.py +++ b/pyiceberg/io/pyarrow.py @@ -1157,6 +1157,165 @@ def _read_deletes(io: FileIO, data_file: DataFile) -> dict[str, pa.ChunkedArray] raise ValueError(f"Delete file format not supported: {data_file.file_format}") +def _read_equality_deletes(io: FileIO, delete_file: DataFile) -> pa.Table: + """Read an equality delete file while preserving its field-ID metadata.""" + if delete_file.file_format not in (FileFormat.PARQUET, FileFormat.ORC): + raise ValueError(f"Equality delete file format not supported: {delete_file.file_format}") + if not delete_file.equality_ids: + raise ValueError(f"Equality delete file has no equality IDs: {delete_file.file_path}") + + with io.new_input(delete_file.file_path).open() as fin: + fragment = _get_file_format(delete_file.file_format, pre_buffer=True, buffer_size=ONE_MEGABYTE).make_fragment(fin) + physical_schema = fragment.physical_schema + equality_ids = set(delete_file.equality_ids) + projected_columns = [field.name for field in physical_schema if _get_field_id(field) in equality_ids] + # Older writers may not retain field IDs in the physical schema. Reading + # all columns is the conservative fallback; alignment below still rejects + # a file that cannot supply every declared equality field. + columns = projected_columns if len(projected_columns) == len(equality_ids) else None + # Delete files are already read concurrently by PyIceberg's executor; + # disable nested Arrow threading to avoid oversubscription per file. + return ds.Scanner.from_fragment( + fragment=fragment, schema=physical_schema, columns=columns, use_threads=False + ).to_table() + + +def _column_name_for_field_id(table: pa.Table, field_id: int, schema: Schema) -> str | None: + try: + current_name = schema.find_field(field_id).name + except ValueError: + current_name = None + if current_name is not None and current_name in table.column_names: + return current_name + for field in table.schema: + if _get_field_id(field) == field_id: + return field.name + return None + + +def _align_equality_delete_table( + data: pa.Table, deletes: pa.Table, equality_ids: Iterable[int], table_schema: Schema +) -> tuple[pa.Table, pa.Table, list[str]]: + """Align equality keys by stable field ID across rename, drop, and promotion.""" + join_keys: list[str] = [] + for field_id in equality_ids: + data_name = _column_name_for_field_id(data, field_id, table_schema) + delete_name = _column_name_for_field_id(deletes, field_id, table_schema) + if delete_name is None: + raise ValueError(f"Equality delete file is missing field ID {field_id}") + + if data_name is None: + # A field absent from this data file is projected as null under Iceberg + # schema evolution rules. It may therefore match a null delete key. + data = data.append_column(delete_name, pa.nulls(data.num_rows, type=deletes.schema.field(delete_name).type)) + data_name = delete_name + if delete_name != data_name: + names = list(deletes.column_names) + names[names.index(delete_name)] = data_name + deletes = deletes.rename_columns(names) + + data_type = data.schema.field(data_name).type + if deletes.schema.field(data_name).type != data_type: + deletes = deletes.set_column( + deletes.schema.get_field_index(data_name), data_name, pc.cast(deletes[data_name], data_type) + ) + join_keys.append(data_name) + + return data, deletes, join_keys + + +def _fill_null_value(data_type: pa.DataType) -> Any: + if pa.types.is_integer(data_type) or pa.types.is_floating(data_type) or pa.types.is_decimal(data_type): + return 0 + if pa.types.is_boolean(data_type): + return False + if pa.types.is_string(data_type) or pa.types.is_large_string(data_type): + return "" + if pa.types.is_binary(data_type) or pa.types.is_large_binary(data_type): + return b"" + if pa.types.is_fixed_size_binary(data_type): + return b"\x00" * data_type.byte_width + if isinstance(data_type, pa.UuidType): + return pa.scalar(b"\x00" * 16, type=data_type) + if pa.types.is_timestamp(data_type) or pa.types.is_date(data_type) or pa.types.is_time(data_type): + return pa.scalar(0, type=data_type) + if storage_type := getattr(data_type, "storage_type", None): + return _fill_null_value(storage_type) + raise TypeError(f"Unsupported equality field type: {data_type}") + + +def _null_safe_left_anti_join(data: pa.Table, deletes: pa.Table, keys: list[str]) -> pa.Table: + """Return data rows without an Iceberg null/NaN-safe typed key match.""" + if data.num_rows == 0 or deletes.num_rows == 0: + return data + + data_join = data + delete_join = deletes.select(keys) + join_columns: list[str] = [] + temporary_columns: list[str] = [] + temporary_index = 0 + + def temporary_name(kind: str) -> str: + nonlocal temporary_index + while True: + name = f"__pyiceberg_equality_{kind}_{temporary_index}" + temporary_index += 1 + if name not in data_join.column_names and name not in delete_join.column_names: + return name + + for key in keys: + key_type = data_join.schema.field(key).type + data_values = data_join[key] + delete_values = delete_join[key] + if storage_type := getattr(key_type, "storage_type", None): + data_values = data_values.cast(storage_type) + delete_values = delete_values.cast(storage_type) + fill_type = storage_type + else: + fill_type = key_type + + null_key = temporary_name("null") + value_key = temporary_name("value") + fill_value = _fill_null_value(fill_type) + data_join = data_join.append_column(null_key, pc.is_null(data_values)) + delete_join = delete_join.append_column(null_key, pc.is_null(delete_values)) + data_join = data_join.append_column(value_key, pc.fill_null(data_values, fill_value)) + delete_join = delete_join.append_column(value_key, pc.fill_null(delete_values, fill_value)) + join_columns.extend((null_key, value_key)) + temporary_columns.extend((null_key, value_key)) + + if pa.types.is_floating(key_type): + nan_key = temporary_name("nan") + data_join = data_join.append_column(nan_key, pc.fill_null(pc.is_nan(data_join[key]), False)) + delete_join = delete_join.append_column(nan_key, pc.fill_null(pc.is_nan(delete_join[key]), False)) + join_columns.append(nan_key) + temporary_columns.append(nan_key) + + joined = data_join.join( + delete_join.select(join_columns), keys=join_columns, join_type="left anti", use_threads=False + ) + return joined.drop(temporary_columns) + + +def _apply_equality_deletes( + data: pa.Table, equality_groups: dict[frozenset[int], list[pa.Table]], table_schema: Schema +) -> pa.Table: + """Apply each distinct equality-key layout once to a record batch.""" + for equality_ids, delete_tables in equality_groups.items(): + aligned_deletes: list[pa.Table] = [] + join_keys: list[str] | None = None + for delete_table in delete_tables: + data, aligned, keys = _align_equality_delete_table(data, delete_table, equality_ids, table_schema) + join_keys = join_keys or keys + aligned_deletes.append(aligned.select(keys)) + if aligned_deletes: + deletes = pa.concat_tables(aligned_deletes, promote_options="permissive") + data = _null_safe_left_anti_join(data, deletes, join_keys or []) + if data.num_rows == 0: + break + return data + + def _combine_positional_deletes(positional_deletes: list[pa.ChunkedArray], start_index: int, end_index: int) -> pa.Array: if len(positional_deletes) == 1: all_chunks = positional_deletes[0] @@ -1640,6 +1799,7 @@ def _task_to_record_batches( format_version: TableVersion = TableProperties.DEFAULT_FORMAT_VERSION, downcast_ns_timestamp_to_us: bool | None = None, dictionary_columns: tuple[str, ...] = (), + equality_delete_tables: dict[DataFile, pa.Table] | None = None, ) -> Iterator[pa.RecordBatch]: format_kwargs: dict[str, Any] = {"pre_buffer": True, "buffer_size": ONE_MEGABYTE * 8} if dictionary_columns and task.file.file_format == FileFormat.PARQUET: @@ -1672,14 +1832,25 @@ def _task_to_record_batches( bound_file_filter = bind(file_schema, translated_row_filter, case_sensitive=case_sensitive) pyarrow_filter = expression_to_pyarrow(bound_file_filter, file_schema) - file_project_schema = prune_columns(file_schema, projected_field_ids, select_full_types=False) + equality_groups: dict[frozenset[int], list[pa.Table]] = {} + equality_field_ids: set[int] = set() + if equality_delete_tables: + for delete_file in task.delete_files: + if delete_file.content != DataFileContent.EQUALITY_DELETES: + continue + equality_table = equality_delete_tables.get(delete_file) + if equality_table is not None: + equality_ids = frozenset(delete_file.equality_ids or ()) + equality_field_ids.update(equality_ids) + equality_groups.setdefault(equality_ids, []).append(equality_table) + file_project_schema = prune_columns(file_schema, projected_field_ids.union(equality_field_ids), select_full_types=False) fragment_scanner = ds.Scanner.from_fragment( fragment=fragment, schema=physical_schema, # This will push down the query to Arrow. # But in case there are positional deletes, we have to apply them first - filter=pyarrow_filter if not positional_deletes else None, + filter=pyarrow_filter if not positional_deletes and not equality_groups else None, columns=[col.name for col in file_project_schema.columns], ) @@ -1694,16 +1865,18 @@ def _task_to_record_batches( # Create the mask of indices that we're interested in indices = _combine_positional_deletes(positional_deletes, current_index, current_index + len(batch)) current_batch = current_batch.take(indices) - if pyarrow_filter is not None: - # Temporary fix until PyArrow 21 is the minimum supported version - # (https://github.com/apache/arrow/pull/46057): RecordBatch.filter raises - # IndexError on PyArrow <21 when the result is empty; Table.filter does not. - table = pa.Table.from_batches([current_batch]) - table = table.filter(pyarrow_filter) - if table.num_rows == 0: - current_batch = current_batch.slice(0, 0) - else: - current_batch = table.combine_chunks().to_batches()[0] + + if equality_groups and current_batch.num_rows > 0: + table = pa.Table.from_batches([current_batch]) + table = _apply_equality_deletes(table, equality_groups, table_schema) + current_batch = table.combine_chunks().to_batches()[0] if table.num_rows > 0 else current_batch.slice(0, 0) + + if (positional_deletes or equality_groups) and pyarrow_filter is not None and current_batch.num_rows > 0: + # Temporary fix until PyArrow 21 is the minimum supported version + # (https://github.com/apache/arrow/pull/46057): RecordBatch.filter raises + # IndexError on PyArrow <21 when the result is empty; Table.filter does not. + table = pa.Table.from_batches([current_batch]).filter(pyarrow_filter) + current_batch = table.combine_chunks().to_batches()[0] if table.num_rows > 0 else current_batch.slice(0, 0) # skip empty batches if current_batch.num_rows == 0: @@ -1721,7 +1894,11 @@ def _task_to_record_batches( def _read_all_delete_files(io: FileIO, tasks: Iterable[FileScanTask]) -> dict[str, list[ChunkedArray]]: deletes_per_file: dict[str, list[ChunkedArray]] = {} - unique_deletes = set(itertools.chain.from_iterable([task.delete_files for task in tasks])) + unique_deletes = { + delete_file + for delete_file in itertools.chain.from_iterable(task.delete_files for task in tasks) + if delete_file.content == DataFileContent.POSITION_DELETES + } if len(unique_deletes) > 0: executor = ExecutorFactory.get_or_create() deletes_per_files: Iterator[dict[str, ChunkedArray]] = executor.map( @@ -1738,6 +1915,20 @@ def _read_all_delete_files(io: FileIO, tasks: Iterable[FileScanTask]) -> dict[st return deletes_per_file +def _read_all_equality_delete_files(io: FileIO, tasks: Iterable[FileScanTask]) -> dict[DataFile, pa.Table]: + equality_delete_files = { + delete_file + for delete_file in itertools.chain.from_iterable(task.delete_files for task in tasks) + if delete_file.content == DataFileContent.EQUALITY_DELETES + } + if not equality_delete_files: + return {} + + executor = ExecutorFactory.get_or_create() + results = executor.map(lambda delete_file: _read_equality_deletes(io, delete_file), equality_delete_files) + return dict(zip(equality_delete_files, results, strict=True)) + + class ArrowScan: _table_metadata: TableMetadata _io: FileIO @@ -1841,7 +2032,9 @@ def to_record_batches(self, tasks: Iterable[FileScanTask]) -> Iterator[pa.Record ResolveError: When a required field cannot be found in the file ValueError: When a field type in the file cannot be projected to the schema type """ - deletes_per_file = _read_all_delete_files(self._io, tasks) + planned_tasks = list(tasks) + deletes_per_file = _read_all_delete_files(self._io, planned_tasks) + equality_deletes = _read_all_equality_delete_files(self._io, planned_tasks) total_row_count = 0 executor = ExecutorFactory.get_or_create() @@ -1850,10 +2043,10 @@ def batches_for_task(task: FileScanTask) -> list[pa.RecordBatch]: # Materialize the iterator here to ensure execution happens within the executor. # Otherwise, the iterator would be lazily consumed later (in the main thread), # defeating the purpose of using executor.map. - return list(self._record_batches_from_scan_tasks_and_deletes([task], deletes_per_file)) + return list(self._record_batches_from_scan_tasks_and_deletes([task], deletes_per_file, equality_deletes)) limit_reached = False - for batches in executor.map(batches_for_task, tasks): + for batches in executor.map(batches_for_task, planned_tasks): for batch in batches: current_batch_size = len(batch) if self._limit is not None and total_row_count + current_batch_size >= self._limit: @@ -1870,7 +2063,10 @@ def batches_for_task(task: FileScanTask) -> list[pa.RecordBatch]: break def _record_batches_from_scan_tasks_and_deletes( - self, tasks: Iterable[FileScanTask], deletes_per_file: dict[str, list[ChunkedArray]] + self, + tasks: Iterable[FileScanTask], + deletes_per_file: dict[str, list[ChunkedArray]], + equality_deletes: dict[DataFile, pa.Table] | None = None, ) -> Iterator[pa.RecordBatch]: total_row_count = 0 for task in tasks: @@ -1890,6 +2086,7 @@ def _record_batches_from_scan_tasks_and_deletes( self._table_metadata.format_version, self._downcast_ns_timestamp_to_us, self._dictionary_columns, + equality_deletes, ) for batch in batches: if self._limit is not None: diff --git a/pyiceberg/table/__init__.py b/pyiceberg/table/__init__.py index b545be702e..31e7689bc0 100644 --- a/pyiceberg/table/__init__.py +++ b/pyiceberg/table/__init__.py @@ -2255,19 +2255,13 @@ def from_rest_response( Returns: A FileScanTask with the converted data and delete files. - Raises: - NotImplementedError: If equality delete files are encountered. """ - from pyiceberg.catalog.rest.scan_planning import RESTEqualityDeleteFile - data_file = _rest_file_to_data_file(rest_task.data_file) resolved_deletes: set[DataFile] = set() if rest_task.delete_file_references: for idx in rest_task.delete_file_references: delete_file = delete_files[idx] - if isinstance(delete_file, RESTEqualityDeleteFile): - raise NotImplementedError(f"PyIceberg does not yet support equality deletes: {delete_file.file_path}") resolved_deletes.add(_rest_file_to_data_file(delete_file)) return FileScanTask( @@ -2279,7 +2273,7 @@ def from_rest_response( def _rest_file_to_data_file(rest_file: RESTContentFile) -> DataFile: """Convert a REST content file to a manifest DataFile.""" - from pyiceberg.catalog.rest.scan_planning import RESTDataFile + from pyiceberg.catalog.rest.scan_planning import RESTDataFile, RESTEqualityDeleteFile if isinstance(rest_file, RESTDataFile): column_sizes = rest_file.column_sizes.to_dict() if rest_file.column_sizes else None @@ -2292,6 +2286,8 @@ def _rest_file_to_data_file(rest_file: RESTContentFile) -> DataFile: null_value_counts = None nan_value_counts = None + equality_ids = rest_file.equality_ids if isinstance(rest_file, RESTEqualityDeleteFile) else None + data_file = DataFile.from_args( content=DataFileContent.from_rest_type(rest_file.content), file_path=rest_file.file_path, @@ -2305,6 +2301,7 @@ def _rest_file_to_data_file(rest_file: RESTContentFile) -> DataFile: nan_value_counts=nan_value_counts, split_offsets=rest_file.split_offsets, sort_order_id=rest_file.sort_order_id, + equality_ids=equality_ids, ) data_file.spec_id = rest_file.spec_id return data_file diff --git a/tests/catalog/test_scan_planning_models.py b/tests/catalog/test_scan_planning_models.py index f2c80cfb9b..dc147e67c8 100644 --- a/tests/catalog/test_scan_planning_models.py +++ b/tests/catalog/test_scan_planning_models.py @@ -38,7 +38,7 @@ ValueMap, ) from pyiceberg.expressions import AlwaysTrue, EqualTo, Reference -from pyiceberg.manifest import FileFormat +from pyiceberg.manifest import DataFileContent, FileFormat TEST_URI = "https://iceberg-test-catalog/" @@ -545,7 +545,7 @@ def test_plan_scan_cancelled(rest_scan_catalog: RestCatalog, requests_mock: Mock list(rest_scan_catalog.plan_scan(("db", "tbl"), request)) -def test_plan_scan_equality_deletes_not_supported(rest_scan_catalog: RestCatalog, requests_mock: Mocker) -> None: +def test_plan_scan_with_equality_deletes(rest_scan_catalog: RestCatalog, requests_mock: Mocker) -> None: file_one = _rest_data_file(file_path="s3://bucket/tbl/data/file1.parquet") equality_delete = _rest_equality_delete_file(equality_ids=[1, 2]) requests_mock.post( @@ -566,5 +566,9 @@ def test_plan_scan_equality_deletes_not_supported(rest_scan_catalog: RestCatalog ) request = PlanTableScanRequest() - with pytest.raises(NotImplementedError, match="PyIceberg does not yet support equality deletes"): - list(rest_scan_catalog.plan_scan(("db", "tbl"), request)) + tasks = list(rest_scan_catalog.plan_scan(("db", "tbl"), request)) + assert len(tasks) == 1 + assert len(tasks[0].delete_files) == 1 + delete_file = next(iter(tasks[0].delete_files)) + assert delete_file.content == DataFileContent.EQUALITY_DELETES + assert delete_file.equality_ids == [1, 2] diff --git a/tests/io/test_pyarrow.py b/tests/io/test_pyarrow.py index b31c18949b..ea67c72e6b 100644 --- a/tests/io/test_pyarrow.py +++ b/tests/io/test_pyarrow.py @@ -74,6 +74,7 @@ _check_pyarrow_schema_compatible, _ConvertToArrowSchema, _determine_partitions, + _null_safe_left_anti_join, _primitive_to_physical, _read_deletes, _task_to_record_batches, @@ -117,6 +118,7 @@ TimestampType, TimestamptzType, TimeType, + UUIDType, ) from tests.catalog.test_base import InMemoryCatalog from tests.conftest import UNIFIED_AWS_SESSION_PROPERTIES @@ -1914,6 +1916,235 @@ def test_delete_duplicates(deletes_file: str, request: pytest.FixtureRequest, ta assert str(with_deletes) == expected_str +def test_composite_equality_delete_with_nulls_and_hidden_projection(tmp_path: Path) -> None: + table_schema = Schema( + NestedField(1, "id", IntegerType(), required=True), + NestedField(2, "sub_id", StringType(), required=False), + NestedField(3, "payload", StringType(), required=False), + schema_id=1, + ) + projected_schema = Schema(table_schema.find_field(3), schema_id=1) + equality_schema = Schema(table_schema.find_field(1), table_schema.find_field(2), schema_id=1) + + data_path = str(tmp_path / "data.parquet") + pq.write_table( + pa.table( + { + "id": pa.array([1, 2, 3, 4], type=pa.int32()), + "sub_id": ["a", "b", None, None], + "payload": ["keep-a", "delete-b", "delete-null", "keep-null"], + }, + schema=schema_to_pyarrow(table_schema), + ), + data_path, + ) + + delete_path = str(tmp_path / "eq-delete.parquet") + pq.write_table( + pa.table( + { + "id": pa.array([2, 3], type=pa.int32()), + "sub_id": ["b", None], + }, + schema=schema_to_pyarrow(equality_schema), + ), + delete_path, + ) + + data_file = DataFile.from_args( + content=DataFileContent.DATA, + file_path=data_path, + file_format=FileFormat.PARQUET, + partition=Record(), + record_count=4, + file_size_in_bytes=Path(data_path).stat().st_size, + ) + data_file.spec_id = 0 + equality_delete_file = DataFile.from_args( + content=DataFileContent.EQUALITY_DELETES, + file_path=delete_path, + file_format=FileFormat.PARQUET, + partition=Record(), + record_count=2, + file_size_in_bytes=Path(delete_path).stat().st_size, + equality_ids=[1, 2], + ) + equality_delete_file.spec_id = 0 + + result = ArrowScan( + table_metadata=TableMetadataV2( + location=str(tmp_path), + last_column_id=3, + format_version=2, + current_schema_id=1, + schemas=[table_schema], + partition_specs=[PartitionSpec()], + ), + io=PyArrowFileIO(), + projected_schema=projected_schema, + row_filter=AlwaysTrue(), + ).to_table(tasks=[FileScanTask(data_file=data_file, delete_files={equality_delete_file})]) + + assert result.to_pydict() == {"payload": ["keep-a", "keep-null"]} + + +def _scan_equality_delete_files( + schema: Schema, + data_path: Path, + delete_paths: list[Path], + equality_ids: list[int], + file_format: FileFormat = FileFormat.PARQUET, +) -> pa.Table: + data_file = DataFile.from_args( + content=DataFileContent.DATA, + file_path=str(data_path), + file_format=file_format, + partition=Record(), + record_count=1, + file_size_in_bytes=data_path.stat().st_size, + ) + data_file.spec_id = 0 + delete_files = { + DataFile.from_args( + content=DataFileContent.EQUALITY_DELETES, + file_path=str(delete_path), + file_format=file_format, + partition=Record(), + record_count=1, + file_size_in_bytes=delete_path.stat().st_size, + equality_ids=equality_ids, + ) + for delete_path in delete_paths + } + return ArrowScan( + table_metadata=TableMetadataV2( + location=str(data_path.parent), + last_column_id=max(schema.field_ids), + format_version=2, + current_schema_id=schema.schema_id, + schemas=[schema], + partition_specs=[PartitionSpec()], + ), + io=PyArrowFileIO(), + projected_schema=schema, + row_filter=AlwaysTrue(), + ).to_table(tasks=(task for task in [FileScanTask(data_file=data_file, delete_files=delete_files)])) + + +def test_equality_delete_treats_nan_as_equal(tmp_path: Path) -> None: + schema = Schema(NestedField(1, "value", FloatType(), required=False), schema_id=1) + data_path = tmp_path / "data.parquet" + delete_path = tmp_path / "delete.parquet" + pq.write_table( + pa.table({"value": pa.array([1.0, float("nan")], type=pa.float32())}, schema=schema_to_pyarrow(schema)), + data_path, + ) + pq.write_table(pa.table({"value": pa.array([float("nan")], type=pa.float32())}), delete_path) + + result = _scan_equality_delete_files(schema, data_path, [delete_path], [1]) + assert result.column("value").to_pylist() == [1.0] + + +def test_equality_delete_treats_uuid_nulls_as_equal(tmp_path: Path) -> None: + schema = Schema(NestedField(1, "id", UUIDType(), required=False), schema_id=1) + data_path = tmp_path / "data.parquet" + delete_path = tmp_path / "delete.parquet" + kept = uuid.UUID("00000000-0000-0000-0000-000000000001") + pq.write_table( + pa.table({"id": pa.array([kept, None], type=pa.uuid())}, schema=schema_to_pyarrow(schema)), data_path + ) + pq.write_table(pa.table({"id": pa.array([None], type=pa.uuid())}), delete_path) + + result = _scan_equality_delete_files(schema, data_path, [delete_path], [1]) + assert result.column("id").to_pylist() == [kept] + + +def test_equality_delete_aligns_renamed_fields_by_id(tmp_path: Path) -> None: + current_schema = Schema(NestedField(1, "id", IntegerType(), required=True), schema_id=1) + renamed_schema = Schema(NestedField(1, "old_id", IntegerType(), required=True), schema_id=1) + data_path = tmp_path / "data.parquet" + current_delete_path = tmp_path / "current-delete.parquet" + renamed_delete_path = tmp_path / "renamed-delete.parquet" + pq.write_table( + pa.table({"id": pa.array([1, 2, 3], type=pa.int32())}, schema=schema_to_pyarrow(current_schema)), data_path + ) + pq.write_table( + pa.table({"id": pa.array([2], type=pa.int32())}, schema=schema_to_pyarrow(current_schema)), + current_delete_path, + ) + pq.write_table( + pa.table({"old_id": pa.array([3], type=pa.int32())}, schema=schema_to_pyarrow(renamed_schema)), + renamed_delete_path, + ) + + result = _scan_equality_delete_files( + current_schema, data_path, [current_delete_path, renamed_delete_path], [1] + ) + assert result.column("id").to_pylist() == [1] + + +def test_equality_delete_dropped_field_does_not_match_a_partial_key(tmp_path: Path) -> None: + current_schema = Schema(NestedField(1, "id", IntegerType(), required=True), schema_id=1) + delete_schema = Schema( + NestedField(1, "id", IntegerType(), required=True), + NestedField(2, "dropped", StringType(), required=False), + schema_id=1, + ) + data_path = tmp_path / "data.parquet" + delete_path = tmp_path / "delete.parquet" + pq.write_table( + pa.table({"id": pa.array([1], type=pa.int32())}, schema=schema_to_pyarrow(current_schema)), data_path + ) + pq.write_table( + pa.table( + {"id": pa.array([1], type=pa.int32()), "dropped": ["not-null"]}, + schema=schema_to_pyarrow(delete_schema), + ), + delete_path, + ) + + result = _scan_equality_delete_files(current_schema, data_path, [delete_path], [1, 2]) + assert result.column("id").to_pylist() == [1] + + +def test_orc_equality_delete(tmp_path: Path) -> None: + schema = Schema(NestedField(1, "id", IntegerType(), required=True), schema_id=1) + data_path = tmp_path / "data.orc" + delete_path = tmp_path / "delete.orc" + orc.write_table( + pa.table({"id": pa.array([1, 2, 3], type=pa.int32())}, schema=schema_to_pyarrow(schema)), data_path + ) + orc.write_table(pa.table({"id": pa.array([2], type=pa.int32())}, schema=schema_to_pyarrow(schema)), delete_path) + + result = _scan_equality_delete_files(schema, data_path, [delete_path], [1], FileFormat.ORC) + assert result.column("id").to_pylist() == [1, 3] + + +def test_equality_delete_casts_promoted_key_types(tmp_path: Path) -> None: + current_schema = Schema(NestedField(1, "id", LongType(), required=True), schema_id=1) + old_schema = Schema(NestedField(1, "id", IntegerType(), required=True), schema_id=1) + data_path = tmp_path / "data.parquet" + delete_path = tmp_path / "delete.parquet" + pq.write_table( + pa.table({"id": pa.array([1, 2, 3], type=pa.int64())}, schema=schema_to_pyarrow(current_schema)), data_path + ) + pq.write_table( + pa.table({"id": pa.array([2], type=pa.int32())}, schema=schema_to_pyarrow(old_schema)), delete_path + ) + + result = _scan_equality_delete_files(current_schema, data_path, [delete_path], [1]) + assert result.column("id").to_pylist() == [1, 3] + + +def test_equality_delete_temporary_columns_do_not_collide() -> None: + helper_like_name = "__pyiceberg_equality_null_0" + data = pa.table({"id": [1, None, 3], helper_like_name: ["a", "b", "c"]}) + deletes = pa.table({"id": pa.array([None], type=pa.int64())}) + + result = _null_safe_left_anti_join(data, deletes, ["id"]) + assert result.to_pydict() == {"id": [1, 3], helper_like_name: ["a", "c"]} + + def test_pyarrow_wrap_fsspec(example_task: FileScanTask, table_schema_simple: Schema) -> None: metadata_location = "file://a/b/c.json" From 9190d565c36565efec6af59728b2d71908b3196a Mon Sep 17 00:00:00 2001 From: Nico Ritschel Date: Mon, 17 Aug 2026 20:16:50 -0700 Subject: [PATCH 3/6] Write atomic equality-delete upserts --- pyiceberg/io/pyarrow.py | 91 +++++++++++++++++++- pyiceberg/manifest.py | 16 +++- pyiceberg/table/__init__.py | 114 +++++++++++++++++++++++++ pyiceberg/table/update/snapshot.py | 54 ++++++++---- tests/table/test_upsert.py | 129 ++++++++++++++++++++++++++++- 5 files changed, 381 insertions(+), 23 deletions(-) diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py index f7f4d10898..9f58cadb0a 100644 --- a/pyiceberg/io/pyarrow.py +++ b/pyiceberg/io/pyarrow.py @@ -2926,10 +2926,17 @@ def write_file(io: FileIO, table_metadata: TableMetadata, tasks: Iterator[WriteT def write_data_file(task: WriteTask) -> DataFile: table_schema = table_metadata.schema() - if (sanitized_schema := sanitize_column_names(table_schema)) != table_schema: + if task.content == DataFileContent.EQUALITY_DELETES: + if not task.equality_ids: + raise ValueError("Equality delete write task requires equality IDs") + requested_schema = prune_columns(table_schema, set(task.equality_ids), select_full_types=False) + else: + requested_schema = table_schema + + if (sanitized_schema := sanitize_column_names(requested_schema)) != requested_schema: file_schema = sanitized_schema else: - file_schema = table_schema + file_schema = requested_schema downcast_ns_timestamp_to_us = Config().get_bool(DOWNCAST_NS_TIMESTAMP_TO_US_ON_WRITE) or False batches = [ @@ -2955,7 +2962,7 @@ def write_data_file(task: WriteTask) -> DataFile: statistics = writer.result() return DataFile.from_args( - content=DataFileContent.DATA, + content=task.content, file_path=file_path, file_format=file_format, partition=task.partition_key.partition if task.partition_key else Record(), @@ -2966,7 +2973,7 @@ def write_data_file(task: WriteTask) -> DataFile: sort_order_id=None, # Just copy these from the table for now spec_id=table_metadata.default_spec_id, - equality_ids=None, + equality_ids=list(task.equality_ids) if task.equality_ids is not None else None, key_metadata=None, **statistics.to_serialized_dict(), ) @@ -3241,6 +3248,82 @@ def _dataframe_to_data_files( ) +def _dataframe_to_equality_delete_files( + table_metadata: TableMetadata, + df: pa.Table, + equality_ids: list[int], + io: FileIO, + write_uuid: uuid.UUID | None = None, + counter: itertools.count[int] | None = None, +) -> Iterable[DataFile]: + """Write equality delete files from rows containing the table's current columns.""" + from pyiceberg.table import WriteTask + + if table_metadata.format_version < 2: + raise ValueError("Equality deletes require an Iceberg table using format version 2 or later") + if not equality_ids: + raise ValueError("At least one equality field ID is required") + + table_schema = table_metadata.schema() + missing_partition_sources = { + field.source_id for field in table_metadata.spec().fields if field.source_id not in equality_ids + } + if missing_partition_sources: + missing_names = sorted(table_schema.find_field(field_id).name for field_id in missing_partition_sources) + raise ValueError( + "Equality-delete upserts on partitioned tables require every partition source column in the equality keys; " + f"missing: {missing_names}" + ) + equality_names = [] + for field_id in equality_ids: + field_name = table_schema.find_column_name(field_id) + if field_name is None: + raise ValueError(f"Could not find equality field ID: {field_id}") + if "." in field_name: + raise NotImplementedError("Writing nested equality fields is not yet supported") + equality_names.append(field_name) + + counter = counter or itertools.count(0) + write_uuid = write_uuid or uuid.uuid4() + target_file_size = property_as_int( + properties=table_metadata.properties, + property_name=TableProperties.WRITE_TARGET_FILE_SIZE_BYTES, + default=TableProperties.WRITE_TARGET_FILE_SIZE_BYTES_DEFAULT, + ) + name_mapping = table_schema.name_mapping + + def write_partition(delete_rows: pa.Table, partition_key: PartitionKey | None) -> Iterable[DataFile]: + equality_rows = delete_rows.select(equality_names) + task_schema = pyarrow_to_schema( + equality_rows.schema, + name_mapping=name_mapping, + downcast_ns_timestamp_to_us=Config().get_bool(DOWNCAST_NS_TIMESTAMP_TO_US_ON_WRITE) or False, + format_version=table_metadata.format_version, + ) + return write_file( + io=io, + table_metadata=table_metadata, + tasks=( + WriteTask( + write_uuid=write_uuid, + task_id=next(counter), + record_batches=batches, + schema=task_schema, + partition_key=partition_key, + content=DataFileContent.EQUALITY_DELETES, + equality_ids=tuple(equality_ids), + ) + for batches in bin_pack_arrow_table(equality_rows, target_file_size) + ), + ) + + if table_metadata.spec().is_unpartitioned(): + yield from write_partition(df, None) + else: + for partition in _determine_partitions(spec=table_metadata.spec(), schema=table_schema, arrow_table=df): + yield from write_partition(partition.arrow_table_partition, partition.partition_key) + + @dataclass(frozen=True) class _TablePartition: partition_key: PartitionKey diff --git a/pyiceberg/manifest.py b/pyiceberg/manifest.py index 37dbd04b13..3fd8244fa4 100644 --- a/pyiceberg/manifest.py +++ b/pyiceberg/manifest.py @@ -1260,11 +1260,13 @@ def __init__( output_file: OutputFile, snapshot_id: int, avro_compression: AvroCompressionCodec, + content: ManifestContent = ManifestContent.DATA, ): super().__init__(spec, schema, output_file, snapshot_id, avro_compression) + self._content = content def content(self) -> ManifestContent: - return ManifestContent.DATA + return self._content @property def version(self) -> TableVersion: @@ -1274,10 +1276,15 @@ def version(self) -> TableVersion: def _meta(self) -> dict[str, str]: return { **super()._meta, - "content": "data", + "content": "data" if self._content == ManifestContent.DATA else "deletes", } def prepare_entry(self, entry: ManifestEntry) -> ManifestEntry: + entry_content = ManifestContent.DATA if entry.data_file.content == DataFileContent.DATA else ManifestContent.DELETES + if entry_content != self._content: + raise ValueError( + f"Cannot write {entry.data_file.content.name.lower()} file to a {self._content.name.lower()} manifest" + ) if entry.sequence_number is None: if entry.snapshot_id is not None and entry.snapshot_id != self._snapshot_id: raise ValueError(f"Found unassigned sequence number for an entry from snapshot: {entry.snapshot_id}") @@ -1293,11 +1300,14 @@ def write_manifest( output_file: OutputFile, snapshot_id: int, avro_compression: AvroCompressionCodec, + content: ManifestContent = ManifestContent.DATA, ) -> ManifestWriter: if format_version == 1: + if content != ManifestContent.DATA: + raise ValueError("Cannot write delete manifests in a v1 table") return ManifestWriterV1(spec, schema, output_file, snapshot_id, avro_compression) elif format_version == 2: - return ManifestWriterV2(spec, schema, output_file, snapshot_id, avro_compression) + return ManifestWriterV2(spec, schema, output_file, snapshot_id, avro_compression, content) else: raise ValueError(f"Cannot write manifest for table version: {format_version}") diff --git a/pyiceberg/table/__init__.py b/pyiceberg/table/__init__.py index 31e7689bc0..721209b3ed 100644 --- a/pyiceberg/table/__init__.py +++ b/pyiceberg/table/__init__.py @@ -989,6 +989,100 @@ def upsert( return UpsertResult(rows_updated=update_row_cnt, rows_inserted=insert_row_cnt) + def upsert_by_equality_delete( + self, + df: pa.Table, + join_cols: list[str] | None = None, + case_sensitive: bool = True, + branch: str | None = MAIN_BRANCH, + snapshot_properties: dict[str, str] = EMPTY_DICT, + ) -> None: + """Atomically commit source-key deletes and replacement rows in one snapshot. + + Snapshot visibility is atomic: readers see either the previous snapshot or + both the equality deletes and replacement data. Physical files are written + before the metadata commit, so a failed write or catalog commit can leave + unreferenced files for normal orphan-file maintenance. + + Concurrent commits are ordered by their Iceberg snapshot sequence. A later + equality-delete commit can remove matching rows from an earlier concurrent + commit; rows committed later survive. Partitioned tables require every + partition source column in ``join_cols`` and currently reject evolved + partition specs because one partition-scoped delete cannot safely cover + data written under a different spec. + """ + try: + import pyarrow as pa + except ModuleNotFoundError as e: + raise ModuleNotFoundError("For writes PyArrow needs to be installed") from e + + from pyiceberg.io.pyarrow import ( + _check_pyarrow_schema_compatible, + _dataframe_to_data_files, + _dataframe_to_equality_delete_files, + ) + from pyiceberg.table import upsert_util + + if not isinstance(df, pa.Table): + raise ValueError(f"Expected pa.Table, got: {df}") + + if join_cols is None: + join_cols = [] + for field_id in self.table_metadata.schema().identifier_field_ids: + column_name = self.table_metadata.schema().find_column_name(field_id) + if column_name is None: + raise ValueError(f"Field ID could not be found: {field_id}") + join_cols.append(column_name) + if not join_cols: + raise ValueError("Join columns could not be found, please set identifier-field-ids or pass in explicitly.") + if upsert_util.has_duplicate_rows(df, join_cols): + raise ValueError("Duplicate rows found in source dataset based on the key columns. No upsert executed") + + downcast_ns_timestamp_to_us = Config().get_bool(DOWNCAST_NS_TIMESTAMP_TO_US_ON_WRITE) or False + _check_pyarrow_schema_compatible( + self.table_metadata.schema(), + provided_schema=df.schema, + downcast_ns_timestamp_to_us=downcast_ns_timestamp_to_us, + format_version=self.table_metadata.format_version, + ) + if df.num_rows == 0: + return + if not self.table_metadata.spec().is_unpartitioned() and len(self.table_metadata.specs()) > 1: + raise NotImplementedError("Equality-delete upserts do not yet support evolved partition specs") + + equality_ids = [ + self.table_metadata.schema().find_field(column_name, case_sensitive=case_sensitive).field_id + for column_name in join_cols + ] + commit_uuid = uuid.uuid4() + counter = itertools.count(0) + equality_delete_files = list( + _dataframe_to_equality_delete_files( + table_metadata=self.table_metadata, + df=df, + equality_ids=equality_ids, + io=self._table.io, + write_uuid=commit_uuid, + counter=counter, + ) + ) + data_files = list( + _dataframe_to_data_files( + table_metadata=self.table_metadata, + df=df, + io=self._table.io, + write_uuid=commit_uuid, + counter=counter, + ) + ) + + with self.update_snapshot(snapshot_properties=snapshot_properties, branch=branch).row_delta() as row_delta: + row_delta.commit_uuid = commit_uuid + for equality_delete_file in equality_delete_files: + row_delta.append_data_file(equality_delete_file) + for data_file in data_files: + row_delta.append_data_file(data_file) + def _find_referenced_data_files(self, file_paths: list[str]) -> list[str]: """Return file_paths already referenced by data files in the current snapshot.""" snapshot = self.table_metadata.current_snapshot() @@ -1702,6 +1796,24 @@ def upsert( snapshot_properties=snapshot_properties, ) + def upsert_by_equality_delete( + self, + df: pa.Table, + join_cols: list[str] | None = None, + case_sensitive: bool = True, + branch: str | None = MAIN_BRANCH, + snapshot_properties: dict[str, str] = EMPTY_DICT, + ) -> None: + """Atomically commit equality deletes and replacement rows; see the transaction API for semantics.""" + with self.transaction() as tx: + tx.upsert_by_equality_delete( + df=df, + join_cols=join_cols, + case_sensitive=case_sensitive, + branch=branch, + snapshot_properties=snapshot_properties, + ) + def append( self, df: pa.Table | pa.RecordBatchReader, @@ -2907,6 +3019,8 @@ class WriteTask: record_batches: list[pa.RecordBatch] sort_order_id: int | None = None partition_key: PartitionKey | None = None + content: DataFileContent = DataFileContent.DATA + equality_ids: tuple[int, ...] | None = None def generate_data_file_filename(self, extension: str) -> str: # Mimics the behavior in the Java API: diff --git a/pyiceberg/table/update/snapshot.py b/pyiceberg/table/update/snapshot.py index 3c58f8ff44..157fd97d5a 100644 --- a/pyiceberg/table/update/snapshot.py +++ b/pyiceberg/table/update/snapshot.py @@ -228,11 +228,19 @@ def _process_manifests(self, manifests: list[ManifestFile]) -> list[ManifestFile def _manifests(self) -> list[ManifestFile]: def _write_added_manifest() -> list[ManifestFile]: - if self._added_data_files: + files_by_content: dict[ManifestContent, list[DataFile]] = defaultdict(list) + for data_file in self._added_data_files: + manifest_content = ( + ManifestContent.DATA if data_file.content == DataFileContent.DATA else ManifestContent.DELETES + ) + files_by_content[manifest_content].append(data_file) + + manifests = [] + for manifest_content, data_files in files_by_content.items(): with self.new_manifest_writer( - spec=self._transaction.table_metadata.spec(), + spec=self._transaction.table_metadata.spec(), content=manifest_content ) as writer: - for data_file in self._added_data_files: + for data_file in data_files: writer.add( ManifestEntry.from_args( status=ManifestEntryStatus.ADDED, @@ -242,9 +250,8 @@ def _write_added_manifest() -> list[ManifestFile]: data_file=data_file, ) ) - return [writer.to_manifest_file()] - else: - return [] + manifests.append(writer.to_manifest_file()) + return manifests def _write_delete_manifest() -> list[ManifestFile]: # Check if we need to mark the files as deleted @@ -391,7 +398,9 @@ def schema(self) -> Schema: def spec(self, spec_id: int) -> PartitionSpec: return self._transaction.table_metadata.specs()[spec_id] - def new_manifest_writer(self, spec: PartitionSpec) -> ManifestWriter: + def new_manifest_writer( + self, spec: PartitionSpec, content: ManifestContent = ManifestContent.DATA + ) -> ManifestWriter: return write_manifest( format_version=self._transaction.table_metadata.format_version, spec=spec, @@ -399,6 +408,7 @@ def new_manifest_writer(self, spec: PartitionSpec) -> ManifestWriter: output_file=self.new_manifest_output(), snapshot_id=self._snapshot_id, avro_compression=self._compression, + content=content, ) def new_manifest_output(self) -> OutputFile: @@ -888,6 +898,18 @@ def merge_append(self) -> _MergeAppendFiles: snapshot_properties=self._snapshot_properties, ) + def row_delta(self) -> _FastAppendFiles: + """Create a snapshot that atomically adds data and delete files.""" + return _FastAppendFiles( + operation=Operation.OVERWRITE + if self._transaction.table_metadata.snapshot_by_name(name=self._branch) is not None + else Operation.APPEND, + transaction=self._transaction, + io=self._io, + branch=self._branch, + snapshot_properties=self._snapshot_properties, + ) + def overwrite(self, commit_uuid: uuid.UUID | None = None) -> _OverwriteFiles: return _OverwriteFiles( commit_uuid=commit_uuid, @@ -924,14 +946,18 @@ def __init__( self._merge_enabled = merge_enabled self._snapshot_producer = snapshot_producer - def _group_by_spec(self, manifests: list[ManifestFile]) -> dict[int, list[ManifestFile]]: + def _group_by_spec_and_content( + self, manifests: list[ManifestFile] + ) -> dict[tuple[int, ManifestContent], list[ManifestFile]]: groups = defaultdict(list) for manifest in manifests: - groups[manifest.partition_spec_id].append(manifest) + groups[(manifest.partition_spec_id, manifest.content)].append(manifest) return groups def _create_manifest(self, spec_id: int, manifest_bin: list[ManifestFile]) -> ManifestFile: - with self._snapshot_producer.new_manifest_writer(spec=self._snapshot_producer.spec(spec_id)) as writer: + with self._snapshot_producer.new_manifest_writer( + spec=self._snapshot_producer.spec(spec_id), content=manifest_bin[0].content + ) as writer: for manifest in manifest_bin: for entry in self._snapshot_producer.fetch_manifest_entry(manifest=manifest, discard_deleted=False): if entry.status == ManifestEntryStatus.DELETED and entry.snapshot_id == self._snapshot_producer.snapshot_id: @@ -974,12 +1000,12 @@ def merge_manifests(self, manifests: list[ManifestFile]) -> list[ManifestFile]: if not self._merge_enabled or len(manifests) == 0: return manifests - first_manifest = manifests[0] - groups = self._group_by_spec(manifests) + groups = self._group_by_spec_and_content(manifests) merged_manifests = [] - for spec_id in reversed(groups.keys()): - merged_manifests.extend(self._merge_group(first_manifest, spec_id, groups[spec_id])) + for spec_id, content in reversed(groups.keys()): + group = groups[(spec_id, content)] + merged_manifests.extend(self._merge_group(group[0], spec_id, group)) return merged_manifests diff --git a/tests/table/test_upsert.py b/tests/table/test_upsert.py index 08f90c6600..86ca1c77bf 100644 --- a/tests/table/test_upsert.py +++ b/tests/table/test_upsert.py @@ -20,16 +20,20 @@ import pytest from datafusion import SessionContext from pyarrow import Table as pa_table +from pytest_mock import MockerFixture from pyiceberg.catalog import Catalog -from pyiceberg.exceptions import NoSuchTableError +from pyiceberg.exceptions import CommitFailedException, NoSuchTableError from pyiceberg.expressions import AlwaysTrue, And, EqualTo, Reference from pyiceberg.expressions.literals import LongLiteral from pyiceberg.io.pyarrow import schema_to_pyarrow +from pyiceberg.manifest import DataFileContent, ManifestContent +from pyiceberg.partitioning import PartitionField, PartitionSpec from pyiceberg.schema import Schema -from pyiceberg.table import Table, UpsertResult +from pyiceberg.table import Table, TableProperties, UpsertResult from pyiceberg.table.snapshots import Operation from pyiceberg.table.upsert_util import create_match_filter +from pyiceberg.transforms import IdentityTransform from pyiceberg.types import IntegerType, NestedField, StringType, StructType from tests.catalog.test_base import InMemoryCatalog @@ -119,6 +123,127 @@ def assert_upsert_result(res: UpsertResult, expected_updated: int, expected_inse assert res.rows_inserted == expected_inserted, f"rows inserted should be {expected_inserted}, but got {res.rows_inserted}" +def test_composite_equality_delete_upsert_15k_rows(catalog: InMemoryCatalog) -> None: + initial_count = 20_000 + source_start = 10_000 + source_end = 25_000 + initial = pa.table( + { + "account_id": pa.array([row // 10 for row in range(initial_count)], type=pa.int64()), + "item_id": pa.array([row % 10 for row in range(initial_count)], type=pa.int64()), + "value": pa.array(range(initial_count), type=pa.int64()), + } + ) + source = pa.table( + { + "account_id": pa.array([row // 10 for row in range(source_start, source_end)], type=pa.int64()), + "item_id": pa.array([row % 10 for row in range(source_start, source_end)], type=pa.int64()), + "value": pa.array([-row for row in range(source_start, source_end)], type=pa.int64()), + } + ) + + table = catalog.create_table("default.composite_equality_upsert", initial.schema) + table.append(initial) + table.upsert_by_equality_delete(source, join_cols=["account_id", "item_id"]) + + result = table.scan().to_arrow().sort_by([("account_id", "ascending"), ("item_id", "ascending")]) + assert result.num_rows == source_end + values = result.column("value").to_pylist() + assert values[:source_start] == list(range(source_start)) + assert values[source_start:] == [-row for row in range(source_start, source_end)] + + snapshot = table.current_snapshot() + assert snapshot is not None + manifests = snapshot.manifests(table.io) + assert {manifest.content for manifest in manifests} == {ManifestContent.DATA, ManifestContent.DELETES} + delete_entries = [ + entry + for manifest in manifests + if manifest.content == ManifestContent.DELETES + for entry in manifest.fetch_manifest_entry(table.io, discard_deleted=True) + ] + assert len(delete_entries) == 1 + assert delete_entries[0].data_file.content == DataFileContent.EQUALITY_DELETES + assert delete_entries[0].data_file.equality_ids == [1, 2] + + +def test_equality_delete_upsert_is_metadata_atomic_on_commit_failure( + catalog: InMemoryCatalog, mocker: MockerFixture +) -> None: + initial = pa.table({"id": pa.array([1, 2], type=pa.int64()), "value": ["one", "two"]}) + source = pa.table({"id": pa.array([2, 3], type=pa.int64()), "value": ["updated", "three"]}) + table = catalog.create_table( + "default.equality_atomic_failure", + initial.schema, + properties={TableProperties.COMMIT_NUM_RETRIES: "0"}, + ) + table.append(initial) + original_snapshot = table.current_snapshot() + assert original_snapshot is not None + + mocker.patch.object(catalog, "commit_table", side_effect=CommitFailedException("forced failure")) + with pytest.raises(CommitFailedException, match="forced failure"): + table.upsert_by_equality_delete(source, join_cols=["id"]) + + table.refresh() + assert table.current_snapshot().snapshot_id == original_snapshot.snapshot_id # type: ignore[union-attr] + assert table.scan().to_arrow().sort_by("id").to_pydict() == {"id": [1, 2], "value": ["one", "two"]} + + +def test_equality_delete_upsert_survives_manifest_merging(catalog: InMemoryCatalog) -> None: + initial = pa.table({"id": pa.array([1, 2], type=pa.int64()), "value": ["one", "two"]}) + table = catalog.create_table( + "default.equality_manifest_merge", + initial.schema, + properties={ + TableProperties.MANIFEST_MERGE_ENABLED: "true", + TableProperties.MANIFEST_MIN_MERGE_COUNT: "1", + TableProperties.MANIFEST_TARGET_SIZE_BYTES: str(64 * 1024 * 1024), + }, + ) + table.append(initial) + table.upsert_by_equality_delete( + pa.table({"id": pa.array([2, 3], type=pa.int64()), "value": ["updated", "three"]}), join_cols=["id"] + ) + table.append(pa.table({"id": pa.array([4], type=pa.int64()), "value": ["four"]})) + + snapshot = table.current_snapshot() + assert snapshot is not None + assert {manifest.content for manifest in snapshot.manifests(table.io)} == { + ManifestContent.DATA, + ManifestContent.DELETES, + } + assert table.scan().to_arrow().sort_by("id").to_pydict() == { + "id": [1, 2, 3, 4], + "value": ["one", "updated", "three", "four"], + } + + +def test_empty_equality_delete_upsert_does_not_create_snapshot(catalog: InMemoryCatalog) -> None: + schema = pa.schema([pa.field("id", pa.int64()), pa.field("value", pa.string())]) + table = catalog.create_table("default.empty_equality_upsert", schema) + table.upsert_by_equality_delete(schema.empty_table(), join_cols=["id"]) + assert table.current_snapshot() is None + + +def test_partitioned_equality_delete_upsert_requires_partition_source_in_key(catalog: InMemoryCatalog) -> None: + schema = Schema( + NestedField(1, "id", IntegerType(), required=True), + NestedField(2, "region", StringType(), required=True), + NestedField(3, "value", StringType(), required=False), + ) + spec = PartitionSpec(PartitionField(2, 1000, IdentityTransform(), "region")) + table = catalog.create_table("default.partitioned_equality_upsert", schema, partition_spec=spec) + source = pa.table( + {"id": pa.array([1], type=pa.int32()), "region": ["west"], "value": ["one"]}, + schema=schema_to_pyarrow(schema, include_field_ids=False), + ) + + with pytest.raises(ValueError, match="require every partition source column"): + table.upsert_by_equality_delete(source, join_cols=["id"]) + assert table.current_snapshot() is None + + @pytest.mark.parametrize( ( "join_cols, src_start_row, src_end_row, target_start_row, target_end_row, " From d9f2cbd330c35366612c9cde1ef3c6c2f6027b5e Mon Sep 17 00:00:00 2001 From: Nico Ritschel Date: Mon, 17 Aug 2026 20:16:51 -0700 Subject: [PATCH 4/6] Encode equality field IDs as Avro integers --- pyiceberg/manifest.py | 4 +- .../test_writes/test_equality_deletes.py | 80 +++++++++++++++++++ tests/utils/test_manifest.py | 10 ++- 3 files changed, 91 insertions(+), 3 deletions(-) create mode 100644 tests/integration/test_writes/test_equality_deletes.py diff --git a/pyiceberg/manifest.py b/pyiceberg/manifest.py index 3fd8244fa4..726803b5c5 100644 --- a/pyiceberg/manifest.py +++ b/pyiceberg/manifest.py @@ -295,7 +295,7 @@ def __repr__(self) -> str: NestedField( field_id=135, name="equality_ids", - field_type=ListType(element_id=136, element_type=LongType(), element_required=True), + field_type=ListType(element_id=136, element_type=IntegerType(), element_required=True), required=False, doc="Field ids used to determine row equality in equality delete files.", ), @@ -390,7 +390,7 @@ def __repr__(self) -> str: NestedField( field_id=135, name="equality_ids", - field_type=ListType(element_id=136, element_type=LongType(), element_required=True), + field_type=ListType(element_id=136, element_type=IntegerType(), element_required=True), required=False, doc="Field ids used to determine row equality in equality delete files.", ), diff --git a/tests/integration/test_writes/test_equality_deletes.py b/tests/integration/test_writes/test_equality_deletes.py new file mode 100644 index 0000000000..8ed5e90098 --- /dev/null +++ b/tests/integration/test_writes/test_equality_deletes.py @@ -0,0 +1,80 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import pyarrow as pa +import pytest +from pyspark.sql import SparkSession + +from pyiceberg.catalog import Catalog +from pyiceberg.manifest import ManifestContent + + +@pytest.mark.integration +def test_large_composite_equality_upsert_is_visible_to_spark( + session_catalog: Catalog, spark: SparkSession +) -> None: + identifier = "default.large_composite_equality_upsert" + if session_catalog.table_exists(identifier): + session_catalog.drop_table(identifier) + + initial_count = 200_000 + source_start = 100_000 + source_end = 300_000 + initial = pa.table( + { + "account_id": pa.array((row // 10 for row in range(initial_count)), type=pa.int64()), + "item_id": pa.array((row % 10 for row in range(initial_count)), type=pa.int64()), + "value": pa.array(range(initial_count), type=pa.int64()), + } + ) + source = pa.table( + { + "account_id": pa.array((row // 10 for row in range(source_start, source_end)), type=pa.int64()), + "item_id": pa.array((row % 10 for row in range(source_start, source_end)), type=pa.int64()), + "value": pa.array((-row for row in range(source_start, source_end)), type=pa.int64()), + } + ) + + try: + table = session_catalog.create_table(identifier, initial.schema, properties={"format-version": "2"}) + table.append(initial) + table.upsert_by_equality_delete(source, join_cols=["account_id", "item_id"]) + + result = table.scan().to_arrow() + assert result.num_rows == source_end + result_by_key = result.sort_by([("account_id", "ascending"), ("item_id", "ascending")]) + assert result_by_key.column("value")[source_start - 1].as_py() == source_start - 1 + assert result_by_key.column("value")[source_start].as_py() == -source_start + assert result_by_key.column("value")[-1].as_py() == -(source_end - 1) + + snapshot = table.current_snapshot() + assert snapshot is not None + assert {manifest.content for manifest in snapshot.manifests(table.io)} == { + ManifestContent.DATA, + ManifestContent.DELETES, + } + + spark_table = spark.table(identifier) + assert spark_table.count() == source_end + sample = spark_table.where( + (spark_table.account_id == source_start // 10) & (spark_table.item_id == source_start % 10) + ).collect() + assert len(sample) == 1 + assert sample[0].value == -source_start + finally: + if session_catalog.table_exists(identifier): + session_catalog.drop_table(identifier) diff --git a/tests/utils/test_manifest.py b/tests/utils/test_manifest.py index 0535ec01ed..65f908a6ff 100644 --- a/tests/utils/test_manifest.py +++ b/tests/utils/test_manifest.py @@ -28,6 +28,7 @@ from pyiceberg.io import load_file_io from pyiceberg.io.pyarrow import PyArrowFileIO from pyiceberg.manifest import ( + DATA_FILE_TYPE, DataFile, DataFileContent, FileFormat, @@ -47,7 +48,7 @@ from pyiceberg.schema import Schema from pyiceberg.table.snapshots import Operation, Snapshot, Summary from pyiceberg.typedef import Record, TableVersion -from pyiceberg.types import IntegerType, NestedField +from pyiceberg.types import IntegerType, ListType, NestedField @pytest.fixture(autouse=True) @@ -55,6 +56,13 @@ def reset_global_manifests_cache() -> None: clear_manifest_cache() +@pytest.mark.parametrize("format_version", [2, 3]) +def test_equality_ids_use_iceberg_int_wire_type(format_version: int) -> None: + equality_ids = DATA_FILE_TYPE[format_version].field_by_name("equality_ids") + assert isinstance(equality_ids.field_type, ListType) + assert isinstance(equality_ids.field_type.element_type, IntegerType) + + def _verify_metadata_with_fastavro(avro_file: str, expected_metadata: dict[str, str]) -> None: with open(avro_file, "rb") as f: reader = fastavro.reader(f) From 4e7b5ea40585cf103ce30b4413530f624fc63979 Mon Sep 17 00:00:00 2001 From: Nico Ritschel Date: Mon, 17 Aug 2026 20:16:52 -0700 Subject: [PATCH 5/6] Normalize the equality-delete implementation --- pyiceberg/io/pyarrow.py | 13 +++------ pyiceberg/table/delete_file_index.py | 7 ++--- pyiceberg/table/update/snapshot.py | 16 +++-------- .../test_writes/test_equality_deletes.py | 4 +-- tests/io/test_pyarrow.py | 28 +++++-------------- tests/table/test_delete_file_index.py | 28 +++++-------------- tests/table/test_upsert.py | 8 +++--- tests/utils/test_manifest.py | 1 + 8 files changed, 31 insertions(+), 74 deletions(-) diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py index 9f58cadb0a..399a4b3431 100644 --- a/pyiceberg/io/pyarrow.py +++ b/pyiceberg/io/pyarrow.py @@ -1175,9 +1175,7 @@ def _read_equality_deletes(io: FileIO, delete_file: DataFile) -> pa.Table: columns = projected_columns if len(projected_columns) == len(equality_ids) else None # Delete files are already read concurrently by PyIceberg's executor; # disable nested Arrow threading to avoid oversubscription per file. - return ds.Scanner.from_fragment( - fragment=fragment, schema=physical_schema, columns=columns, use_threads=False - ).to_table() + return ds.Scanner.from_fragment(fragment=fragment, schema=physical_schema, columns=columns, use_threads=False).to_table() def _column_name_for_field_id(table: pa.Table, field_id: int, schema: Schema) -> str | None: @@ -1291,9 +1289,7 @@ def temporary_name(kind: str) -> str: join_columns.append(nan_key) temporary_columns.append(nan_key) - joined = data_join.join( - delete_join.select(join_columns), keys=join_columns, join_type="left anti", use_threads=False - ) + joined = data_join.join(delete_join.select(join_columns), keys=join_columns, join_type="left anti", use_threads=False) return joined.drop(temporary_columns) @@ -3265,9 +3261,7 @@ def _dataframe_to_equality_delete_files( raise ValueError("At least one equality field ID is required") table_schema = table_metadata.schema() - missing_partition_sources = { - field.source_id for field in table_metadata.spec().fields if field.source_id not in equality_ids - } + missing_partition_sources = {field.source_id for field in table_metadata.spec().fields if field.source_id not in equality_ids} if missing_partition_sources: missing_names = sorted(table_schema.find_field(field_id).name for field_id in missing_partition_sources) raise ValueError( @@ -3290,6 +3284,7 @@ def _dataframe_to_equality_delete_files( property_name=TableProperties.WRITE_TARGET_FILE_SIZE_BYTES, default=TableProperties.WRITE_TARGET_FILE_SIZE_BYTES_DEFAULT, ) + assert target_file_size is not None name_mapping = table_schema.name_mapping def write_partition(delete_rows: pa.Table, partition_key: PartitionKey | None) -> Iterable[DataFile]: diff --git a/pyiceberg/table/delete_file_index.py b/pyiceberg/table/delete_file_index.py index 88a8695c14..e9b43aeabf 100644 --- a/pyiceberg/table/delete_file_index.py +++ b/pyiceberg/table/delete_file_index.py @@ -153,10 +153,9 @@ def _equality_delete_applies_to_data_file(delete_file: DataFile, data_file: Data and field_id in data_upper ): field_type = field.field_type - if ( - from_bytes(field_type, delete_upper[field_id]) < from_bytes(field_type, data_lower[field_id]) - or from_bytes(field_type, delete_lower[field_id]) > from_bytes(field_type, data_upper[field_id]) - ): + if from_bytes(field_type, delete_upper[field_id]) < from_bytes(field_type, data_lower[field_id]) or from_bytes( + field_type, delete_lower[field_id] + ) > from_bytes(field_type, data_upper[field_id]): return False return True diff --git a/pyiceberg/table/update/snapshot.py b/pyiceberg/table/update/snapshot.py index 157fd97d5a..6359163767 100644 --- a/pyiceberg/table/update/snapshot.py +++ b/pyiceberg/table/update/snapshot.py @@ -230,16 +230,12 @@ def _manifests(self) -> list[ManifestFile]: def _write_added_manifest() -> list[ManifestFile]: files_by_content: dict[ManifestContent, list[DataFile]] = defaultdict(list) for data_file in self._added_data_files: - manifest_content = ( - ManifestContent.DATA if data_file.content == DataFileContent.DATA else ManifestContent.DELETES - ) + manifest_content = ManifestContent.DATA if data_file.content == DataFileContent.DATA else ManifestContent.DELETES files_by_content[manifest_content].append(data_file) manifests = [] for manifest_content, data_files in files_by_content.items(): - with self.new_manifest_writer( - spec=self._transaction.table_metadata.spec(), content=manifest_content - ) as writer: + with self.new_manifest_writer(spec=self._transaction.table_metadata.spec(), content=manifest_content) as writer: for data_file in data_files: writer.add( ManifestEntry.from_args( @@ -398,9 +394,7 @@ def schema(self) -> Schema: def spec(self, spec_id: int) -> PartitionSpec: return self._transaction.table_metadata.specs()[spec_id] - def new_manifest_writer( - self, spec: PartitionSpec, content: ManifestContent = ManifestContent.DATA - ) -> ManifestWriter: + def new_manifest_writer(self, spec: PartitionSpec, content: ManifestContent = ManifestContent.DATA) -> ManifestWriter: return write_manifest( format_version=self._transaction.table_metadata.format_version, spec=spec, @@ -946,9 +940,7 @@ def __init__( self._merge_enabled = merge_enabled self._snapshot_producer = snapshot_producer - def _group_by_spec_and_content( - self, manifests: list[ManifestFile] - ) -> dict[tuple[int, ManifestContent], list[ManifestFile]]: + def _group_by_spec_and_content(self, manifests: list[ManifestFile]) -> dict[tuple[int, ManifestContent], list[ManifestFile]]: groups = defaultdict(list) for manifest in manifests: groups[(manifest.partition_spec_id, manifest.content)].append(manifest) diff --git a/tests/integration/test_writes/test_equality_deletes.py b/tests/integration/test_writes/test_equality_deletes.py index 8ed5e90098..aa696816fb 100644 --- a/tests/integration/test_writes/test_equality_deletes.py +++ b/tests/integration/test_writes/test_equality_deletes.py @@ -24,9 +24,7 @@ @pytest.mark.integration -def test_large_composite_equality_upsert_is_visible_to_spark( - session_catalog: Catalog, spark: SparkSession -) -> None: +def test_large_composite_equality_upsert_is_visible_to_spark(session_catalog: Catalog, spark: SparkSession) -> None: identifier = "default.large_composite_equality_upsert" if session_catalog.table_exists(identifier): session_catalog.drop_table(identifier) diff --git a/tests/io/test_pyarrow.py b/tests/io/test_pyarrow.py index ea67c72e6b..9c86dc6eeb 100644 --- a/tests/io/test_pyarrow.py +++ b/tests/io/test_pyarrow.py @@ -2050,9 +2050,7 @@ def test_equality_delete_treats_uuid_nulls_as_equal(tmp_path: Path) -> None: data_path = tmp_path / "data.parquet" delete_path = tmp_path / "delete.parquet" kept = uuid.UUID("00000000-0000-0000-0000-000000000001") - pq.write_table( - pa.table({"id": pa.array([kept, None], type=pa.uuid())}, schema=schema_to_pyarrow(schema)), data_path - ) + pq.write_table(pa.table({"id": pa.array([kept, None], type=pa.uuid())}, schema=schema_to_pyarrow(schema)), data_path) pq.write_table(pa.table({"id": pa.array([None], type=pa.uuid())}), delete_path) result = _scan_equality_delete_files(schema, data_path, [delete_path], [1]) @@ -2065,9 +2063,7 @@ def test_equality_delete_aligns_renamed_fields_by_id(tmp_path: Path) -> None: data_path = tmp_path / "data.parquet" current_delete_path = tmp_path / "current-delete.parquet" renamed_delete_path = tmp_path / "renamed-delete.parquet" - pq.write_table( - pa.table({"id": pa.array([1, 2, 3], type=pa.int32())}, schema=schema_to_pyarrow(current_schema)), data_path - ) + pq.write_table(pa.table({"id": pa.array([1, 2, 3], type=pa.int32())}, schema=schema_to_pyarrow(current_schema)), data_path) pq.write_table( pa.table({"id": pa.array([2], type=pa.int32())}, schema=schema_to_pyarrow(current_schema)), current_delete_path, @@ -2077,9 +2073,7 @@ def test_equality_delete_aligns_renamed_fields_by_id(tmp_path: Path) -> None: renamed_delete_path, ) - result = _scan_equality_delete_files( - current_schema, data_path, [current_delete_path, renamed_delete_path], [1] - ) + result = _scan_equality_delete_files(current_schema, data_path, [current_delete_path, renamed_delete_path], [1]) assert result.column("id").to_pylist() == [1] @@ -2092,9 +2086,7 @@ def test_equality_delete_dropped_field_does_not_match_a_partial_key(tmp_path: Pa ) data_path = tmp_path / "data.parquet" delete_path = tmp_path / "delete.parquet" - pq.write_table( - pa.table({"id": pa.array([1], type=pa.int32())}, schema=schema_to_pyarrow(current_schema)), data_path - ) + pq.write_table(pa.table({"id": pa.array([1], type=pa.int32())}, schema=schema_to_pyarrow(current_schema)), data_path) pq.write_table( pa.table( {"id": pa.array([1], type=pa.int32()), "dropped": ["not-null"]}, @@ -2111,9 +2103,7 @@ def test_orc_equality_delete(tmp_path: Path) -> None: schema = Schema(NestedField(1, "id", IntegerType(), required=True), schema_id=1) data_path = tmp_path / "data.orc" delete_path = tmp_path / "delete.orc" - orc.write_table( - pa.table({"id": pa.array([1, 2, 3], type=pa.int32())}, schema=schema_to_pyarrow(schema)), data_path - ) + orc.write_table(pa.table({"id": pa.array([1, 2, 3], type=pa.int32())}, schema=schema_to_pyarrow(schema)), data_path) orc.write_table(pa.table({"id": pa.array([2], type=pa.int32())}, schema=schema_to_pyarrow(schema)), delete_path) result = _scan_equality_delete_files(schema, data_path, [delete_path], [1], FileFormat.ORC) @@ -2125,12 +2115,8 @@ def test_equality_delete_casts_promoted_key_types(tmp_path: Path) -> None: old_schema = Schema(NestedField(1, "id", IntegerType(), required=True), schema_id=1) data_path = tmp_path / "data.parquet" delete_path = tmp_path / "delete.parquet" - pq.write_table( - pa.table({"id": pa.array([1, 2, 3], type=pa.int64())}, schema=schema_to_pyarrow(current_schema)), data_path - ) - pq.write_table( - pa.table({"id": pa.array([2], type=pa.int32())}, schema=schema_to_pyarrow(old_schema)), delete_path - ) + pq.write_table(pa.table({"id": pa.array([1, 2, 3], type=pa.int64())}, schema=schema_to_pyarrow(current_schema)), data_path) + pq.write_table(pa.table({"id": pa.array([2], type=pa.int32())}, schema=schema_to_pyarrow(old_schema)), delete_path) result = _scan_equality_delete_files(current_schema, data_path, [delete_path], [1]) assert result.column("id").to_pylist() == [1, 3] diff --git a/tests/table/test_delete_file_index.py b/tests/table/test_delete_file_index.py index 7ee3c34552..7d48532934 100644 --- a/tests/table/test_delete_file_index.py +++ b/tests/table/test_delete_file_index.py @@ -292,15 +292,9 @@ def test_equality_delete_metrics_filtering() -> None: ) index.add_delete_file(equality_delete) - before = _create_data_file( - lower_bounds={1: to_bytes(IntegerType(), 0)}, upper_bounds={1: to_bytes(IntegerType(), 5)} - ) - overlap = _create_data_file( - lower_bounds={1: to_bytes(IntegerType(), 15)}, upper_bounds={1: to_bytes(IntegerType(), 25)} - ) - after = _create_data_file( - lower_bounds={1: to_bytes(IntegerType(), 25)}, upper_bounds={1: to_bytes(IntegerType(), 30)} - ) + before = _create_data_file(lower_bounds={1: to_bytes(IntegerType(), 0)}, upper_bounds={1: to_bytes(IntegerType(), 5)}) + overlap = _create_data_file(lower_bounds={1: to_bytes(IntegerType(), 15)}, upper_bounds={1: to_bytes(IntegerType(), 25)}) + after = _create_data_file(lower_bounds={1: to_bytes(IntegerType(), 25)}, upper_bounds={1: to_bytes(IntegerType(), 30)}) assert index.for_data_file(1, before) == set() assert index.for_data_file(1, overlap) == {equality_delete.data_file} assert index.for_data_file(1, after) == set() @@ -310,9 +304,7 @@ def test_equality_delete_metrics_filtering() -> None: ("delete_nulls", "data_nulls"), [((10, 10), (0, 100)), ((0, 10), (100, 100))], ) -def test_equality_delete_prunes_disjoint_null_populations( - delete_nulls: tuple[int, int], data_nulls: tuple[int, int] -) -> None: +def test_equality_delete_prunes_disjoint_null_populations(delete_nulls: tuple[int, int], data_nulls: tuple[int, int]) -> None: index = DeleteFileIndex(Schema(NestedField(1, "id", IntegerType(), required=False))) equality_delete = _create_equality_delete( sequence_number=10, @@ -334,12 +326,8 @@ def test_equality_delete_metrics_after_int_to_long_promotion() -> None: upper_bounds={1: to_bytes(IntegerType(), 20)}, ) index.add_delete_file(equality_delete) - before = _create_data_file( - lower_bounds={1: to_bytes(IntegerType(), 0)}, upper_bounds={1: to_bytes(IntegerType(), 5)} - ) - overlap = _create_data_file( - lower_bounds={1: to_bytes(IntegerType(), 15)}, upper_bounds={1: to_bytes(IntegerType(), 25)} - ) + before = _create_data_file(lower_bounds={1: to_bytes(IntegerType(), 0)}, upper_bounds={1: to_bytes(IntegerType(), 5)}) + overlap = _create_data_file(lower_bounds={1: to_bytes(IntegerType(), 15)}, upper_bounds={1: to_bytes(IntegerType(), 25)}) assert index.for_data_file(1, before) == set() assert index.for_data_file(1, overlap) == {equality_delete.data_file} @@ -353,9 +341,7 @@ def test_equality_delete_dropped_field_is_not_pruned() -> None: upper_bounds={1: to_bytes(IntegerType(), 20)}, ) index.add_delete_file(equality_delete) - data_file = _create_data_file( - lower_bounds={1: to_bytes(IntegerType(), 15)}, upper_bounds={1: to_bytes(IntegerType(), 25)} - ) + data_file = _create_data_file(lower_bounds={1: to_bytes(IntegerType(), 15)}, upper_bounds={1: to_bytes(IntegerType(), 25)}) assert index.for_data_file(1, data_file) == {equality_delete.data_file} diff --git a/tests/table/test_upsert.py b/tests/table/test_upsert.py index 86ca1c77bf..624f15678a 100644 --- a/tests/table/test_upsert.py +++ b/tests/table/test_upsert.py @@ -167,9 +167,7 @@ def test_composite_equality_delete_upsert_15k_rows(catalog: InMemoryCatalog) -> assert delete_entries[0].data_file.equality_ids == [1, 2] -def test_equality_delete_upsert_is_metadata_atomic_on_commit_failure( - catalog: InMemoryCatalog, mocker: MockerFixture -) -> None: +def test_equality_delete_upsert_is_metadata_atomic_on_commit_failure(catalog: InMemoryCatalog, mocker: MockerFixture) -> None: initial = pa.table({"id": pa.array([1, 2], type=pa.int64()), "value": ["one", "two"]}) source = pa.table({"id": pa.array([2, 3], type=pa.int64()), "value": ["updated", "three"]}) table = catalog.create_table( @@ -186,7 +184,9 @@ def test_equality_delete_upsert_is_metadata_atomic_on_commit_failure( table.upsert_by_equality_delete(source, join_cols=["id"]) table.refresh() - assert table.current_snapshot().snapshot_id == original_snapshot.snapshot_id # type: ignore[union-attr] + current_snapshot = table.current_snapshot() + assert current_snapshot is not None + assert current_snapshot.snapshot_id == original_snapshot.snapshot_id assert table.scan().to_arrow().sort_by("id").to_pydict() == {"id": [1, 2], "value": ["one", "two"]} diff --git a/tests/utils/test_manifest.py b/tests/utils/test_manifest.py index 65f908a6ff..2a81f64d91 100644 --- a/tests/utils/test_manifest.py +++ b/tests/utils/test_manifest.py @@ -59,6 +59,7 @@ def reset_global_manifests_cache() -> None: @pytest.mark.parametrize("format_version", [2, 3]) def test_equality_ids_use_iceberg_int_wire_type(format_version: int) -> None: equality_ids = DATA_FILE_TYPE[format_version].field_by_name("equality_ids") + assert equality_ids is not None assert isinstance(equality_ids.field_type, ListType) assert isinstance(equality_ids.field_type.element_type, IntegerType) From 968fb9677987d7a308ce4aeb70e8c441fe8f4e44 Mon Sep 17 00:00:00 2001 From: Nico Ritschel Date: Mon, 17 Aug 2026 20:16:53 -0700 Subject: [PATCH 6/6] Support delete-only equality keys --- pyiceberg/table/__init__.py | 46 +++++++++++++++++++++++++++---------- tests/table/test_upsert.py | 36 +++++++++++++++++++++++++++++ 2 files changed, 70 insertions(+), 12 deletions(-) diff --git a/pyiceberg/table/__init__.py b/pyiceberg/table/__init__.py index 721209b3ed..bf9470e2b4 100644 --- a/pyiceberg/table/__init__.py +++ b/pyiceberg/table/__init__.py @@ -996,9 +996,14 @@ def upsert_by_equality_delete( case_sensitive: bool = True, branch: str | None = MAIN_BRANCH, snapshot_properties: dict[str, str] = EMPTY_DICT, + delete_df: pa.Table | None = None, ) -> None: """Atomically commit source-key deletes and replacement rows in one snapshot. + ``df`` contains replacement rows: their keys are deleted and the rows are + appended. ``delete_df`` optionally contains delete-only rows; only its key + columns are written to the equality-delete file. + Snapshot visibility is atomic: readers see either the previous snapshot or both the equality deletes and replacement data. Physical files are written before the metadata commit, so a failed write or catalog commit can leave @@ -1025,6 +1030,8 @@ def upsert_by_equality_delete( if not isinstance(df, pa.Table): raise ValueError(f"Expected pa.Table, got: {df}") + if delete_df is not None and not isinstance(delete_df, pa.Table): + raise ValueError(f"Expected delete_df to be pa.Table, got: {delete_df}") if join_cols is None: join_cols = [] @@ -1035,9 +1042,6 @@ def upsert_by_equality_delete( join_cols.append(column_name) if not join_cols: raise ValueError("Join columns could not be found, please set identifier-field-ids or pass in explicitly.") - if upsert_util.has_duplicate_rows(df, join_cols): - raise ValueError("Duplicate rows found in source dataset based on the key columns. No upsert executed") - downcast_ns_timestamp_to_us = Config().get_bool(DOWNCAST_NS_TIMESTAMP_TO_US_ON_WRITE) or False _check_pyarrow_schema_compatible( self.table_metadata.schema(), @@ -1045,7 +1049,19 @@ def upsert_by_equality_delete( downcast_ns_timestamp_to_us=downcast_ns_timestamp_to_us, format_version=self.table_metadata.format_version, ) - if df.num_rows == 0: + delete_key_tables = [df.select(join_cols)] + if delete_df is not None: + _check_pyarrow_schema_compatible( + self.table_metadata.schema().select(*join_cols, case_sensitive=case_sensitive), + provided_schema=delete_df.select(join_cols).schema, + downcast_ns_timestamp_to_us=downcast_ns_timestamp_to_us, + format_version=self.table_metadata.format_version, + ) + delete_key_tables.append(delete_df.select(join_cols)) + delete_rows = pa.concat_tables(delete_key_tables) + if upsert_util.has_duplicate_rows(delete_rows, join_cols): + raise ValueError("Duplicate rows found in source datasets based on the key columns. No upsert executed") + if delete_rows.num_rows == 0: return if not self.table_metadata.spec().is_unpartitioned() and len(self.table_metadata.specs()) > 1: raise NotImplementedError("Equality-delete upserts do not yet support evolved partition specs") @@ -1059,21 +1075,25 @@ def upsert_by_equality_delete( equality_delete_files = list( _dataframe_to_equality_delete_files( table_metadata=self.table_metadata, - df=df, + df=delete_rows, equality_ids=equality_ids, io=self._table.io, write_uuid=commit_uuid, counter=counter, ) ) - data_files = list( - _dataframe_to_data_files( - table_metadata=self.table_metadata, - df=df, - io=self._table.io, - write_uuid=commit_uuid, - counter=counter, + data_files = ( + list( + _dataframe_to_data_files( + table_metadata=self.table_metadata, + df=df, + io=self._table.io, + write_uuid=commit_uuid, + counter=counter, + ) ) + if df.num_rows > 0 + else [] ) with self.update_snapshot(snapshot_properties=snapshot_properties, branch=branch).row_delta() as row_delta: @@ -1803,12 +1823,14 @@ def upsert_by_equality_delete( case_sensitive: bool = True, branch: str | None = MAIN_BRANCH, snapshot_properties: dict[str, str] = EMPTY_DICT, + delete_df: pa.Table | None = None, ) -> None: """Atomically commit equality deletes and replacement rows; see the transaction API for semantics.""" with self.transaction() as tx: tx.upsert_by_equality_delete( df=df, join_cols=join_cols, + delete_df=delete_df, case_sensitive=case_sensitive, branch=branch, snapshot_properties=snapshot_properties, diff --git a/tests/table/test_upsert.py b/tests/table/test_upsert.py index 624f15678a..3306391b5c 100644 --- a/tests/table/test_upsert.py +++ b/tests/table/test_upsert.py @@ -226,6 +226,42 @@ def test_empty_equality_delete_upsert_does_not_create_snapshot(catalog: InMemory assert table.current_snapshot() is None +def test_equality_delete_upsert_supports_delete_only_rows(catalog: InMemoryCatalog) -> None: + initial = pa.table({"id": pa.array([1, 2, 3], type=pa.int64()), "value": ["one", "two", "three"]}) + table = catalog.create_table("default.equality_delete_only", initial.schema) + table.append(initial) + + table.upsert_by_equality_delete( + initial.schema.empty_table(), + join_cols=["id"], + delete_df=pa.table({"id": pa.array([2], type=pa.int64())}), + ) + + assert table.scan().to_arrow().sort_by("id").to_pydict() == {"id": [1, 3], "value": ["one", "three"]} + + +def test_equality_delete_upsert_commits_replacements_and_delete_only_rows_together(catalog: InMemoryCatalog) -> None: + initial = pa.table({"id": pa.array([1, 2, 3], type=pa.int64()), "value": ["one", "two", "three"]}) + table = catalog.create_table("default.equality_replace_and_delete", initial.schema) + table.append(initial) + previous_snapshot = table.current_snapshot() + assert previous_snapshot is not None + + table.upsert_by_equality_delete( + pa.table({"id": pa.array([1, 4], type=pa.int64()), "value": ["updated", "four"]}), + join_cols=["id"], + delete_df=pa.table({"id": pa.array([2], type=pa.int64())}), + ) + + current_snapshot = table.current_snapshot() + assert current_snapshot is not None + assert current_snapshot.parent_snapshot_id == previous_snapshot.snapshot_id + assert table.scan().to_arrow().sort_by("id").to_pydict() == { + "id": [1, 3, 4], + "value": ["updated", "three", "four"], + } + + def test_partitioned_equality_delete_upsert_requires_partition_source_in_key(catalog: InMemoryCatalog) -> None: schema = Schema( NestedField(1, "id", IntegerType(), required=True),