diff --git a/pyiceberg/table/maintenance.py b/pyiceberg/table/maintenance.py index 0fcda35ae9..d0f4e0d0ea 100644 --- a/pyiceberg/table/maintenance.py +++ b/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) diff --git a/pyiceberg/table/snapshots.py b/pyiceberg/table/snapshots.py index 68dfcc868d..c44b5b58d4 100644 --- a/pyiceberg/table/snapshots.py +++ b/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: diff --git a/pyiceberg/table/update/snapshot.py b/pyiceberg/table/update/snapshot.py index 7931edacdd..cdeb6fe986 100644 --- a/pyiceberg/table/update/snapshot.py +++ b/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 diff --git a/tests/integration/test_rewrite_manifests_interop.py b/tests/integration/test_rewrite_manifests_interop.py new file mode 100644 index 0000000000..bb44c3bfd8 --- /dev/null +++ b/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)) diff --git a/tests/table/test_rewrite_manifests.py b/tests/table/test_rewrite_manifests.py new file mode 100644 index 0000000000..2c1f841c37 --- /dev/null +++ b/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 diff --git a/tests/table/test_snapshots.py b/tests/table/test_snapshots.py index 5f1680ed59..9e23669cd0 100644 --- a/tests/table/test_snapshots.py +++ b/tests/table/test_snapshots.py @@ -399,10 +399,23 @@ def test_merge_snapshot_summaries_overwrite_summary() -> None: assert actual.additional_properties == expected -def test_invalid_operation() -> None: - with pytest.raises(ValueError) as e: - update_snapshot_summaries(summary=Summary(Operation.REPLACE)) - assert "Operation not implemented: Operation.REPLACE" in str(e.value) +def test_replace_operation_carries_totals() -> None: + actual = update_snapshot_summaries( + summary=Summary(Operation.REPLACE), + previous_summary={ + "total-data-files": "3", + "total-delete-files": "0", + "total-records": "9", + "total-files-size": "1234", + "total-position-deletes": "0", + "total-equality-deletes": "0", + }, + ) + + # a replace operation does not change any of the totals + assert actual["total-data-files"] == "3" + assert actual["total-records"] == "9" + assert actual["total-files-size"] == "1234" def test_invalid_type() -> None: