Support rewrite_manifests table maintenance

pyiceberg/table/maintenance.py

@@ -24,7 +24,7 @@
 
 if TYPE_CHECKING:
     from pyiceberg.table import Table
-    from pyiceberg.table.update.snapshot import ExpireSnapshots
+    from pyiceberg.table.update.snapshot import ExpireSnapshots, RewriteManifests
 
 
 class MaintenanceTable:
@@ -43,3 +43,17 @@ def expire_snapshots(self) -> ExpireSnapshots:
         from pyiceberg.table.update.snapshot import ExpireSnapshots
 
         return ExpireSnapshots(transaction=Transaction(self.tbl, autocommit=True))
+
+    def rewrite_manifests(self) -> RewriteManifests:
+        """Return a RewriteManifests operation that merges the current snapshot's data manifests.
+
+        Entries are rewritten as EXISTING, keeping their sequence numbers; delete
+        manifests are kept as-is. The result is committed as a `replace` snapshot.
+
+        Returns:
+            RewriteManifests operation; call commit() to execute it.
+        """
+        from pyiceberg.table import Transaction
+        from pyiceberg.table.update.snapshot import RewriteManifests
+
+        return RewriteManifests(transaction=Transaction(self.tbl, autocommit=True), io=self.tbl.io)

pyiceberg/table/snapshots.py

@@ -344,7 +344,7 @@ def _partition_summary(self, update_metrics: UpdateMetrics) -> str:
 
 
 def update_snapshot_summaries(summary: Summary, previous_summary: Mapping[str, str] | None = None) -> Summary:
-    if summary.operation not in {Operation.APPEND, Operation.OVERWRITE, Operation.DELETE}:
+    if summary.operation not in {Operation.APPEND, Operation.OVERWRITE, Operation.DELETE, Operation.REPLACE}:
         raise ValueError(f"Operation not implemented: {summary.operation}")
 
     if not previous_summary:

pyiceberg/table/update/snapshot.py

@@ -1131,3 +1131,131 @@ def older_than(self, dt: datetime) -> ExpireSnapshots:
             if snapshot.timestamp_ms < expire_from and snapshot.snapshot_id not in protected_ids:
                 self._snapshot_ids_to_expire.add(snapshot.snapshot_id)
         return self
+
+
+class RewriteManifests(_SnapshotProducer["RewriteManifests"]):
+    """Rewrite the current snapshot's data manifests without changing data.
+
+    Live entries from the rewritten data manifests are regrouped into new
+    manifests sized by `commit.manifest.target-size-bytes`, written as EXISTING
+    entries that keep their sequence numbers. Entries with status DELETED are
+    dropped, matching the reference implementation, which rewrites live entries
+    only. Delete manifests are kept as-is. The result is committed as a
+    `replace` snapshot; if no manifests need merging, no snapshot is committed.
+    """
+
+    _computed_manifests: list[ManifestFile] | None
+
+    _rewritten_count: int
+    _created_count: int
+    _kept_count: int
+    _entries_processed: int
+
+    def __init__(
+        self,
+        transaction: Transaction,
+        io: FileIO,
+        commit_uuid: uuid.UUID | None = None,
+        snapshot_properties: dict[str, str] = EMPTY_DICT,
+        branch: str | None = MAIN_BRANCH,
+    ) -> None:
+        super().__init__(Operation.REPLACE, transaction, io, commit_uuid, snapshot_properties, branch)
+        if transaction.table_metadata.format_version >= 3:
+            raise NotImplementedError(
+                "Rewriting manifests is not yet supported for V3 tables: "
+                "the first-row-id of rewritten manifests must be preserved, "
+                "see: https://github.com/apache/iceberg-python/issues/3621"
+            )
+        self._rewritten_count = 0
+        self._created_count = 0
+        self._kept_count = 0
+        self._entries_processed = 0
+        self._computed_manifests = None
+
+    def _deleted_entries(self) -> list[ManifestEntry]:
+        return []
+
+    def _target_size_bytes(self) -> int:
+        from pyiceberg.table import TableProperties
+
+        return property_as_int(  # type: ignore
+            self._transaction.table_metadata.properties,
+            TableProperties.MANIFEST_TARGET_SIZE_BYTES,
+            TableProperties.MANIFEST_TARGET_SIZE_BYTES_DEFAULT,
+        )
+
+    def _group_by_target_size(self, manifests: list[ManifestFile]) -> list[list[ManifestFile]]:
+        """Pack manifests into groups whose source sizes add up to roughly the target size."""
+        target_size = self._target_size_bytes()
+        groups: list[list[ManifestFile]] = []
+        current_group: list[ManifestFile] = []
+        current_size = 0
+        for manifest in manifests:
+            if current_group and current_size + manifest.manifest_length > target_size:
+                groups.append(current_group)
+                current_group = []
+                current_size = 0
+            current_group.append(manifest)
+            current_size += manifest.manifest_length
+        if current_group:
+            groups.append(current_group)
+        return groups
+
+    def _existing_manifests(self) -> list[ManifestFile]:
+        if self._computed_manifests is not None:
+            return self._computed_manifests
+
+        snapshot = self._transaction.table_metadata.snapshot_by_name(self._target_branch or MAIN_BRANCH)
+        if snapshot is None:
+            self._computed_manifests = []
+            return self._computed_manifests
+
+        data_manifests_by_spec: defaultdict[int, list[ManifestFile]] = defaultdict(list)
+        kept_manifests: list[ManifestFile] = []
+        for manifest in snapshot.manifests(self._io):
+            if manifest.content == ManifestContent.DATA:
+                data_manifests_by_spec[manifest.partition_spec_id].append(manifest)
+            else:
+                kept_manifests.append(manifest)
+
+        new_manifests: list[ManifestFile] = []
+        for spec_id, manifests in data_manifests_by_spec.items():
+            for group in self._group_by_target_size(manifests):
+                if len(group) == 1:
+                    # nothing to merge; keep the manifest as-is
+                    kept_manifests.append(group[0])
+                    continue
+                with self.new_manifest_writer(self.spec(spec_id)) as writer:
+                    for manifest in group:
+                        for entry in manifest.fetch_manifest_entry(self._io, discard_deleted=True):
+                            writer.existing(entry)
+                            self._entries_processed += 1
+                new_manifests.append(writer.to_manifest_file())
+                self._rewritten_count += len(group)
+                self._created_count += 1
+
+        self._kept_count = len(kept_manifests)
+        self.snapshot_properties = {
+            **self.snapshot_properties,
+            "manifests-created": str(self._created_count),
+            "manifests-kept": str(self._kept_count),
+            "manifests-replaced": str(self._rewritten_count),
+            "entries-processed": str(self._entries_processed),
+        }
+        self._computed_manifests = new_manifests + kept_manifests
+        return self._computed_manifests
+
+    def _commit(self) -> UpdatesAndRequirements:
+        self._existing_manifests()
+        if self._created_count == 0:
+            # nothing was merged; committing would only produce a pointless replace snapshot
+            return (), ()
+        return super()._commit()
+
+    def rewrites_needed(self) -> bool:
+        """Return whether the current snapshot has more than one data manifest to merge."""
+        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]
+        return len(data_manifests) > 1

tests/integration/test_rewrite_manifests_interop.py

@@ -0,0 +1,74 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint:disable=redefined-outer-name
+from typing import TYPE_CHECKING
+
+import pyarrow as pa
+import pytest
+
+from pyiceberg.catalog import Catalog
+from pyiceberg.exceptions import NoSuchTableError
+from pyiceberg.manifest import ManifestContent
+
+if TYPE_CHECKING:
+    from pyspark.sql import SparkSession
+
+
+@pytest.mark.integration
+def test_spark_reads_table_after_rewrite_manifests(session_catalog: Catalog, spark: "SparkSession") -> None:
+    identifier = "default.test_rewrite_manifests_interop"
+    try:
+        session_catalog.drop_table(identifier)
+    except NoSuchTableError:
+        pass
+
+    table = session_catalog.create_table(identifier, schema=pa.schema([pa.field("id", pa.int64())]))
+    for i in range(3):
+        table.append(pa.table({"id": pa.array([i * 3 + 1, i * 3 + 2, i * 3 + 3], type=pa.int64())}))
+
+    table = session_catalog.load_table(identifier)
+    snapshot = table.current_snapshot()
+    assert snapshot is not None
+    assert len([m for m in snapshot.manifests(table.io) if m.content == ManifestContent.DATA]) == 3
+
+    table.maintenance.rewrite_manifests().commit()
+
+    table = session_catalog.load_table(identifier)
+    snapshot = table.current_snapshot()
+    assert snapshot is not None
+    assert len([m for m in snapshot.manifests(table.io) if m.content == ManifestContent.DATA]) == 1
+
+    # Spark must read the rewritten table with the same data
+    spark_rows = spark.table(f"{identifier}").collect()
+    assert sorted(row.id for row in spark_rows) == list(range(1, 10))
+
+    # Spark must see the replace snapshot and the preserved data files
+    snapshots = spark.sql(f"SELECT operation FROM {identifier}.snapshots ORDER BY committed_at").collect()
+    assert [row.operation for row in snapshots] == ["append", "append", "append", "replace"]
+
+    files = spark.sql(f"SELECT file_path FROM {identifier}.files").collect()
+    assert len(files) == 3
+
+    # Spark sees the same manifest consolidation
+    manifests = spark.sql(f"SELECT path FROM {identifier}.manifests").collect()
+    assert len(manifests) == 1
+
+    # time travel to the pre-rewrite snapshot still works from Spark
+    previous_snapshot_id = snapshot.parent_snapshot_id
+    assert previous_snapshot_id is not None
+    previous_rows = spark.sql(f"SELECT id FROM {identifier} VERSION AS OF {previous_snapshot_id}").collect()
+    assert sorted(row.id for row in previous_rows) == list(range(1, 10))

tests/table/test_rewrite_manifests.py

