Add ManageSnapshots.cherry_pick_snapshot

mkdocs/docs/api.md

@@ -1483,6 +1483,39 @@ Remove an existing branch:
 table.manage_snapshots().remove_branch("dev").commit()
 ```
 
+### Write-Audit-Publish
+
+Stage a write on a branch, validate it, then publish it to the table with
+`cherry_pick_snapshot`. The staged data is invisible to readers of the table until it is
+published.
+
+```python
+# Write: stage the changes on an audit branch
+table.manage_snapshots().create_branch(
+    snapshot_id=table.metadata.current_snapshot_id,
+    branch_name="audit-2024-01-15",
+).commit()
+
+table = catalog.load_table("db.table")
+table.append(new_rows, branch="audit-2024-01-15", snapshot_properties={"wap.id": "etl-2024-01-15"})
+
+# Audit: validate the staged data without affecting readers of the table
+table = catalog.load_table("db.table")
+staged = table.metadata.refs["audit-2024-01-15"].snapshot_id
+assert len(table.scan(snapshot_id=staged).to_arrow()) > 0
+
+# Publish: replay the staged changes onto the current table state
+table.manage_snapshots().cherry_pick_snapshot(staged).commit()
+```
+
+The published snapshot records `source-snapshot-id`, and `published-wap-id` when the staged
+snapshot carried a `wap.id`. A given `wap.id` can only be published once.
+
+Append snapshots are always replayed, so the wap trail is recorded even when the table has not
+changed since the branch was cut. A snapshot with any other operation is fast-forwarded to when
+its parent is already the current snapshot, and raises otherwise. Picking a snapshot that is
+already an ancestor of the current state does nothing.
+
 ## Table Maintenance
 
 PyIceberg provides table maintenance operations through the `table.maintenance` API. This provides a clean interface for performing maintenance tasks like snapshot expiration.

pyiceberg/table/update/snapshot.py

@@ -80,6 +80,12 @@
 if TYPE_CHECKING:
     from pyiceberg.table import Transaction
 
+# Snapshot summary properties for Write-Audit-Publish, mirroring Java's SnapshotSummary:
+# https://github.com/apache/iceberg/blob/0d60b2bb8780d781f4ceb69032418d222b649ae5/core/src/main/java/org/apache/iceberg/SnapshotSummary.java#L60-L62
+_STAGED_WAP_ID_PROP = "wap.id"
+_PUBLISHED_WAP_ID_PROP = "published-wap-id"
+_SOURCE_SNAPSHOT_ID_PROP = "source-snapshot-id"
+
 
 def _new_manifest_file_name(num: int, commit_uuid: uuid.UUID) -> str:
     return f"{commit_uuid}-m{num}.avro"
@@ -974,6 +980,87 @@ def set_current_snapshot(self, snapshot_id: int | None = None, ref_name: str | N
         self._transaction._stage(update, requirement)
         return self
 
+    def cherry_pick_snapshot(self, snapshot_id: int) -> ManageSnapshots:
+        """Apply the changes in a snapshot to the current table state.
+
+        Creates a new snapshot on top of the current one carrying the data files the picked
+        snapshot added, which is how a Write-Audit-Publish staged write is published. The picked
+        snapshot itself is left in place.
+
+        Append snapshots are always replayed, matching Java's ``CherryPickOperation``. A snapshot
+        with another operation is fast-forwarded to when its parent is already the current
+        snapshot, and rejected otherwise. Picking a snapshot that is already an ancestor of the
+        current state does nothing.
+
+        A snapshot staged with a ``wap.id`` can only be published once. The new snapshot records
+        ``source-snapshot-id``, and ``published-wap-id`` when the picked snapshot carried one.
+
+        Args:
+            snapshot_id: The ID of the snapshot to cherry-pick.
+
+        Returns:
+            This for method chaining.
+
+        Raises:
+            ValueError: If the snapshot does not exist, its ``wap.id`` was already published, or
+                its operation cannot be cherry-picked.
+        """
+        self._commit_if_ref_updates_exist()
+
+        metadata = self._transaction.table_metadata
+        picked = metadata.snapshot_by_id(snapshot_id)
+        if picked is None:
+            raise ValueError(f"Cannot cherry-pick unknown snapshot id: {snapshot_id}")
+
+        if self._is_current_ancestor(snapshot_id):
+            return self
+
+        wap_id = self._validate_wap_publish(picked)
+        operation = picked.summary.operation if picked.summary else None
+
+        if operation == Operation.APPEND:
+            snapshot_properties = {_SOURCE_SNAPSHOT_ID_PROP: str(snapshot_id)}
+            if wap_id is not None:
+                snapshot_properties[_PUBLISHED_WAP_ID_PROP] = wap_id
+
+            with self._transaction.update_snapshot(snapshot_properties=snapshot_properties).fast_append() as append:
+                for data_file in self._added_data_files(picked):
+                    append.append_data_file(data_file)
+            return self
+
+        if picked.parent_snapshot_id == metadata.current_snapshot_id:
+            return self.set_current_snapshot(snapshot_id=snapshot_id)
+
+        raise ValueError(
+            f"Cannot cherry-pick snapshot {snapshot_id}: not append, dynamic overwrite, or fast-forward "
+            f"(operation: {operation}). Only append snapshots can be replayed."
+        )
+
+    def _added_data_files(self, snapshot: Snapshot) -> list[DataFile]:
+        """Return the data files the given snapshot added, excluding entries carried over from its ancestors."""
+        io = self._transaction._table.io
+        return [
+            entry.data_file
+            for manifest in snapshot.manifests(io)
+            if manifest.added_snapshot_id == snapshot.snapshot_id
+            for entry in manifest.fetch_manifest_entry(io, discard_deleted=True)
+            if entry.status == ManifestEntryStatus.ADDED and entry.snapshot_id == snapshot.snapshot_id
+        ]
+
+    def _validate_wap_publish(self, picked: Snapshot) -> str | None:
+        """Return the picked snapshot's staged wap id, rejecting one that an ancestor already published."""
+        wap_id = picked.summary.additional_properties.get(_STAGED_WAP_ID_PROP) if picked.summary else None
+        if not wap_id:
+            return None
+
+        metadata = self._transaction.table_metadata
+        for ancestor in ancestors_of(metadata.current_snapshot(), metadata):
+            props = ancestor.summary.additional_properties if ancestor.summary else {}
+            if wap_id in (props.get(_STAGED_WAP_ID_PROP), props.get(_PUBLISHED_WAP_ID_PROP)):
+                raise ValueError(f"Duplicate request to cherry pick wap id that was published already: {wap_id}")
+
+        return wap_id
+
     def rollback_to_snapshot(self, snapshot_id: int) -> ManageSnapshots:
         """Rollback the table to the given snapshot id.
 

tests/integration/test_snapshot_operations.py

@@ -332,3 +332,48 @@ def test_rollback_to_timestamp_chained_with_tag(table_with_snapshots: Table) ->
     assert table_with_snapshots.metadata.refs[tag_name] == SnapshotRef(
         snapshot_id=current_snapshot.snapshot_id, snapshot_ref_type="tag"
     )
+
+
+@pytest.mark.integration
+@pytest.mark.parametrize("catalog", [lf("session_catalog_hive"), lf("session_catalog")])
+def test_cherry_pick_snapshot_publishes_staged_write(catalog: Catalog) -> None:
+    """Publish a branch-staged write onto a main that has moved on since the branch was cut."""
+    catalog.create_namespace_if_not_exists("default")
+    identifier = f"default.test_cherry_pick_{uuid.uuid4().hex[:8]}"
+    arrow_schema = pa.schema([pa.field("id", pa.int64(), nullable=False)])
+    tbl = catalog.create_table(identifier=identifier, schema=arrow_schema)
+
+    tbl.append(pa.Table.from_pylist([{"id": 1}], schema=arrow_schema))
+    tbl = catalog.load_table(identifier)
+    current_snapshot_id = tbl.metadata.current_snapshot_id
+    assert current_snapshot_id is not None
+    tbl.manage_snapshots().create_branch(snapshot_id=current_snapshot_id, branch_name="audit").commit()
+
+    tbl = catalog.load_table(identifier)
+    tbl.append(
+        pa.Table.from_pylist([{"id": 2}, {"id": 3}], schema=arrow_schema),
+        branch="audit",
+        snapshot_properties={"wap.id": "etl-001"},
+    )
+    tbl = catalog.load_table(identifier)
+    staged = tbl.metadata.refs["audit"].snapshot_id
+
+    # main advances independently, so this is a replay rather than a fast-forward
+    tbl.append(pa.Table.from_pylist([{"id": 9}], schema=arrow_schema))
+    tbl = catalog.load_table(identifier)
+    assert sorted(tbl.scan().to_arrow().column("id").to_pylist()) == [1, 9]
+
+    tbl.manage_snapshots().cherry_pick_snapshot(staged).commit()
+
+    tbl = catalog.load_table(identifier)
+    assert sorted(tbl.scan().to_arrow().column("id").to_pylist()) == [1, 2, 3, 9]
+    published = tbl.current_snapshot()
+    assert published is not None and published.summary is not None
+    summary = published.summary.additional_properties
+    assert summary["source-snapshot-id"] == str(staged)
+    assert summary["published-wap-id"] == "etl-001"
+
+    with pytest.raises(ValueError, match="Duplicate request to cherry pick wap id"):
+        tbl.manage_snapshots().cherry_pick_snapshot(staged).commit()
+
+    catalog.drop_table(identifier)

tests/table/test_manage_snapshots.py

@@ -17,8 +17,10 @@
 from unittest.mock import MagicMock
 from uuid import uuid4
 
+import pyarrow as pa
 import pytest
 
+from pyiceberg.catalog import Catalog
 from pyiceberg.table import CommitTableResponse, Table
 from pyiceberg.table.update import SetSnapshotRefUpdate, TableUpdate
 
@@ -177,3 +179,172 @@ def test_set_current_snapshot_chained_with_create_tag(table_v2: Table) -> None:
     # The main branch should point to the same snapshot as the tag
     main_update = next(u for u in set_ref_updates if u.ref_name == "main")
     assert main_update.snapshot_id == snapshot_one
+
+
+def _staged_wap_table(
+    catalog_with_warehouse: Catalog,
+    namespace: str,
+    advance_main: bool = True,
+    wap_id: str | None = None,
+) -> tuple[Table, int]:
+    """Build the WAP shape: a staged branch write, with main optionally moved on since.
+
+    Returns the reloaded table and the staged snapshot id.
+    """
+    catalog_with_warehouse.create_namespace(namespace)
+    schema = pa.schema([pa.field("id", pa.int64())])
+    tbl = catalog_with_warehouse.create_table(f"{namespace}.tbl", schema=schema)
+
+    tbl.append(pa.table({"id": [1]}, schema=schema))
+    tbl = catalog_with_warehouse.load_table(f"{namespace}.tbl")
+
+    tbl.manage_snapshots().create_branch(snapshot_id=_current_snapshot_id(tbl), branch_name="audit").commit()
+    tbl = catalog_with_warehouse.load_table(f"{namespace}.tbl")
+
+    tbl.append(
+        pa.table({"id": [2, 3]}, schema=schema),
+        branch="audit",
+        snapshot_properties={"wap.id": wap_id} if wap_id else {},
+    )
+    tbl = catalog_with_warehouse.load_table(f"{namespace}.tbl")
+    staged = tbl.metadata.refs["audit"].snapshot_id
+
+    if advance_main:
+        tbl.append(pa.table({"id": [9]}, schema=schema))
+        tbl = catalog_with_warehouse.load_table(f"{namespace}.tbl")
+
+    return tbl, staged
+
+
+def _ids(tbl: Table) -> list[int]:
+    return sorted(tbl.scan().to_arrow().column("id").to_pylist())
+
+
+def _current_summary_props(tbl: Table) -> dict[str, str]:
+    snapshot = tbl.current_snapshot()
+    assert snapshot is not None and snapshot.summary is not None
+    return snapshot.summary.additional_properties
+
+
+def _current_snapshot_id(tbl: Table) -> int:
+    snapshot_id = tbl.metadata.current_snapshot_id
+    assert snapshot_id is not None
+    return snapshot_id
+
+
+def test_cherry_pick_snapshot_replays_staged_append(catalog_with_warehouse: Catalog) -> None:
+    """The staged rows land on main even though main moved on after the branch was cut."""
+    tbl, staged = _staged_wap_table(catalog_with_warehouse, "cp_replay")
+    assert _ids(tbl) == [1, 9]
+
+    tbl.manage_snapshots().cherry_pick_snapshot(staged).commit()
+
+    tbl = catalog_with_warehouse.load_table("cp_replay.tbl")
+    assert _ids(tbl) == [1, 2, 3, 9]
+    assert _current_summary_props(tbl)["source-snapshot-id"] == str(staged)
+
+
+def test_cherry_pick_snapshot_records_published_wap_id(catalog_with_warehouse: Catalog) -> None:
+    tbl, staged = _staged_wap_table(catalog_with_warehouse, "cp_wap", wap_id="etl-001")
+
+    tbl.manage_snapshots().cherry_pick_snapshot(staged).commit()
+
+    tbl = catalog_with_warehouse.load_table("cp_wap.tbl")
+    assert _current_summary_props(tbl)["published-wap-id"] == "etl-001"
+
+
+def test_cherry_pick_snapshot_rejects_duplicate_wap_publish(catalog_with_warehouse: Catalog) -> None:
+    """A wap id may only be published once."""
+    tbl, staged = _staged_wap_table(catalog_with_warehouse, "cp_dup", wap_id="etl-001")
+    tbl.manage_snapshots().cherry_pick_snapshot(staged).commit()
+
+    tbl = catalog_with_warehouse.load_table("cp_dup.tbl")
+    with pytest.raises(ValueError, match="Duplicate request to cherry pick wap id"):
+        tbl.manage_snapshots().cherry_pick_snapshot(staged).commit()
+
+
+def test_cherry_pick_snapshot_replays_append_even_when_parent_is_current(
+    catalog_with_warehouse: Catalog,
+) -> None:
+    """Appends are always replayed, so the wap trail is recorded even on an unmoved table."""
+    tbl, staged = _staged_wap_table(catalog_with_warehouse, "cp_ff", advance_main=False, wap_id="etl-001")
+    assert _ids(tbl) == [1]
+
+    tbl.manage_snapshots().cherry_pick_snapshot(staged).commit()
+
+    tbl = catalog_with_warehouse.load_table("cp_ff.tbl")
+    assert _ids(tbl) == [1, 2, 3]
+    assert tbl.metadata.current_snapshot_id != staged
+    assert _current_summary_props(tbl)["published-wap-id"] == "etl-001"
+
+
+def test_cherry_pick_snapshot_fast_forwards_non_append(catalog_with_warehouse: Catalog) -> None:
+    """A non-append whose parent is current cannot be replayed, so the ref advances to it."""
+    catalog_with_warehouse.create_namespace("cp_ffna")
+    schema = pa.schema([pa.field("id", pa.int64())])
+    tbl = catalog_with_warehouse.create_table("cp_ffna.tbl", schema=schema)
+    tbl.append(pa.table({"id": [1, 2]}, schema=schema))
+    tbl = catalog_with_warehouse.load_table("cp_ffna.tbl")
+
+    tbl.manage_snapshots().create_branch(snapshot_id=_current_snapshot_id(tbl), branch_name="audit").commit()
+    tbl = catalog_with_warehouse.load_table("cp_ffna.tbl")
+    tbl.delete("id = 1", branch="audit")
+    tbl = catalog_with_warehouse.load_table("cp_ffna.tbl")
+    staged = tbl.metadata.refs["audit"].snapshot_id
+
+    tbl.manage_snapshots().cherry_pick_snapshot(staged).commit()
+
+    tbl = catalog_with_warehouse.load_table("cp_ffna.tbl")
+    assert tbl.metadata.current_snapshot_id == staged
+    assert _ids(tbl) == [2]
+
+
+def test_cherry_pick_snapshot_is_noop_for_ancestor(catalog_with_warehouse: Catalog) -> None:
+    """Picking a snapshot already in main's history changes nothing."""
+    catalog_with_warehouse.create_namespace("cp_anc")
+    schema = pa.schema([pa.field("id", pa.int64())])
+    tbl = catalog_with_warehouse.create_table("cp_anc.tbl", schema=schema)
+    tbl.append(pa.table({"id": [1]}, schema=schema))
+    tbl = catalog_with_warehouse.load_table("cp_anc.tbl")
+    first = _current_snapshot_id(tbl)
+    tbl.append(pa.table({"id": [2]}, schema=schema))
+    tbl = catalog_with_warehouse.load_table("cp_anc.tbl")
+    before = _current_snapshot_id(tbl)
+
+    tbl.manage_snapshots().cherry_pick_snapshot(first).commit()
+
+    tbl = catalog_with_warehouse.load_table("cp_anc.tbl")
+    assert tbl.metadata.current_snapshot_id == before
+    assert _ids(tbl) == [1, 2]
+
+
+def test_cherry_pick_snapshot_rejects_unknown_snapshot(catalog_with_warehouse: Catalog) -> None:
+    catalog_with_warehouse.create_namespace("cp_unknown")
+    schema = pa.schema([pa.field("id", pa.int64())])
+    tbl = catalog_with_warehouse.create_table("cp_unknown.tbl", schema=schema)
+    tbl.append(pa.table({"id": [1]}, schema=schema))
+    tbl = catalog_with_warehouse.load_table("cp_unknown.tbl")
+
+    with pytest.raises(ValueError, match="Cannot cherry-pick unknown snapshot id"):
+        tbl.manage_snapshots().cherry_pick_snapshot(1234567890).commit()
+
+
+def test_cherry_pick_snapshot_rejects_non_append(catalog_with_warehouse: Catalog) -> None:
+    """Only append snapshots can be replayed; anything else must say so rather than silently skip."""
+    catalog_with_warehouse.create_namespace("cp_nonappend")
+    schema = pa.schema([pa.field("id", pa.int64())])
+    tbl = catalog_with_warehouse.create_table("cp_nonappend.tbl", schema=schema)
+    tbl.append(pa.table({"id": [1, 2]}, schema=schema))
+    tbl = catalog_with_warehouse.load_table("cp_nonappend.tbl")
+
+    tbl.manage_snapshots().create_branch(snapshot_id=_current_snapshot_id(tbl), branch_name="audit").commit()
+    tbl = catalog_with_warehouse.load_table("cp_nonappend.tbl")
+    tbl.delete("id = 1", branch="audit")
+    tbl = catalog_with_warehouse.load_table("cp_nonappend.tbl")
+    staged = tbl.metadata.refs["audit"].snapshot_id
+
+    tbl.append(pa.table({"id": [9]}, schema=schema))
+    tbl = catalog_with_warehouse.load_table("cp_nonappend.tbl")
+
+    with pytest.raises(ValueError, match="not append, dynamic overwrite, or fast-forward"):
+        tbl.manage_snapshots().cherry_pick_snapshot(staged).commit()