qzyu999 commented on code in PR #3131:
URL: https://github.com/apache/iceberg-python/pull/3131#discussion_r3089966996
##########
pyiceberg/table/update/snapshot.py:
##########
@@ -667,6 +667,96 @@ 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 (), ()
+
+ def _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 _existing_manifests(self) -> list[ManifestFile]:
+ """To determine if there are any existing manifests."""
Review Comment:
Hi @geruh, thank you for the feedback! I agree with this idea that the two
sets of code are highly similar, in b0a770c I create a
`_get_existing_manifests` function in `_SnapshotProducer` that is reused in
both `_OverwriteFiles` and `_RewriteFiles`.
--
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]
