diff --git a/pyiceberg/table/update/snapshot.py b/pyiceberg/table/update/snapshot.py index 7931edacdd..afbeef7e99 100644 --- a/pyiceberg/table/update/snapshot.py +++ b/pyiceberg/table/update/snapshot.py @@ -26,6 +26,7 @@ from typing import TYPE_CHECKING, Generic from pyiceberg.avro.codecs import AvroCompressionCodec +from pyiceberg.exceptions import ValidationException from pyiceberg.expressions import AlwaysFalse, BooleanExpression, Or from pyiceberg.expressions.visitors import ( ROWS_MIGHT_NOT_MATCH, @@ -196,8 +197,6 @@ def _write_added_manifest() -> list[ManifestFile]: return [] def _write_delete_manifest() -> list[ManifestFile]: - # Check if we need to mark the files as deleted - deleted_entries = self._deleted_entries() if len(deleted_entries) > 0: deleted_manifests = [] partition_groups: dict[int, list[ManifestEntry]] = defaultdict(list) @@ -214,6 +213,8 @@ def _write_delete_manifest() -> list[ManifestFile]: # Updates self._predicate with computed partition predicate for manifest pruning self._build_delete_files_partition_predicate() + # Plan deletes before starting manifest writers so validation failures do not leave orphaned manifests + deleted_entries = self._deleted_entries() executor = ExecutorFactory.get_or_create() @@ -667,9 +668,30 @@ def _get_entries(manifest: ManifestFile) -> list[ManifestEntry]: ] list_of_entries = executor.map(_get_entries, previous_snapshot.manifests(self._io)) - return list(itertools.chain(*list_of_entries)) + deleted_entries = list(itertools.chain(*list_of_entries)) else: - return [] + deleted_entries = [] + + self._validate_required_deletes(deleted_entries) + + return deleted_entries + + def _validate_required_deletes(self, deleted_entries: list[ManifestEntry]) -> None: + """Validate that explicitly deleted data files exist in the parent snapshot. + + Files passed to `delete_data_file` are required deletes. If one is absent, an overwrite + could commit replacement files without the corresponding deletion and produce incorrect + snapshot summary totals. + + Args: + deleted_entries: Live parent-snapshot entries selected for deletion. + + Raises: + ValidationException: If a required data file is missing. + """ + found_data_files = {entry.data_file for entry in deleted_entries} + if missing := [data_file.file_path for data_file in self._deleted_data_files if data_file not in found_data_files]: + raise ValidationException(f"Missing required files to delete: {', '.join(sorted(missing))}") class UpdateSnapshot: diff --git a/tests/table/test_snapshots.py b/tests/table/test_snapshots.py index 5f1680ed59..0f72b08087 100644 --- a/tests/table/test_snapshots.py +++ b/tests/table/test_snapshots.py @@ -15,10 +15,18 @@ # specific language governing permissions and limitations # under the License. # pylint:disable=redefined-outer-name,eval-used +import re +import uuid +from pathlib import Path from typing import cast +from urllib.parse import urlparse +import pyarrow as pa import pytest +from pyiceberg.catalog import Catalog +from pyiceberg.exceptions import ValidationException +from pyiceberg.io.pyarrow import _dataframe_to_data_files from pyiceberg.manifest import DataFile, DataFileContent, ManifestContent, ManifestFile from pyiceberg.partitioning import PartitionField, PartitionSpec from pyiceberg.schema import Schema @@ -649,3 +657,83 @@ def summary_calls(n_files: int) -> int: f"_MergeAppendFiles.__init__ made {merge_init - fast_init} extra update_table_metadata " "calls over its superclass; expected 1 (hoisted)" ) + + +@pytest.fixture +def overwrite_table(catalog: Catalog, arrow_table_simple: pa.Table) -> Table: + catalog.create_namespace("default") + table = catalog.create_table("default.overwrite", arrow_table_simple.schema) + table.append(arrow_table_simple) + return table + + +def _write_data_file(table: Table, rows: pa.Table) -> DataFile: + return next( + iter( + _dataframe_to_data_files( + table_metadata=table.metadata, + df=rows, + io=table.io, + write_uuid=uuid.uuid4(), + ) + ) + ) + + +def _total_data_file_count(table: Table) -> int: + snapshot = table.current_snapshot() + assert snapshot is not None and snapshot.summary is not None + return int(snapshot.summary.additional_properties["total-data-files"]) + + +def test_overwrite_replaces_existing_file(overwrite_table: Table, arrow_table_simple: pa.Table) -> None: + original_file = next(iter(overwrite_table.scan().plan_files())).file + replacement = _write_data_file(overwrite_table, arrow_table_simple.slice(0, 1)) + + with overwrite_table.transaction() as tx: + with tx.update_snapshot().overwrite() as overwrite: + overwrite.delete_data_file(original_file) + overwrite.append_data_file(replacement) + + assert overwrite_table.scan().to_arrow()["foo"].to_pylist() == ["a"] + assert _total_data_file_count(overwrite_table) == 1 + + +def test_overwrite_rejects_explicit_delete_missing_from_base_snapshot(catalog: Catalog, overwrite_table: Table) -> None: + stale_file = next(iter(overwrite_table.scan().plan_files())).file + stale_rows = overwrite_table.scan().to_arrow() + + # Delete the file before the replacement transaction begins + with catalog.load_table(overwrite_table.name()).transaction() as tx: + with tx.update_snapshot().overwrite() as overwrite: + overwrite.delete_data_file(stale_file) + + current = catalog.load_table(overwrite_table.name()) + replacement = _write_data_file(current, stale_rows) + expected_error = re.escape(f"Missing required files to delete: {stale_file.file_path}") + + with pytest.raises(ValidationException, match=expected_error): + with current.transaction() as tx: + with tx.update_snapshot().overwrite() as overwrite: + overwrite.delete_data_file(stale_file) + overwrite.append_data_file(replacement) + + metadata_path = Path(urlparse(current.location()).path) / "metadata" + assert not any(metadata_path.glob(f"{overwrite.commit_uuid}-m*.avro")) + + committed = catalog.load_table(overwrite_table.name()) + assert committed.scan().to_arrow()["foo"].to_pylist() == [] + assert _total_data_file_count(committed) == 0 + + +def test_overwrite_rejects_explicit_delete_without_parent_snapshot( + catalog: Catalog, overwrite_table: Table, arrow_table_simple: pa.Table +) -> None: + stale_file = next(iter(overwrite_table.scan().plan_files())).file + empty = catalog.create_table("default.empty", arrow_table_simple.schema) + expected_error = re.escape(f"Missing required files to delete: {stale_file.file_path}") + + with pytest.raises(ValidationException, match=expected_error): + with empty.transaction() as tx: + with tx.update_snapshot().overwrite() as overwrite: + overwrite.delete_data_file(stale_file)