Revert

Author: rambleraptor

pyiceberg/table/__init__.py

@@ -390,19 +390,17 @@ def _set_ref_snapshot(
 
         return updates, requirements
 
-    def _build_partition_predicate(
-        self, partition_records: set[Record], spec: PartitionSpec, schema: Schema
-    ) -> BooleanExpression:
+    def _build_partition_predicate(self, partition_records: set[Record]) -> BooleanExpression:
         """Build a filter predicate matching any of the input partition records.
 
         Args:
             partition_records: A set of partition records to match
-            spec: An optional partition spec, if none then defaults to current
-            schema: An optional schema, if none then defaults to current
         Returns:
             A predicate matching any of the input partition records.
         """
-        partition_fields = [schema.find_field(field.source_id).name for field in spec.fields]
+        partition_spec = self.table_metadata.spec()
+        schema = self.table_metadata.schema()
+        partition_fields = [schema.find_field(field.source_id).name for field in partition_spec.fields]
         if not partition_records or not partition_fields:
             return AlwaysFalse()
 
@@ -622,9 +620,7 @@ def dynamic_partition_overwrite(
         )
 
         partitions_to_overwrite = {data_file.partition for data_file in data_files}
-        delete_filter = self._build_partition_predicate(
-            partition_records=partitions_to_overwrite, spec=self.table_metadata.spec(), schema=self.table_metadata.schema()
-        )
+        delete_filter = self._build_partition_predicate(partition_records=partitions_to_overwrite)
         self.delete(
             delete_filter=delete_filter,
             snapshot_properties=snapshot_properties,

pyiceberg/table/update/snapshot.py

@@ -51,7 +51,6 @@
     write_manifest_list,
 )
 from pyiceberg.partitioning import PartitionSpec
-from pyiceberg.schema import Schema
 from pyiceberg.table.refs import MAIN_BRANCH, SnapshotRefType
 from pyiceberg.table.snapshots import (
     Operation,
@@ -74,7 +73,7 @@
     UpdatesAndRequirements,
     UpdateTableMetadata,
 )
-from pyiceberg.typedef import EMPTY_DICT, KeyDefaultDict, Record
+from pyiceberg.typedef import EMPTY_DICT, KeyDefaultDict
 from pyiceberg.utils.bin_packing import ListPacker
 from pyiceberg.utils.concurrent import ExecutorFactory
 from pyiceberg.utils.datetime import datetime_to_millis
@@ -229,8 +228,13 @@ def _process_manifests(self, manifests: list[ManifestFile]) -> list[ManifestFile
     def _manifests(self) -> list[ManifestFile]:
         def _write_added_manifest() -> list[ManifestFile]:
             if self._added_data_files:
-                with self.new_manifest_writer(
+                with write_manifest(
+                    format_version=self._transaction.table_metadata.format_version,
                     spec=self._transaction.table_metadata.spec(),
+                    schema=self._transaction.table_metadata.schema(),
+                    output_file=self.new_manifest_output(),
+                    snapshot_id=self._snapshot_id,
+                    avro_compression=self._compression,
                 ) as writer:
                     for data_file in self._added_data_files:
                         writer.add(
@@ -255,17 +259,21 @@ def _write_delete_manifest() -> list[ManifestFile]:
                 for deleted_entry in deleted_entries:
                     partition_groups[deleted_entry.data_file.spec_id].append(deleted_entry)
                 for spec_id, entries in partition_groups.items():
-                    with self.new_manifest_writer(self.spec(spec_id)) as writer:
+                    with write_manifest(
+                        format_version=self._transaction.table_metadata.format_version,
+                        spec=self._transaction.table_metadata.specs()[spec_id],
+                        schema=self._transaction.table_metadata.schema(),
+                        output_file=self.new_manifest_output(),
+                        snapshot_id=self._snapshot_id,
+                        avro_compression=self._compression,
+                    ) as writer:
                         for entry in entries:
                             writer.add_entry(entry)
                     deleted_manifests.append(writer.to_manifest_file())
                 return deleted_manifests
             else:
                 return []
 
-        # Updates self._predicate with computed partition predicate for manifest pruning
-        self._build_delete_files_partition_predicate()
-
         executor = ExecutorFactory.get_or_create()
 
         added_manifests = executor.submit(_write_added_manifest)
@@ -385,17 +393,14 @@ def _commit(self) -> UpdatesAndRequirements:
     def snapshot_id(self) -> int:
         return self._snapshot_id
 
-    def schema(self) -> Schema:
-        return self._transaction.table_metadata.schema()
-
     def spec(self, spec_id: int) -> PartitionSpec:
         return self._transaction.table_metadata.specs()[spec_id]
 
     def new_manifest_writer(self, spec: PartitionSpec) -> ManifestWriter:
         return write_manifest(
             format_version=self._transaction.table_metadata.format_version,
             spec=spec,
-            schema=self.schema(),
+            schema=self._transaction.table_metadata.schema(),
             output_file=self.new_manifest_output(),
             snapshot_id=self._snapshot_id,
             avro_compression=self._compression,
@@ -506,35 +511,22 @@ def _validate_concurrency(self) -> None:
             )
 
     def _build_partition_projection(self, spec_id: int) -> BooleanExpression:
-        project = inclusive_projection(self.schema(), self.spec(spec_id), self._case_sensitive)
+        schema = self._transaction.table_metadata.schema()
+        project = inclusive_projection(schema, self.spec(spec_id), self._case_sensitive)
         return project(self._predicate)
 
     @cached_property
     def partition_filters(self) -> KeyDefaultDict[int, BooleanExpression]:
         return KeyDefaultDict(self._build_partition_projection)
 
     def _build_manifest_evaluator(self, spec_id: int) -> Callable[[ManifestFile], bool]:
-        return manifest_evaluator(self.spec(spec_id), self.schema(), self.partition_filters[spec_id], self._case_sensitive)
+        schema = self._transaction.table_metadata.schema()
+        return manifest_evaluator(self.spec(spec_id), schema, self.partition_filters[spec_id], self._case_sensitive)
 
     def delete_by_predicate(self, predicate: BooleanExpression, case_sensitive: bool = True) -> None:
         self._predicate = Or(self._predicate, predicate)
         self._case_sensitive = case_sensitive
 
-    def _build_delete_files_partition_predicate(self) -> None:
-        """Build BooleanExpression based on deleted data files partitions."""
-        partition_to_overwrite: dict[int, set[Record]] = {}
-        for data_file in self._deleted_data_files:
-            group = partition_to_overwrite.setdefault(data_file.spec_id, set())
-            group.add(data_file.partition)
-
-        for spec_id, partition_records in partition_to_overwrite.items():
-            self.delete_by_predicate(
-                self._transaction._build_partition_predicate(
-                    partition_records=partition_records, schema=self.schema(), spec=self.spec(spec_id)
-                ),
-                self._case_sensitive,
-            )
-
 
 class _DeleteFiles(_SnapshotProducer["_DeleteFiles"]):
     """Will delete manifest entries from the current snapshot based on the predicate.
@@ -637,7 +629,14 @@ def _copy_with_new_status(entry: ManifestEntry, status: ManifestEntryStatus) ->
 
                                 # Rewrite the manifest
                                 if len(existing_entries) > 0:
-                                    with self.new_manifest_writer(spec=self.spec(manifest_file.partition_spec_id)) as writer:
+                                    with write_manifest(
+                                        format_version=self._transaction.table_metadata.format_version,
+                                        spec=self._transaction.table_metadata.specs()[manifest_file.partition_spec_id],
+                                        schema=self._transaction.table_metadata.schema(),
+                                        output_file=self.new_manifest_output(),
+                                        snapshot_id=self._snapshot_id,
+                                        avro_compression=self._compression,
+                                    ) as writer:
                                         for existing_entry in existing_entries:
                                             writer.add_entry(existing_entry)
                                     existing_manifests.append(writer.to_manifest_file())
@@ -772,46 +771,36 @@ def _existing_manifests(self) -> list[ManifestFile]:
         """Determine if there are any existing manifest files."""
         existing_files = []
 
-        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()
+                entries = manifest_file.fetch_manifest_entry(io=self._io, discard_deleted=True)
+                found_deleted_data_files = [entry.data_file for entry in entries if entry.data_file in self._deleted_data_files]
 
-                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)
-
-                # Is the intercept the empty set?
-                if len(found_deleted_entries) == 0:
+                if len(found_deleted_data_files) == 0:
                     existing_files.append(manifest_file)
-                    continue
-
-                # Delete all files from manifest
-                if len(entries_to_write) == 0:
-                    continue
-
-                # 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,
-                            )
-                        )
-                existing_files.append(writer.to_manifest_file())
-
+                else:
+                    # We have to rewrite the manifest file without the deleted data files
+                    if any(entry.data_file not in found_deleted_data_files for entry in entries):
+                        with write_manifest(
+                            format_version=self._transaction.table_metadata.format_version,
+                            spec=self._transaction.table_metadata.specs()[manifest_file.partition_spec_id],
+                            schema=self._transaction.table_metadata.schema(),
+                            output_file=self.new_manifest_output(),
+                            snapshot_id=self._snapshot_id,
+                            avro_compression=self._compression,
+                        ) as writer:
+                            for entry in entries:
+                                if entry.data_file not in found_deleted_data_files:
+                                    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,
+                                        )
+                                    )
+                        existing_files.append(writer.to_manifest_file())
         return existing_files
 
     def _deleted_entries(self) -> list[ManifestEntry]:
@@ -828,12 +817,8 @@ def _deleted_entries(self) -> list[ManifestEntry]:
                 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,

tests/table/test_init.py

@@ -1982,11 +1982,7 @@ def test_check_uuid_passes_when_match(table_v2: Table, example_table_metadata_v2
 
 def test_build_large_partition_predicate(table_v2: Table) -> None:
     with table_v2.transaction() as tx:
-        expr = tx._build_partition_predicate(
-            partition_records={Record(i) for i in range(5000)},
-            spec=table_v2.metadata.spec(),
-            schema=table_v2.metadata.schema(),
-        )
+        expr = tx._build_partition_predicate(partition_records={Record(i) for i in range(5000)})
 
     bind(table_v2.metadata.schema(), expr, case_sensitive=True)