diff --git a/pyiceberg/table/__init__.py b/pyiceberg/table/__init__.py index c0adce84dc..cb1b69fa9a 100644 --- a/pyiceberg/table/__init__.py +++ b/pyiceberg/table/__init__.py @@ -26,12 +26,25 @@ from functools import cached_property from itertools import chain from types import TracebackType -from typing import TYPE_CHECKING, Any, TypeVar +from typing import ( + TYPE_CHECKING, + Any, + TypeVar, +) from pydantic import Field import pyiceberg.expressions.parser as parser -from pyiceberg.expressions import AlwaysFalse, AlwaysTrue, And, BooleanExpression, EqualTo, IsNull, Or, Reference +from pyiceberg.expressions import ( + AlwaysFalse, + AlwaysTrue, + And, + BooleanExpression, + EqualTo, + IsNull, + Or, + Reference, +) from pyiceberg.expressions.visitors import ( ResidualEvaluator, _InclusiveMetricsEvaluator, @@ -48,8 +61,13 @@ from pyiceberg.table.inspect import InspectTable from pyiceberg.table.locations import LocationProvider, load_location_provider from pyiceberg.table.maintenance import MaintenanceTable -from pyiceberg.table.metadata import INITIAL_SEQUENCE_NUMBER, TableMetadata -from pyiceberg.table.name_mapping import NameMapping +from pyiceberg.table.metadata import ( + INITIAL_SEQUENCE_NUMBER, + TableMetadata, +) +from pyiceberg.table.name_mapping import ( + NameMapping, +) from pyiceberg.table.refs import MAIN_BRANCH, SnapshotRef from pyiceberg.table.snapshots import ( Operation, @@ -82,7 +100,11 @@ update_table_metadata, ) from pyiceberg.table.update.schema import UpdateSchema -from pyiceberg.table.update.snapshot import ManageSnapshots, UpdateSnapshot, _FastAppendFiles +from pyiceberg.table.update.snapshot import ( + ManageSnapshots, + UpdateSnapshot, + _FastAppendFiles, +) from pyiceberg.table.update.sorting import UpdateSortOrder from pyiceberg.table.update.spec import UpdateSpec from pyiceberg.table.update.statistics import UpdateStatistics @@ -97,7 +119,9 @@ Record, TableVersion, ) -from pyiceberg.types import strtobool +from pyiceberg.types import ( + strtobool, +) from pyiceberg.utils.concurrent import ExecutorFactory from pyiceberg.utils.config import Config from pyiceberg.utils.properties import property_as_bool @@ -113,7 +137,11 @@ from pyiceberg_core.datafusion import IcebergDataFusionTable from pyiceberg.catalog import Catalog - from pyiceberg.catalog.rest.scan_planning import RESTContentFile, RESTDeleteFile, RESTFileScanTask + from pyiceberg.catalog.rest.scan_planning import ( + RESTContentFile, + RESTDeleteFile, + RESTFileScanTask, + ) ALWAYS_TRUE = AlwaysTrue() DOWNCAST_NS_TIMESTAMP_TO_US_ON_WRITE = "downcast-ns-timestamp-to-us-on-write" @@ -361,19 +389,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() @@ -593,9 +619,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, branch=branch) with self._append_snapshot_producer(snapshot_properties, branch=branch) as append_files: @@ -721,7 +745,11 @@ def delete( case_sensitive: A bool determine if the provided `delete_filter` is case-sensitive branch: Branch Reference to run the delete operation """ - from pyiceberg.io.pyarrow import ArrowScan, _dataframe_to_data_files, _expression_to_complementary_pyarrow + from pyiceberg.io.pyarrow import ( + ArrowScan, + _dataframe_to_data_files, + _expression_to_complementary_pyarrow, + ) if ( self.table_metadata.properties.get(TableProperties.DELETE_MODE, TableProperties.DELETE_MODE_DEFAULT) diff --git a/pyiceberg/table/update/snapshot.py b/pyiceberg/table/update/snapshot.py index 7931edacdd..c75e7f6740 100644 --- a/pyiceberg/table/update/snapshot.py +++ b/pyiceberg/table/update/snapshot.py @@ -26,7 +26,11 @@ from typing import TYPE_CHECKING, Generic from pyiceberg.avro.codecs import AvroCompressionCodec -from pyiceberg.expressions import AlwaysFalse, BooleanExpression, Or +from pyiceberg.expressions import ( + AlwaysFalse, + BooleanExpression, + Or, +) from pyiceberg.expressions.visitors import ( ROWS_MIGHT_NOT_MATCH, ROWS_MUST_MATCH, @@ -47,8 +51,9 @@ write_manifest, write_manifest_list, ) -from pyiceberg.partitioning import PartitionSpec -from pyiceberg.schema import Schema +from pyiceberg.partitioning import ( + PartitionSpec, +) from pyiceberg.table.refs import MAIN_BRANCH, SnapshotRefType from pyiceberg.table.snapshots import ( Operation, @@ -71,7 +76,10 @@ 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 @@ -102,8 +110,6 @@ class _SnapshotProducer(UpdateTableMetadata[U], Generic[U]): _deleted_data_files: set[DataFile] _compression: AvroCompressionCodec _target_branch: str | None - _predicate: BooleanExpression - _case_sensitive: bool def __init__( self, @@ -132,8 +138,6 @@ def __init__( self._parent_snapshot_id = ( snapshot.snapshot_id if (snapshot := self._transaction.table_metadata.snapshot_by_name(self._target_branch)) else None ) - self._predicate = AlwaysFalse() - self._case_sensitive = True def _validate_target_branch(self, branch: str | None) -> str | None: # if branch is none, write will be written into a staging snapshot @@ -178,8 +182,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( @@ -204,7 +213,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()) @@ -212,9 +228,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) @@ -333,9 +346,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] @@ -343,7 +353,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, @@ -358,35 +368,6 @@ def new_manifest_output(self) -> OutputFile: def fetch_manifest_entry(self, manifest: ManifestFile, discard_deleted: bool = True) -> list[ManifestEntry]: return manifest.fetch_manifest_entry(io=self._io, discard_deleted=discard_deleted) - def _build_partition_projection(self, spec_id: int) -> BooleanExpression: - project = inclusive_projection(self.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) - - 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) - ) - ) - class _DeleteFiles(_SnapshotProducer["_DeleteFiles"]): """Will delete manifest entries from the current snapshot based on the predicate. @@ -398,6 +379,22 @@ class _DeleteFiles(_SnapshotProducer["_DeleteFiles"]): From the specification """ + _predicate: BooleanExpression + _case_sensitive: bool + + def __init__( + self, + operation: Operation, + transaction: Transaction, + io: FileIO, + branch: str | None = MAIN_BRANCH, + commit_uuid: uuid.UUID | None = None, + snapshot_properties: dict[str, str] = EMPTY_DICT, + ): + super().__init__(operation, transaction, io, commit_uuid, snapshot_properties, branch) + self._predicate = AlwaysFalse() + self._case_sensitive = True + def _commit(self) -> UpdatesAndRequirements: # Only produce a commit when there is something to delete if self.files_affected: @@ -405,6 +402,25 @@ def _commit(self) -> UpdatesAndRequirements: else: return (), () + def _build_partition_projection(self, spec_id: int) -> BooleanExpression: + schema = self._transaction.table_metadata.schema() + spec = self._transaction.table_metadata.specs()[spec_id] + project = inclusive_projection(schema, spec, 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]: + schema = self._transaction.table_metadata.schema() + spec = self._transaction.table_metadata.specs()[spec_id] + return manifest_evaluator(spec, 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 + @cached_property def _compute_deletes(self) -> tuple[list[ManifestFile], list[ManifestEntry], bool]: """Computes all the delete operation and cache it when nothing changes. @@ -472,7 +488,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()) @@ -592,46 +615,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() - - 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) + 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] - # 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]: @@ -648,12 +661,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)