feat: add ManageSnapshots.fast_forward_branch (#3649)

* feat: add ManageSnapshots.fast_forward_branch

test: fast_forward_branch auto-creates missing from_branch
test: fast_forward_branch is a no-op when snapshots already equal
test: fast_forward_branch rejects a tag as the source ref
test: fast_forward_branch rejects missing to_ref
test: fast_forward_branch rejects non-ancestor target
test: fast_forward_branch preserves retention fields
test: fast_forward_branch composes in a manage_snapshots chain
test(integration): fast_forward_branch on real catalogs

docs: add fast_forward_branch section with WAP example
chore: apply ruff-format and add None-checks for mypy

* feat: intra-chain ref lookups via _effective_refs in ManageSnapshots

* test(integration): fast_forward_branch preserves retention intra-chain

* docs(fix): Cleanup Example & Add Mermaid Diagram

* tests(PR Comments 1): Add fast_forward tag case
tests(PR Comments 2): Add more noop cases
tests(cleanup): Rename tests with naming format `test__{method}__with_{condition}__{outcome}`.
tests(cleanup): Move tests into test class groupings
tests(cleanup): Added explicit `table_v2_main_behind` fixture & update the tests to have main be fast-forwarded to reduce confusion against WAP process.

* NIT: snapshot.py | Using walrus notation for single lookup.

Co-authored-by: Fokko Driesprong <fokko@apache.org>

* Trigger Build

* fix: Bug introduced by single lookup walrus operation

---------

Co-authored-by: Fokko Driesprong <fokko@apache.org>

Author: JeroenSchmidt

mkdocs/docs/api.md

@@ -1483,6 +1483,110 @@ Remove an existing branch:
 table.manage_snapshots().remove_branch("dev").commit()
 ```
 
+#### Fast-forwarding a branch
+
+Fast-forward the `main` branch to the `audit-branch` branch:
+
+```python
+with table.manage_snapshots() as ms:
+    ms.fast_forward_branch(from_branch="main", to_ref="audit-branch")
+```
+
+Fast-forward `from_branch` to point at the snapshot referenced by `to_ref`.
+`to_ref` may be a branch or tag. `from_branch` must be a branch.
+
+<!-- markdownlint-disable MD046 -- Allowing indented multi-line formatting in admonition-->
+
+!!! info "Fast Forward Behavior"
+
+    * Case 1: If `from_branch` does not yet exist it is created and pointing at `to_ref`'s
+        snapshot. The default retention properties are applied on the auto-created snapshot.
+    * Case 2:** If both already point at the same snapshot the call is a no-op.
+    * Case 3: Otherwise `from_branch`'s current snapshot must be an ancestor of `to_ref`'s snapshot;
+    if not, `NotAncestorError` is raised.
+
+<!-- markdownlint-enable MD046 -->
+
+#### Example Use-Case: write-audit-publish (WAP)
+
+The use of branching & fast-forwarding enable the usage of the write-audit-publish (WAP) process:
+
+1. Writes proceed on a side branch
+2. Audit Validation runs against that branch
+3. Publish the new data by fast-forwarding the main branch
+
+```mermaid
+---
+title: Conceptually Illustration of the WAP Process
+---
+flowchart LR
+
+    subgraph audit [audit branch]
+        s1_audit["snapshot_1"] -- "1.2 append(new_rows)" --> s2_audit["snapshot_2"]
+        v@{ shape: comment, label: '2. Validation Performed & Passed' }
+         s2_audit ~~~ v
+         v -.-> s2_audit
+    end
+
+    subgraph main [main branch]
+        s1["snapshot_1"]
+        s2_main["snapshot_2"]
+
+    end
+
+    s1 -. "1.1 create_branch" .-> s1_audit
+    s2_audit -. "3. fast_forward_branch" .-> s2_main
+
+```
+
+If validation fails, callers simply skip the fast-forward step. The
+audit branch (and its data files) can then be inspected, rewritten,
+or removed via `remove_branch` and subsequent snapshot expiration -
+without ever having polluted the data on `main`.
+
+##### Programmatic Example
+
+```python
+import pyarrow as pa
+import pyarrow.compute as pc
+from pyiceberg.catalog import load_catalog
+
+catalog = load_catalog("prod")
+table = catalog.load_table("sales.orders")
+
+# 1. WRITE — create a side branch off main and append to it.
+# 1.1 Create the Branch
+main_snapshot_id = table.current_snapshot().snapshot_id
+table.manage_snapshots().create_branch(
+    snapshot_id=main_snapshot_id,
+    branch_name="audit",
+).commit()
+
+new_rows = pa.table({
+    "order_id": [1001, 1002, 1003],
+    "amount":   [ 49.99, 129.00, 12.50],
+})
+
+# 1.2 Write into the branch
+table.append(new_rows, branch="audit")
+
+# 2. AUDIT — scan the audit branch and run whatever validation your data-quality contract requires.
+#    Nothing on `main` has changed yet, so readers of `main` still see the pre-write state.
+audit_snapshot_id = table.refs()["audit"].snapshot_id
+audit_data = table.scan(snapshot_id=audit_snapshot_id).to_arrow()
+
+assert audit_data.num_rows > 0, "audit branch is empty"
+assert pc.all(pc.greater(audit_data["amount"], 0)).as_py(), \
+    "found non-positive amounts"
+
+# 3. PUBLISH — validation passed; fast-forward main to audit.
+#    fast-forward can occur because the audit branch was created from main and
+#    main's current snapshot is an ancestor of audit's snapshot.
+with table.manage_snapshots() as ms:
+    ms.fast_forward_branch("main", "audit")
+    ms.remove_branch("audit")  # optional: clean up the side branch
+```
+
 ## Table Maintenance
 
 PyIceberg provides table maintenance operations through the `table.maintenance` API. This provides a clean interface for performing maintenance tasks like snapshot expiration.

mkdocs/mkdocs.yml

@@ -61,5 +61,15 @@ markdown_extensions:
   - pymdownx.highlight:
       anchor_linenums: true
   - pymdownx.superfences
+  - pymdownx.superfences:
+      preserve_tabs: true
+      custom_fences:
+        # Mermaid diagrams
+        #   Needed so that Superfences doesn't break Mermaid
+        #   See: https://facelessuser.github.io/pymdown-extensions/extras/mermaid/#using-in-mkdocs
+        - name: mermaid
+          class: mermaid
+          format: !!python/name:pymdownx.superfences.fence_code_format
+
   - toc:
       permalink: true

pyiceberg/exceptions.py

@@ -146,3 +146,15 @@ class WaitingForLockException(Exception):
 
 class ValidationException(Exception):
     """Raised when validation fails."""
+
+
+class NoSuchSnapshotRefError(ValueError):
+    """Raised when a named snapshot ref (branch or tag) does not exist."""
+
+
+class SnapshotRefTypeError(ValueError):
+    """Raised when an operation expects a branch and gets a tag (or vice versa)."""
+
+
+class NotAncestorError(ValueError):
+    """Raised when an operation requires ancestry between two snapshots and it does not hold."""

pyiceberg/table/update/snapshot.py

@@ -28,7 +28,12 @@
 from typing import TYPE_CHECKING, Generic
 
 from pyiceberg.avro.codecs import AvroCompressionCodec
-from pyiceberg.exceptions import ValidationException
+from pyiceberg.exceptions import (
+    NoSuchSnapshotRefError,
+    NotAncestorError,
+    SnapshotRefTypeError,
+    ValidationException,
+)
 from pyiceberg.expressions import AlwaysFalse, BooleanExpression, Or
 from pyiceberg.expressions.visitors import (
     ROWS_MIGHT_NOT_MATCH,
@@ -52,13 +57,14 @@
 )
 from pyiceberg.partitioning import PartitionSpec
 from pyiceberg.schema import Schema
-from pyiceberg.table.refs import MAIN_BRANCH, SnapshotRefType
+from pyiceberg.table.refs import MAIN_BRANCH, SnapshotRef, SnapshotRefType
 from pyiceberg.table.snapshots import (
     Operation,
     Snapshot,
     SnapshotSummaryCollector,
     Summary,
     ancestors_of,
+    is_ancestor_of,
     latest_ancestor_before_timestamp,
     update_snapshot_summaries,
 )
@@ -1043,6 +1049,35 @@ def _commit_if_ref_updates_exist(self) -> None:
             self._updates = ()
             self._requirements = ()
 
+    def _effective_refs(self) -> dict[str, SnapshotRef]:
+        """Return refs as they would appear after all currently-staged updates.
+
+        Committed refs from ``table_metadata.refs`` overlaid with the effects
+        of every ``SetSnapshotRefUpdate`` / ``RemoveSnapshotRefUpdate`` that has
+        been accumulated onto ``self._updates`` in this chain, in order. Later
+        stages win. Callers use this instead of ``table_metadata.refs`` when a
+        decision needs to observe the results of earlier operations in the
+        same ``manage_snapshots()`` chain.
+
+        Note that this projection is for *decision-making* only. Requirements
+        emitted via ``_set_ref_snapshot`` continue to reference committed
+        state, which is what the catalog checks at commit time and what makes
+        concurrent-write detection correct.
+        """
+        refs: dict[str, SnapshotRef] = dict(self._transaction.table_metadata.refs)
+        for update in self._updates:
+            if isinstance(update, SetSnapshotRefUpdate):
+                refs[update.ref_name] = SnapshotRef(
+                    snapshot_id=update.snapshot_id,
+                    snapshot_ref_type=update.type,
+                    max_ref_age_ms=update.max_ref_age_ms,
+                    max_snapshot_age_ms=update.max_snapshot_age_ms,
+                    min_snapshots_to_keep=update.min_snapshots_to_keep,
+                )
+            elif isinstance(update, RemoveSnapshotRefUpdate):
+                refs.pop(update.ref_name, None)
+        return refs
+
     def _remove_ref_snapshot(self, ref_name: str) -> ManageSnapshots:
         """Remove a snapshot ref.
 
@@ -1232,6 +1267,61 @@ def _current_ancestors(self) -> set[int]:
             )
         }
 
+    def fast_forward_branch(self, from_branch: str, to_ref: str) -> ManageSnapshots:
+        """Fast-forward ``from_branch`` to the snapshot referenced by ``to_ref``.
+
+        * If ``from_branch`` does not exist, it is created pointing at ``to_ref``'s snapshot (Java/Spark parity).
+        * If both refs already point to the same snapshot the call is a no-op.
+        * Otherwise ``from_branch`` must be a branch (not a tag) and its current snapshot
+        must be an ancestor of ``to_ref``'s snapshot.
+
+        Within a single ``manage_snapshots()`` chain, ref lookups observe earlier staged
+        operations via :meth:`_effective_refs`. This means that `create_branch(...)` followed
+        by `fast_forward_branch(...)` on the same ref works as expected.
+
+        Args:
+            from_branch: name of the branch to advance.
+            to_ref: name of the branch or tag whose snapshot ``from_branch`` will point to.
+
+        Returns:
+            This for method chaining.
+
+        Raises:
+            NoSuchSnapshotRefError: ``to_ref`` does not exist.
+            SnapshotRefTypeError: ``from_branch`` exists but is a tag.
+            NotAncestorError: ``from_branch``'s snapshot is not an ancestor of ``to_ref``'s snapshot.
+        """
+        refs = self._effective_refs()
+
+        if (to_snapshot_ref := refs.get(to_ref)) is None:
+            raise NoSuchSnapshotRefError(f"Ref does not exist: {to_ref}")
+        to_snapshot_id = to_snapshot_ref.snapshot_id
+
+        if from_branch not in refs:
+            return self.create_branch(snapshot_id=to_snapshot_id, branch_name=from_branch)
+
+        from_ref = refs[from_branch]
+        if from_ref.snapshot_ref_type != SnapshotRefType.BRANCH:
+            raise SnapshotRefTypeError(f"Ref {from_branch} is a tag, not a branch")
+
+        if from_ref.snapshot_id == to_snapshot_id:
+            return self
+
+        if not is_ancestor_of(to_snapshot_id, from_ref.snapshot_id, self._transaction.table_metadata):
+            raise NotAncestorError(f"Cannot fast-forward: {from_branch} is not an ancestor of {to_ref}")
+
+        update, requirement = self._transaction._set_ref_snapshot(
+            snapshot_id=to_snapshot_id,
+            ref_name=from_branch,
+            type=SnapshotRefType.BRANCH,
+            max_ref_age_ms=from_ref.max_ref_age_ms,
+            max_snapshot_age_ms=from_ref.max_snapshot_age_ms,
+            min_snapshots_to_keep=from_ref.min_snapshots_to_keep,
+        )
+        self._updates += update
+        self._requirements += requirement
+        return self
+
 
 class ExpireSnapshots(UpdateTableMetadata["ExpireSnapshots"]):
     """Expire snapshots by ID.

tests/integration/test_snapshot_operations.py

@@ -282,6 +282,80 @@ def test_rollback_to_timestamp_no_valid_snapshot(table_with_snapshots: Table) ->
         table_with_snapshots.manage_snapshots().rollback_to_timestamp(timestamp_ms=oldest_timestamp).commit()
 
 
+@pytest.mark.integration
+@pytest.mark.parametrize("catalog", [lf("session_catalog_hive"), lf("session_catalog")])
+def test_fast_forward_branch(catalog: Catalog) -> None:
+    identifier = "default.test_table_snapshot_operations"
+    tbl = catalog.load_table(identifier)
+    assert len(tbl.history()) > 2
+
+    # Create a side branch off an older snapshot on main, then append to it
+    # so that the side branch is a strict descendant of main's older snapshot
+    # but the *current* main is not yet caught up.
+    current_snapshot = tbl.current_snapshot()
+    assert current_snapshot is not None
+    main_snapshot_id = current_snapshot.snapshot_id
+    side_branch = "audit_ff"
+
+    tbl.manage_snapshots().create_branch(snapshot_id=main_snapshot_id, branch_name=side_branch).commit()
+
+    arrow_schema = tbl.schema().as_arrow()
+    new_rows = pa.Table.from_pylist([{col.name: None for col in arrow_schema}], schema=arrow_schema)
+    tbl.append(new_rows, branch=side_branch)
+
+    # Validate appending to the side branch has advanced its snapshot.
+    tbl = catalog.load_table(identifier)
+    audit_snapshot_id = tbl.refs()[side_branch].snapshot_id
+    assert audit_snapshot_id != main_snapshot_id, "side-branch append should have advanced it"
+
+    # Fast-forward main to the side branch.
+    tbl.manage_snapshots().fast_forward_branch(from_branch="main", to_ref=side_branch).commit()
+
+    tbl = catalog.load_table(identifier)
+    assert tbl.refs()["main"].snapshot_id == audit_snapshot_id
+
+
+@pytest.mark.integration
+@pytest.mark.parametrize("catalog", [lf("session_catalog_hive"), lf("session_catalog")])
+def test_fast_forward_branch_preserves_retention_intra_chain(catalog: Catalog) -> None:
+    identifier = "default.test_table_snapshot_operations"
+    tbl = catalog.load_table(identifier)
+    assert len(tbl.history()) > 2
+
+    # Pick an older snapshot as the branch's starting point so the subsequent
+    # fast-forward to main is a real advance, not a no-op.
+    older_snapshot_id = tbl.history()[-3].snapshot_id
+    current_snapshot = tbl.current_snapshot()
+    assert current_snapshot is not None
+    main_snapshot_id = current_snapshot.snapshot_id
+
+    branch_name = "retention_intra_chain"
+    max_ref = 3_600_000  # 1h — distinct value per field so a swap would be caught
+    max_snap = 7_200_000  # 2h
+    min_keep = 5
+
+    # Chain create_branch (with retention) + fast_forward_branch in one commit.
+    # _effective_refs lets the fast-forward observe the same-chain create and
+    # carry the retention fields onto the second staged SetSnapshotRefUpdate.
+    tbl.manage_snapshots().create_branch(
+        snapshot_id=older_snapshot_id,
+        branch_name=branch_name,
+        max_ref_age_ms=max_ref,
+        max_snapshot_age_ms=max_snap,
+        min_snapshots_to_keep=min_keep,
+    ).fast_forward_branch(
+        from_branch=branch_name,
+        to_ref="main",
+    ).commit()
+
+    tbl = catalog.load_table(identifier)
+    ref = tbl.refs()[branch_name]
+    assert ref.snapshot_id == main_snapshot_id
+    assert ref.max_ref_age_ms == max_ref
+    assert ref.max_snapshot_age_ms == max_snap
+    assert ref.min_snapshots_to_keep == min_keep
+
+
 @pytest.mark.integration
 def test_rollback_to_timestamp(table_with_snapshots: Table) -> None:
     current_snapshot = table_with_snapshots.current_snapshot()