@@ -0,0 +1,143 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from pathlib import Path
+
+import pyarrow as pa
+import pytest
+
+from pyiceberg.catalog import Catalog
+from pyiceberg.catalog.memory import InMemoryCatalog
+from pyiceberg.manifest import ManifestContent, ManifestFile
+from pyiceberg.table import Table
+from pyiceberg.table.snapshots import Operation
+
+
+@pytest.fixture
+def catalog(tmp_path: Path) -> Catalog:
+    catalog = InMemoryCatalog("test.rewrite_manifests", warehouse=f"file://{tmp_path}")
+    catalog.create_namespace("default")
+    return catalog
+
+
+def _arrow_table(offset: int = 0) -> pa.Table:
+    return pa.table({"id": pa.array([offset + 1, offset + 2, offset + 3], type=pa.int64())})
+
+
+def _create_table_with_appends(catalog: Catalog, appends: int = 3) -> Table:
+    table = catalog.create_table("default.test_rewrite", schema=pa.schema([pa.field("id", pa.int64())]))
+    for i in range(appends):
+        table.append(_arrow_table(offset=i * 3))
+    return table
+
+
+def _data_manifests(table: Table) -> list[ManifestFile]:
+    snapshot = table.current_snapshot()
+    assert snapshot is not None
+    return [m for m in snapshot.manifests(table.io) if m.content == ManifestContent.DATA]
+
+
+def test_rewrite_manifests_merges_data_manifests(catalog: Catalog) -> None:
+    table = _create_table_with_appends(catalog, appends=3)
+    assert len(_data_manifests(table)) == 3
+    rows_before = table.scan().to_arrow().sort_by("id")
+
+    table.maintenance.rewrite_manifests().commit()
+
+    table = catalog.load_table("default.test_rewrite")
+    manifests = _data_manifests(table)
+    assert len(manifests) == 1
+    # entries are rewritten as EXISTING
+    assert manifests[0].existing_files_count == 3
+    assert manifests[0].added_files_count == 0
+
+    # data is unchanged
+    assert table.scan().to_arrow().sort_by("id") == rows_before
+
+    snapshot = table.current_snapshot()
+    assert snapshot is not None
+    assert snapshot.summary is not None
+    assert snapshot.summary.operation == Operation.REPLACE
+    assert snapshot.summary["manifests-created"] == "1"
+    assert snapshot.summary["manifests-replaced"] == "3"
+    assert snapshot.summary["entries-processed"] == "3"
+    # totals carry over unchanged
+    assert snapshot.summary["total-data-files"] == "3"
+    assert snapshot.summary["total-records"] == "9"
+
+
+def _sequence_numbers_by_file(table: Table) -> dict[str, int]:
+    result: dict[str, int] = {}
+    for manifest in _data_manifests(table):
+        for entry in manifest.fetch_manifest_entry(table.io, discard_deleted=True):
+            assert entry.sequence_number is not None
+            result[entry.data_file.file_path] = entry.sequence_number
+    return result
+
+
+def test_rewrite_manifests_preserves_sequence_numbers(catalog: Catalog) -> None:
+    table = _create_table_with_appends(catalog, appends=3)
+    entries_before = _sequence_numbers_by_file(table)
+
+    table.maintenance.rewrite_manifests().commit()
+
+    table = catalog.load_table("default.test_rewrite")
+    entries_after = _sequence_numbers_by_file(table)
+    assert entries_after == entries_before
+    # the merged manifest keeps the min sequence number of its entries
+    assert _data_manifests(table)[0].min_sequence_number == min(entries_before.values())
+
+
+def test_rewrite_manifests_single_manifest_is_noop(catalog: Catalog) -> None:
+    table = _create_table_with_appends(catalog, appends=1)
+    snapshot_before = table.current_snapshot()
+    assert snapshot_before is not None
+    manifest_path_before = _data_manifests(table)[0].manifest_path
+
+    table.maintenance.rewrite_manifests().commit()
+
+    table = catalog.load_table("default.test_rewrite")
+    # nothing to merge: no new snapshot is committed and the manifest is untouched
+    snapshot = table.current_snapshot()
+    assert snapshot is not None
+    assert snapshot.snapshot_id == snapshot_before.snapshot_id
+    assert _data_manifests(table)[0].manifest_path == manifest_path_before
+
+
+def test_rewrite_manifests_respects_target_size(catalog: Catalog) -> None:
+    table = _create_table_with_appends(catalog, appends=4)
+    max_manifest_length = max(m.manifest_length for m in _data_manifests(table))
+
+    # allow two source manifests per group (2x fits, 3x exceeds), robust to small size variations
+    with table.transaction() as tx:
+        tx.set_properties({"commit.manifest.target-size-bytes": str(int(max_manifest_length * 2.5))})
+
+    table = catalog.load_table("default.test_rewrite")
+    table.maintenance.rewrite_manifests().commit()
+
+    table = catalog.load_table("default.test_rewrite")
+    manifests = _data_manifests(table)
+    assert len(manifests) == 2
+    assert all(m.existing_files_count == 2 for m in manifests)
+
+
+def test_rewrites_needed(catalog: Catalog) -> None:
+    table = _create_table_with_appends(catalog, appends=1)
+    assert table.maintenance.rewrite_manifests().rewrites_needed() is False
+
+    table.append(_arrow_table(offset=3))
+    table = catalog.load_table("default.test_rewrite")
+    assert table.maintenance.rewrite_manifests().rewrites_needed() is True