diff --git a/pyiceberg/table/__init__.py b/pyiceberg/table/__init__.py index 3dffc2270c..4afdffe571 100644 --- a/pyiceberg/table/__init__.py +++ b/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, diff --git a/pyiceberg/table/update/snapshot.py b/pyiceberg/table/update/snapshot.py index 3c58f8ff44..faeea445d0 100644 --- a/pyiceberg/table/update/snapshot.py +++ b/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,7 +259,14 @@ 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()) @@ -263,9 +274,6 @@ def _write_delete_manifest() -> list[ManifestFile]: 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,9 +393,6 @@ 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] @@ -395,7 +400,7 @@ 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,7 +511,8 @@ 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 @@ -514,27 +520,13 @@ 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, diff --git a/tests/table/test_init.py b/tests/table/test_init.py index 3f1e97768c..9b8001b121 100644 --- a/tests/table/test_init.py +++ b/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)