print stmt

Author: gabeiglio

pyiceberg/table/update/snapshot.py

@@ -17,6 +17,7 @@
 from __future__ import annotations
 
 import itertools
+import os
 import uuid
 from abc import abstractmethod
 from collections import defaultdict
@@ -591,88 +592,77 @@ class _OverwriteFiles(_MergingSnapshotProducer["_OverwriteFiles"]):
     Data and delete files were added and removed in a logical overwrite operation.
     """
 
-    def _existing_manifests(self) -> list[ManifestFile]:
-        """Determine if there are any existing manifest files."""
-        existing_files = []
+    @cached_property
+    def _cache_manifests(self) -> tuple[list[ManifestFile], list[ManifestEntry]]:
+        """Iterate over manifests once, returning (existing_manifests, deleted_entries).
 
-        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):
-                # Manifest does not contain rows that match the files to delete partitions
-                if not manifest_evaluators[manifest_file.partition_spec_id](manifest_file):
-                    existing_files.append(manifest_file)
-                    continue
-
-                entries_to_write: set[ManifestEntry] = set()
-                found_deleted_entries: set[ManifestEntry] = set()
-
-                for entry in manifest_file.fetch_manifest_entry(io=self._io, discard_deleted=True):
-                    if entry.data_file in self._deleted_data_files:
-                        found_deleted_entries.add(entry)
-                    else:
-                        entries_to_write.add(entry)
+        Avoids the double manifest read that occurs when _existing_manifests() and
+        _deleted_entries() are called independently from separate futures.
+        """
+        print(f"CACHE MISS: executing _cache_manifests PID:{os.getpid()}")
 
-                # Is the intercept the empty set?
-                if len(found_deleted_entries) == 0:
-                    existing_files.append(manifest_file)
-                    continue
+        if not (snapshot := self._transaction.table_metadata.snapshot_by_name(name=self._target_branch)):
+            return [], []
 
-                # Delete all files from manifest
-                if len(entries_to_write) == 0:
-                    continue
+        manifest_evaluators: dict[int, Callable[[ManifestFile], bool]] = KeyDefaultDict(self._build_manifest_evaluator)
 
-                # We have to rewrite the manifest file without the deleted data files
-                with self.new_manifest_writer(self.spec(manifest_file.partition_spec_id)) as writer:
-                    for entry in entries_to_write:
-                        writer.add_entry(
-                            ManifestEntry.from_args(
-                                status=ManifestEntryStatus.EXISTING,
-                                snapshot_id=entry.snapshot_id,
-                                sequence_number=entry.sequence_number,
-                                file_sequence_number=entry.file_sequence_number,
-                                data_file=entry.data_file,
-                            )
+        def _process_manifest(manifest: ManifestFile) -> tuple[ManifestFile | None, list[ManifestEntry]]:
+            if not manifest_evaluators[manifest.partition_spec_id](manifest):
+                return manifest, []  # unaffected partition — keep as-is, no deleted entries
+
+            surviving: list[ManifestEntry] = []
+            deleted: list[ManifestEntry] = []
+
+            for entry in manifest.fetch_manifest_entry(io=self._io, discard_deleted=True):
+                if entry.data_file.content == DataFileContent.DATA and entry.data_file in self._deleted_data_files:
+                    deleted.append(
+                        ManifestEntry.from_args(
+                            status=ManifestEntryStatus.DELETED,
+                            snapshot_id=entry.snapshot_id,
+                            sequence_number=entry.sequence_number,
+                            file_sequence_number=entry.file_sequence_number,
+                            data_file=entry.data_file,
                         )
-                existing_files.append(writer.to_manifest_file())
+                    )
+                else:
+                    surviving.append(entry)
+
+            if not deleted:
+                return manifest, []  # nothing deleted in this manifest
+
+            if not surviving:
+                return None, deleted  # all entries deleted — drop manifest
+
+            # Partial deletion — rewrite manifest with surviving entries only
+            with self.new_manifest_writer(self.spec(manifest.partition_spec_id)) as writer:
+                for entry in surviving:
+                    writer.add_entry(
+                        ManifestEntry.from_args(
+                            status=ManifestEntryStatus.EXISTING,
+                            snapshot_id=entry.snapshot_id,
+                            sequence_number=entry.sequence_number,
+                            file_sequence_number=entry.file_sequence_number,
+                            data_file=entry.data_file,
+                        )
+                    )
+            return writer.to_manifest_file(), deleted
 
-        return existing_files
+        executor = ExecutorFactory.get_or_create()
+        results = list(executor.map(_process_manifest, snapshot.manifests(io=self._io)))
 
-    def _deleted_entries(self) -> list[ManifestEntry]:
-        """To determine if we need to record any deleted entries.
+        existing = [m for m, _ in results if m is not None]
+        deleted_entries = [e for _, entries in results for e in entries]
+        return existing, deleted_entries
 
-        With a full overwrite all the entries are considered deleted.
-        With partial overwrites we have to use the predicate to evaluate
-        which entries are affected.
-        """
-        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:
-                # This should never happen since you cannot overwrite an empty table
-                raise ValueError(f"Could not find the previous snapshot: {self._parent_snapshot_id}")
-
-            executor = ExecutorFactory.get_or_create()
-            manifest_evaluators: dict[int, Callable[[ManifestFile], bool]] = KeyDefaultDict(self._build_manifest_evaluator)
-
-            def _get_entries(manifest: ManifestFile) -> list[ManifestEntry]:
-                if not manifest_evaluators[manifest.partition_spec_id](manifest):
-                    return []
-
-                return [
-                    ManifestEntry.from_args(
-                        status=ManifestEntryStatus.DELETED,
-                        snapshot_id=entry.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
-                ]
+    def _existing_manifests(self) -> list[ManifestFile]:
+        hit = "_cache_manifests" in self.__dict__
+        print(f"[_existing_manifests] cache {'HIT' if hit else 'MISS'}")
+        return self._cache_manifests[0]
 
-            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]:
+        hit = "_cache_manifests" in self.__dict__
+        print(f"[_deleted_entries] cache {'HIT' if hit else 'MISS'}")
+        return self._cache_manifests[1]
 
 
 class UpdateSnapshot: