|
15 | 15 | # specific language governing permissions and limitations |
16 | 16 | # under the License. |
17 | 17 | # pylint:disable=redefined-outer-name,eval-used |
| 18 | +import re |
| 19 | +import uuid |
| 20 | +from pathlib import Path |
18 | 21 | from typing import cast |
| 22 | +from urllib.parse import urlparse |
19 | 23 |
|
| 24 | +import pyarrow as pa |
20 | 25 | import pytest |
21 | 26 |
|
| 27 | +from pyiceberg.catalog import Catalog |
| 28 | +from pyiceberg.exceptions import ValidationException |
| 29 | +from pyiceberg.io.pyarrow import _dataframe_to_data_files |
22 | 30 | from pyiceberg.manifest import DataFile, DataFileContent, ManifestContent, ManifestFile |
23 | 31 | from pyiceberg.partitioning import PartitionField, PartitionSpec |
24 | 32 | from pyiceberg.schema import Schema |
@@ -649,3 +657,83 @@ def summary_calls(n_files: int) -> int: |
649 | 657 | f"_MergeAppendFiles.__init__ made {merge_init - fast_init} extra update_table_metadata " |
650 | 658 | "calls over its superclass; expected 1 (hoisted)" |
651 | 659 | ) |
| 660 | + |
| 661 | + |
| 662 | +@pytest.fixture |
| 663 | +def overwrite_table(catalog: Catalog, arrow_table_simple: pa.Table) -> Table: |
| 664 | + catalog.create_namespace("default") |
| 665 | + table = catalog.create_table("default.overwrite", arrow_table_simple.schema) |
| 666 | + table.append(arrow_table_simple) |
| 667 | + return table |
| 668 | + |
| 669 | + |
| 670 | +def _write_data_file(table: Table, rows: pa.Table) -> DataFile: |
| 671 | + return next( |
| 672 | + iter( |
| 673 | + _dataframe_to_data_files( |
| 674 | + table_metadata=table.metadata, |
| 675 | + df=rows, |
| 676 | + io=table.io, |
| 677 | + write_uuid=uuid.uuid4(), |
| 678 | + ) |
| 679 | + ) |
| 680 | + ) |
| 681 | + |
| 682 | + |
| 683 | +def _total_data_file_count(table: Table) -> int: |
| 684 | + snapshot = table.current_snapshot() |
| 685 | + assert snapshot is not None and snapshot.summary is not None |
| 686 | + return int(snapshot.summary.additional_properties["total-data-files"]) |
| 687 | + |
| 688 | + |
| 689 | +def test_overwrite_replaces_existing_file(overwrite_table: Table, arrow_table_simple: pa.Table) -> None: |
| 690 | + original_file = next(iter(overwrite_table.scan().plan_files())).file |
| 691 | + replacement = _write_data_file(overwrite_table, arrow_table_simple.slice(0, 1)) |
| 692 | + |
| 693 | + with overwrite_table.transaction() as tx: |
| 694 | + with tx.update_snapshot().overwrite() as overwrite: |
| 695 | + overwrite.delete_data_file(original_file) |
| 696 | + overwrite.append_data_file(replacement) |
| 697 | + |
| 698 | + assert overwrite_table.scan().to_arrow()["foo"].to_pylist() == ["a"] |
| 699 | + assert _total_data_file_count(overwrite_table) == 1 |
| 700 | + |
| 701 | + |
| 702 | +def test_overwrite_rejects_explicit_delete_missing_from_base_snapshot(catalog: Catalog, overwrite_table: Table) -> None: |
| 703 | + stale_file = next(iter(overwrite_table.scan().plan_files())).file |
| 704 | + stale_rows = overwrite_table.scan().to_arrow() |
| 705 | + |
| 706 | + # Delete the file before the replacement transaction begins |
| 707 | + with catalog.load_table(overwrite_table.name()).transaction() as tx: |
| 708 | + with tx.update_snapshot().overwrite() as overwrite: |
| 709 | + overwrite.delete_data_file(stale_file) |
| 710 | + |
| 711 | + current = catalog.load_table(overwrite_table.name()) |
| 712 | + replacement = _write_data_file(current, stale_rows) |
| 713 | + expected_error = re.escape(f"Missing required files to delete: {stale_file.file_path}") |
| 714 | + |
| 715 | + with pytest.raises(ValidationException, match=expected_error): |
| 716 | + with current.transaction() as tx: |
| 717 | + with tx.update_snapshot().overwrite() as overwrite: |
| 718 | + overwrite.delete_data_file(stale_file) |
| 719 | + overwrite.append_data_file(replacement) |
| 720 | + |
| 721 | + metadata_path = Path(urlparse(current.location()).path) / "metadata" |
| 722 | + assert not any(metadata_path.glob(f"{overwrite.commit_uuid}-m*.avro")) |
| 723 | + |
| 724 | + committed = catalog.load_table(overwrite_table.name()) |
| 725 | + assert committed.scan().to_arrow()["foo"].to_pylist() == [] |
| 726 | + assert _total_data_file_count(committed) == 0 |
| 727 | + |
| 728 | + |
| 729 | +def test_overwrite_rejects_explicit_delete_without_parent_snapshot( |
| 730 | + catalog: Catalog, overwrite_table: Table, arrow_table_simple: pa.Table |
| 731 | +) -> None: |
| 732 | + stale_file = next(iter(overwrite_table.scan().plan_files())).file |
| 733 | + empty = catalog.create_table("default.empty", arrow_table_simple.schema) |
| 734 | + expected_error = re.escape(f"Missing required files to delete: {stale_file.file_path}") |
| 735 | + |
| 736 | + with pytest.raises(ValidationException, match=expected_error): |
| 737 | + with empty.transaction() as tx: |
| 738 | + with tx.update_snapshot().overwrite() as overwrite: |
| 739 | + overwrite.delete_data_file(stale_file) |
0 commit comments