Skip to content

Revert PR #3011 with additional test changes #3776

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
rambleraptor wants to merge 1 commit into
base: main
Choose a base branch
from
+60 −83
Open

Revert PR #3011 with additional test changes #3776

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 5 additions & 9 deletions pyiceberg/table/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand Down Expand Up @@ -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,
Expand Down
123 changes: 54 additions & 69 deletions pyiceberg/table/update/snapshot.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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
Expand Down Expand Up @@ -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(
Expand All @@ -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)
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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())
Expand Down Expand Up @@ -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]:
Expand All @@ -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,
Expand Down
6 changes: 1 addition & 5 deletions tests/table/test_init.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down