Replan the manifest rewrite on commit retry and align rewrites_needed with it

Author: moomindani

mkdocs/docs/api.md

@@ -1535,7 +1535,7 @@ Rewrite the current snapshot's data manifests without changing any data. Live en
 table.maintenance.rewrite_manifests().commit()
 ```
 
-Nothing is committed when there is only one data manifest, so the call is safe to schedule periodically. To find out in advance whether a rewrite would do anything:
+Nothing is committed when no manifests would be merged, so the call is safe to schedule periodically. To find out in advance whether a rewrite would do anything:
 
 ```python
 if table.maintenance.rewrite_manifests().rewrites_needed():

pyiceberg/table/update/snapshot.py

@@ -1452,6 +1452,18 @@ def _existing_manifests(self) -> list[ManifestFile]:
         self._computed_manifests = new_manifests + kept_manifests
         return self._computed_manifests
 
+    def _refresh_for_retry(self) -> None:
+        """Reset state for a retry attempt, discarding the plan built from the stale branch head."""
+        super()._refresh_for_retry()
+        # The plan and its counters depend on the branch head, which changes on retry. Keeping them
+        # would rebuild the replacement snapshot from the manifests of the superseded snapshot,
+        # dropping whatever was committed concurrently.
+        self._computed_manifests = None
+        self._rewritten_count = 0
+        self._created_count = 0
+        self._kept_count = 0
+        self._entries_processed = 0
+
     def _commit(self) -> UpdatesAndRequirements:
         self._existing_manifests()
         if self._created_count == 0:
@@ -1460,11 +1472,25 @@ def _commit(self) -> UpdatesAndRequirements:
         return super()._commit()
 
     def rewrites_needed(self) -> bool:
-        """Return whether the current snapshot has data manifests to rewrite."""
+        """Return whether committing would rewrite any manifest.
+
+        This mirrors the grouping that `_existing_manifests` performs, so a snapshot whose data
+        manifests all end up kept as-is reports False. A predicate that matches only manifests
+        holding no live entries still reports True, because deciding that requires reading them.
+        """
         snapshot = self._transaction.table_metadata.snapshot_by_name(self._target_branch or MAIN_BRANCH)
         if snapshot is None:
             return False
-        data_manifests = [m for m in snapshot.manifests(self._io) if m.content == ManifestContent.DATA]
-        if self._manifest_predicate is not None:
-            return any(self._manifest_predicate(m) for m in data_manifests)
-        return len(data_manifests) > 1
+
+        data_manifests_by_spec: defaultdict[int, list[ManifestFile]] = defaultdict(list)
+        for manifest in snapshot.manifests(self._io):
+            if manifest.content == ManifestContent.DATA and (
+                self._manifest_predicate is None or self._manifest_predicate(manifest)
+            ):
+                data_manifests_by_spec[manifest.partition_spec_id].append(manifest)
+
+        return any(
+            len(group) > 1 or self._manifest_predicate is not None
+            for manifests in data_manifests_by_spec.values()
+            for group in self._group_by_target_size(manifests)
+        )

tests/table/test_commit_retry.py

@@ -302,6 +302,40 @@ def test_delete_files_refresh_clears_compute_deletes_cache(catalog: Catalog) ->
     assert "_compute_deletes" not in producer.__dict__
 
 
+def test_rewrite_manifests_refresh_clears_computed_manifests(catalog: Catalog) -> None:
+    """Verify that _refresh_for_retry discards the plan RewriteManifests computed for the old head."""
+    catalog.create_namespace("default")
+    schema = _test_schema()
+    table = catalog.create_table("default.rewrite_cache_test", schema=schema)
+
+    import pyarrow as pa
+
+    table.append(pa.table({"x": [1, 2, 3]}))
+    table.append(pa.table({"x": [4, 5, 6]}))
+    table = catalog.load_table("default.rewrite_cache_test")
+
+    from pyiceberg.table.update.snapshot import RewriteManifests
+
+    tx = Transaction(table, autocommit=False)
+    producer = RewriteManifests(transaction=tx, io=table.io)
+
+    # Plan against the current head, populating the cache and the counters
+    planned = producer._existing_manifests()
+
+    assert planned
+    assert producer._created_count == 1
+    assert producer._rewritten_count == 2
+
+    producer._refresh_for_retry()
+
+    # Reusing either would rebuild the replacement snapshot from the superseded head
+    assert producer._computed_manifests is None
+    assert producer._created_count == 0
+    assert producer._rewritten_count == 0
+    assert producer._kept_count == 0
+    assert producer._entries_processed == 0
+
+
 def test_concurrent_overwrite_overwrite_raises_validation_exception(catalog: Catalog) -> None:
     """Concurrent overwrites on the same data should fail with ValidationException."""
     catalog.create_namespace("default")

tests/table/test_rewrite_manifests.py

@@ -284,3 +284,22 @@ def test_rewrite_manifests_merges_live_and_fully_deleted_manifests(catalog: Cata
 
     # Verify table data is fully preserved and matches remaining live records
     assert table.scan().to_arrow().sort_by("id") == _arrow_table(offset=3)
+
+
+def test_rewrites_needed_is_false_when_every_group_is_a_single_manifest(catalog: Catalog) -> None:
+    table = _create_table_with_appends(catalog, appends=3)
+    # A tiny target size puts every manifest in its own group, so nothing can be merged
+    with table.transaction() as tx:
+        tx.set_properties({"commit.manifest.target-size-bytes": "1"})
+
+    table = catalog.load_table("default.test_rewrite")
+    assert len(_data_manifests(table)) == 3
+
+    snapshot_before = table.current_snapshot()
+    assert table.maintenance.rewrite_manifests().rewrites_needed() is False
+
+    table.maintenance.rewrite_manifests().commit()
+
+    table = catalog.load_table("default.test_rewrite")
+    assert table.current_snapshot() == snapshot_before
+    assert len(_data_manifests(table)) == 3