rambleraptor commented on code in PR #3631:
URL: https://github.com/apache/iceberg-python/pull/3631#discussion_r4030669692
##########
pyiceberg/table/update/snapshot.py:
##########
@@ -1312,3 +1312,185 @@ 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
+ _manifest_predicate: Callable[[ManifestFile], bool] | 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._manifest_predicate = None
+ self._rewritten_count = 0
+ self._created_count = 0
+ self._kept_count = 0
+ self._entries_processed = 0
+ self._computed_manifests = None
+
+ def rewrite_if(self, predicate: Callable[[ManifestFile], bool]) ->
RewriteManifests:
+ """Filter which manifests should be rewritten.
+
+ Passing a predicate also disables the optimization that keeps
single-manifest
+ groups as-is, allowing single manifests to be rewritten when they
match the predicate.
+
+ Args:
+ predicate: A function that takes a ManifestFile and returns True
if it should be rewritten.
+
+ Returns:
+ This RewriteManifests instance for method chaining.
+ """
+ self._manifest_predicate = predicate
+ return self
+
+ 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]:
Review Comment:
We should check that the manifest rewrites preserve every active file. Java
does this. I think this would be the best place to that.
In general, we should validate as much as we can. If we write data, we
should have some run-time assumptions about how valid the data is.
##########
pyiceberg/table/update/snapshot.py:
##########
@@ -1312,3 +1312,185 @@ 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"]):
Review Comment:
I've got some concerns around concurrency.
Some of the classes override `SnapshotProducer._validate_concurrency` to
remove all validation. I think we want to do the same thing here.
##########
tests/table/test_rewrite_manifests.py:
##########
@@ -0,0 +1,305 @@
+# 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
+
+
[email protected]
+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:
Review Comment:
Can we add in a new test to talk about the concurrency issues I spoke about
earlier? I'll give you one for free:
```
def test_rewrite_manifests_replans_after_concurrent_append(catalog: Catalog)
-> None:
"""A concurrent append must not fail the rewrite: it replans against the
new head."""
catalog.create_namespace("default")
table = catalog.create_table("default.rewrite_concurrent",
schema=_test_schema())
import pyarrow as pa
for i in range(3):
table.append(pa.table({"x": [i]}))
# attempt rewrite
table = catalog.load_table("default.rewrite_concurrent")
rewrite = table.maintenance.rewrite_manifests()
# another writer commits between planning and committing the rewrite
catalog.load_table("default.rewrite_concurrent").append(pa.table({"x":
[99]}))
rewrite.commit()
table = catalog.load_table("default.rewrite_concurrent")
snapshot = table.current_snapshot()
# Do actual data assertions
assert snapshot is not None
assert snapshot.summary.operation == Operation.REPLACE
# the concurrently appended file is merged in, not dropped
assert sorted(row["x"] for row in table.scan().to_arrow().to_pylist())
== [0, 1, 2, 99]
assert len([m for m in snapshot.manifests(table.io) if m.content ==
ManifestContent.DATA]) == 1
assert snapshot.summary["manifests-replaced"] == "4"
```
##########
tests/table/test_rewrite_manifests.py:
##########
@@ -0,0 +1,305 @@
+# 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
+
+
[email protected]
+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
+
+
+def test_rewrite_manifests_with_predicate_selective(catalog: Catalog) -> None:
+ table = _create_table_with_appends(catalog, appends=3)
+ manifests_before = _data_manifests(table)
+ assert len(manifests_before) == 3
+
+ # 1. Extract all manifest paths and underlying data file paths before
rewrite
+ paths_before = [m.manifest_path for m in manifests_before]
+ target_path = paths_before[0]
+ kept_paths_before = set(paths_before[1:])
+ rows_before = table.scan().to_arrow().sort_by("id")
+ data_files_before = [
+ entry.data_file.file_path for m in manifests_before for entry in
m.fetch_manifest_entry(table.io, discard_deleted=True)
+ ]
+
+ # 2. Execute selective rewrite
+ table.maintenance.rewrite_manifests().rewrite_if(lambda m: m.manifest_path
== target_path).commit()
Review Comment:
I really love this rewrite_if syntax.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
