diff --git a/pyiceberg/table/update/snapshot.py b/pyiceberg/table/update/snapshot.py index 57215dca04..e500e09fac 100644 --- a/pyiceberg/table/update/snapshot.py +++ b/pyiceberg/table/update/snapshot.py @@ -472,6 +472,7 @@ def _validate_concurrency(self) -> None: from pyiceberg.table.snapshots import IsolationLevel from pyiceberg.table.update.validate import ( _validate_added_data_files, + _validate_data_files_exist, _validate_deleted_data_files, _validate_no_new_delete_files, _validate_no_new_deletes_for_data_files, @@ -501,6 +502,7 @@ def _validate_concurrency(self) -> None: _validate_deleted_data_files(table, catalog_head, conflict_detection_filter, starting_snapshot) if self._deleted_data_files: + _validate_data_files_exist(table, catalog_head, self._deleted_data_files, starting_snapshot) _validate_no_new_deletes_for_data_files( table, catalog_head, conflict_detection_filter, self._deleted_data_files, starting_snapshot ) diff --git a/pyiceberg/table/update/validate.py b/pyiceberg/table/update/validate.py index df8506aab4..0545182bf0 100644 --- a/pyiceberg/table/update/validate.py +++ b/pyiceberg/table/update/validate.py @@ -14,6 +14,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from collections import defaultdict from collections.abc import Iterator from pyiceberg.exceptions import ValidationException @@ -195,6 +196,33 @@ def _validate_deleted_data_files( raise ValidationException(f"Deleted data files were found matching the filter for snapshots {conflicting_snapshots}!") +def _validate_data_files_exist( + table: Table, + starting_snapshot: Snapshot, + data_files: set[DataFile], + parent_snapshot: Snapshot | None, +) -> None: + """Validate that explicitly replaced data files have not been concurrently deleted. + + Args: + table: Table to validate + starting_snapshot: Snapshot at the end of the validation window + data_files: Data files that must still exist + parent_snapshot: Snapshot at the start of the validation window, excluded from the scan + """ + partition_set: dict[int, set[Record]] = defaultdict(set) + for data_file in data_files: + partition_set[data_file.spec_id].add(data_file.partition) + + conflicting_paths = { + entry.data_file.file_path + for entry in _deleted_data_files(table, starting_snapshot, None, partition_set, parent_snapshot) + if entry.data_file in data_files + } + if conflicting_paths: + raise ValidationException(f"Data files were concurrently deleted: {sorted(conflicting_paths)}") + + def _added_data_files( table: Table, starting_snapshot: Snapshot, diff --git a/tests/table/test_commit_retry.py b/tests/table/test_commit_retry.py index ce5dca96aa..ab95457e23 100644 --- a/tests/table/test_commit_retry.py +++ b/tests/table/test_commit_retry.py @@ -14,6 +14,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +import uuid from typing import Any from unittest.mock import patch @@ -323,6 +324,77 @@ def test_concurrent_overwrite_overwrite_raises_validation_exception(catalog: Cat tbl2.overwrite(pa.table({"x": [40, 50, 60]}), overwrite_filter="x > 0") +@pytest.mark.parametrize( + ("concurrently_deleted_file", "expect_conflict", "expected_values"), + [ + pytest.param("target", True, [1, 3], id="target-file"), + pytest.param("same-partition", False, [1, 2], id="same-partition-file"), + pytest.param("different-partition", False, [2, 3], id="different-partition-file"), + ], +) +def test_file_overwrite_validates_concurrent_file_delete( + catalog: Catalog, + concurrently_deleted_file: str, + expect_conflict: bool, + expected_values: list[int], +) -> None: + """A file replacement must fail only when its target file was concurrently deleted.""" + import pyarrow as pa + + from pyiceberg.io.pyarrow import _dataframe_to_data_files + from pyiceberg.partitioning import PartitionField, PartitionSpec + from pyiceberg.transforms import IdentityTransform + + catalog.create_namespace("default") + schema = Schema( + NestedField(1, "category", StringType(), required=False), + NestedField(2, "value", LongType(), required=False), + ) + spec = PartitionSpec(PartitionField(source_id=1, field_id=1000, transform=IdentityTransform(), name="category")) + identifier = "default.concurrent_file_delete" + table = catalog.create_table(identifier, schema=schema, partition_spec=spec) + table.append(pa.table({"category": ["a", "b"], "value": [0, 1]})) + file_to_replace = next(task.file for task in table.scan().plan_files() if task.file.partition[0] == "a") + table.append(pa.table({"category": ["a"], "value": [3]})) + + replacing_table = catalog.load_table(identifier) + deleting_table = catalog.load_table(identifier) + data_files = [task.file for task in deleting_table.scan().plan_files()] + file_to_delete = { + "target": file_to_replace, + "same-partition": next( + data_file for data_file in data_files if data_file.partition[0] == "a" and data_file != file_to_replace + ), + "different-partition": next(data_file for data_file in data_files if data_file.partition[0] == "b"), + }[concurrently_deleted_file] + + replacement_file = list( + _dataframe_to_data_files( + table_metadata=replacing_table.metadata, + df=pa.table({"category": ["a"], "value": [2]}), + io=replacing_table.io, + write_uuid=uuid.uuid4(), + ) + )[0] + replacing_transaction = replacing_table.transaction() + with replacing_transaction.update_snapshot().overwrite() as overwrite: + overwrite.delete_data_file(file_to_replace) + overwrite.append_data_file(replacement_file) + + with deleting_table.transaction() as deleting_transaction: + with deleting_transaction.update_snapshot().overwrite() as overwrite: + overwrite.delete_data_file(file_to_delete) + + if expect_conflict: + with pytest.raises(ValidationException, match="Data files were concurrently deleted"): + replacing_transaction.commit_transaction() + else: + replacing_transaction.commit_transaction() + + result = catalog.load_table(identifier).scan().to_arrow() + assert sorted(result["value"].to_pylist()) == expected_values + + def test_concurrent_overwrite_append_retries_successfully(catalog: Catalog) -> None: """Append after a concurrent overwrite should succeed via retry.""" catalog.create_namespace("default")