geruh commented on code in PR #3131:
URL: https://github.com/apache/iceberg-python/pull/3131#discussion_r3095389329
##########
pyiceberg/table/update/snapshot.py:
##########
@@ -724,6 +790,14 @@ def delete(self) -> _DeleteFiles:
snapshot_properties=self._snapshot_properties,
)
+ def replace(self) -> _RewriteFiles:
+ return _RewriteFiles(
+ operation=Operation.REPLACE,
+ transaction=self._transaction,
Review Comment:
I noticed that branch is missing here is there a reason for that?
##########
pyiceberg/table/update/snapshot.py:
##########
@@ -724,6 +814,14 @@ def delete(self) -> _DeleteFiles:
snapshot_properties=self._snapshot_properties,
)
+ def replace(self) -> _RewriteFiles:
+ return _RewriteFiles(
Review Comment:
Yeah there is a bit of a distinction here, since rewrite is basically the
rewrite of data files and replace is the logical change to your snapshot
metadata. My thinking is that the users in java today are used to interacting
with this api through:
```java
table.newRewrite()
.deleteFile(old)
.addFile(new)
.commit();
```
So someone coming from Java Iceberg will look for rewrite, not replace. But
ultimately maybe there is more of a history as to why the it follows this
naming convention im missing on.
WDYT @kevinjqliu?
##########
pyiceberg/table/update/snapshot.py:
##########
@@ -667,6 +669,70 @@ def _get_entries(manifest: ManifestFile) ->
list[ManifestEntry]:
return []
+class _RewriteFiles(_SnapshotProducer["_RewriteFiles"]):
+ """A snapshot producer that rewrites data files."""
+
+ def __init__(self, operation: Operation, transaction: Transaction, io:
FileIO, snapshot_properties: dict[str, str]):
+ super().__init__(operation, transaction, io,
snapshot_properties=snapshot_properties)
+
+ def _commit(self) -> UpdatesAndRequirements:
+ # Only produce a commit when there is something to rewrite
+ if self._deleted_data_files or self._added_data_files:
+ # Grab the entries that we actually found in the table's manifests
+ deleted_entries = self._deleted_entries()
+ found_deleted_files = {entry.data_file for entry in
deleted_entries}
+
+ # If the user asked to delete files that aren't in the table,
abort.
+ if len(found_deleted_files) != len(self._deleted_data_files):
+ raise ValueError("Cannot delete files that are not present in
the table")
+
+ added_records = sum(f.record_count for f in self._added_data_files)
+ deleted_records = sum(entry.data_file.record_count for entry in
deleted_entries)
+
+ if added_records > deleted_records:
+ raise ValueError(f"Invalid replace: records added
({added_records}) exceeds records removed ({deleted_records})")
+
+ return super()._commit()
+ else:
+ return (), ()
+
+ @cached_property
+ def _cached_deleted_entries(self) -> list[ManifestEntry]:
+ """Check if we need to mark the files as deleted."""
+ if self._parent_snapshot_id is not None:
+ previous_snapshot =
self._transaction.table_metadata.snapshot_by_id(self._parent_snapshot_id)
+ if previous_snapshot is None:
+ raise ValueError(f"Could not find the previous snapshot:
{self._parent_snapshot_id}")
+
+ executor = ExecutorFactory.get_or_create()
+
+ def _get_entries(manifest: ManifestFile) -> list[ManifestEntry]:
+ return [
+ ManifestEntry.from_args(
+ status=ManifestEntryStatus.DELETED,
+ snapshot_id=self.snapshot_id,
+ sequence_number=entry.sequence_number,
+ file_sequence_number=entry.file_sequence_number,
+ data_file=entry.data_file,
+ )
+ for entry in manifest.fetch_manifest_entry(self._io,
discard_deleted=True)
+ if entry.data_file.content == DataFileContent.DATA and
entry.data_file in self._deleted_data_files
+ ]
+
+ list_of_entries = executor.map(_get_entries,
previous_snapshot.manifests(self._io))
+ return list(itertools.chain(*list_of_entries))
+ else:
+ return []
+
+ def _deleted_entries(self) -> list[ManifestEntry]:
+ """Check if we need to mark the files as deleted."""
+ return self._cached_deleted_entries
+
+ def _existing_manifests(self) -> list[ManifestFile]:
+ """To determine if there are any existing manifests."""
+ return self._get_existing_manifests()
Review Comment:
nit: Looks like these doc strings were copy pasta'd over from the other
classes, and don't fit how they are used here. Either we can remove them or
change to fit their usage.
##########
pyiceberg/table/update/snapshot.py:
##########
@@ -165,6 +165,50 @@ def _calculate_added_rows(self, manifests:
list[ManifestFile]) -> int:
added_rows += manifest.added_rows_count
return added_rows
+ def _get_existing_manifests(self) -> list[ManifestFile]:
+ """Filter existing manifests and rewrite those containing deleted data
files."""
+ existing_files: list[ManifestFile] = []
+ # Use manifest pruning if a predicate is set (primarily for Overwrite)
+ manifest_evaluators: dict[int, Callable[[ManifestFile], bool]] =
KeyDefaultDict(self._build_manifest_evaluator)
+
+ if snapshot :=
self._transaction.table_metadata.snapshot_by_name(name=self._target_branch):
+ for manifest_file in snapshot.manifests(io=self._io):
+ # Skip pruning for rewrite operations unless we want to
optimize later
+ if self._operation == Operation.OVERWRITE and not
manifest_evaluators[manifest_file.partition_spec_id](
Review Comment:
I'm not too keen on the child classes information leaking into the base
here. How about we use an argument like `should_use_manifest_pruning` or
something along those lines.
##########
pyiceberg/table/update/snapshot.py:
##########
@@ -667,6 +669,70 @@ def _get_entries(manifest: ManifestFile) ->
list[ManifestEntry]:
return []
+class _RewriteFiles(_SnapshotProducer["_RewriteFiles"]):
+ """A snapshot producer that rewrites data files."""
+
+ def __init__(self, operation: Operation, transaction: Transaction, io:
FileIO, snapshot_properties: dict[str, str]):
+ super().__init__(operation, transaction, io,
snapshot_properties=snapshot_properties)
Review Comment:
nit: we can remove the constructor here
##########
tests/table/test_replace.py:
##########
@@ -0,0 +1,637 @@
+# 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 typing import cast
+
+import pytest
+
+from pyiceberg.catalog import Catalog
+from pyiceberg.manifest import (
+ DataFile,
+ DataFileContent,
+ FileFormat,
+ ManifestEntry,
+ ManifestEntryStatus,
+)
+from pyiceberg.schema import Schema
+from pyiceberg.table.snapshots import Operation, Snapshot, Summary
+from pyiceberg.typedef import Record
+
+
+def test_replace_internally(catalog: Catalog) -> None:
+ # Setup a basic table using the catalog fixture
+ catalog.create_namespace("default")
+ table = catalog.create_table(
+ identifier="default.test_replace",
+ schema=Schema(),
+ )
+
+ # 1. File we will delete
+ file_to_delete = DataFile.from_args(
Review Comment:
nit: extract this into a helper or see if one exists already.
##########
tests/table/test_snapshots.py:
##########
@@ -398,8 +398,8 @@ def test_merge_snapshot_summaries_overwrite_summary() ->
None:
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)
+
update_snapshot_summaries(summary=Summary.model_construct(operation="unknown_operation"))
Review Comment:
Makes sense if we have support :)
--
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]
