Skip to content

fix: include NULL in delete predicate for evolved partition fields in dynamic_partition_overwrite #3460

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
GayathriSrividya wants to merge 1 commit into
base: main
Choose a base branch
from
+210 −8
Open

fix: include NULL in delete predicate for evolved partition fields in dynamic_partition_overwrite #3460

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
123 changes: 120 additions & 3 deletions pyiceberg/table/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -593,10 +593,127 @@ 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()
current_spec = self.table_metadata.spec()
all_specs = self.table_metadata.specs()
schema = self.table_metadata.schema()

# Keep the existing dynamic overwrite behavior for non-evolution tables.
# We only need per-spec predicate handling when some historical specs are
# missing current partition source IDs.
current_source_ids = {field.source_id for field in current_spec.fields}
has_missing_partition_fields_in_history = any(
current_source_ids - {field.source_id for field in historical_spec.fields} for historical_spec in all_specs.values()
)
self.delete(delete_filter=delete_filter, snapshot_properties=snapshot_properties, branch=branch)

if not has_missing_partition_fields_in_history:
delete_filter = self._build_partition_predicate(
partition_records=partitions_to_overwrite,
spec=current_spec,
schema=schema,
)
self.delete(delete_filter=delete_filter, snapshot_properties=snapshot_properties, branch=branch)
else:
# Build per-spec delete predicates to handle partition spec evolution correctly.
#
# When a partition field was added via spec evolution, data files written under
# older specs carry NULL for that field (because it was absent from the schema at
# write time). A single "category=A AND region=us" predicate would never match
# those files because the strict-metrics evaluator sees region=NULL != "us".
#
# To fix this, we compute a per-spec predicate for every historical spec:
# - For specs that include all current partition fields -> use exact-match predicate.
# - For specs that are missing some current partition fields -> also accept NULL
# for the missing fields.
#
# These per-spec predicates are stored on the delete snapshot producer so that
# _compute_deletes uses the right predicate when evaluating each manifest file.
source_id_to_pos = {field.source_id: pos for pos, field in enumerate(current_spec.fields)}
source_id_to_col = {field.source_id: schema.find_field(field.source_id).name for field in current_spec.fields}
exact_delete_filter = self._build_partition_predicate(
partition_records=partitions_to_overwrite,
spec=current_spec,
schema=schema,
)

per_spec_predicates: dict[int, BooleanExpression] = {}
for spec_id, hist_spec in all_specs.items():
hist_source_ids = {field.source_id for field in hist_spec.fields}
missing_source_ids = current_source_ids - hist_source_ids
has_overlap_with_current = bool(hist_source_ids & current_source_ids)

per_record_exprs: list[BooleanExpression] = []
for partition_record in partitions_to_overwrite:
predicates: list[BooleanExpression] = []
for source_id, col_name in source_id_to_col.items():
value = partition_record[source_id_to_pos[source_id]]
if value is not None:
field_pred: BooleanExpression = EqualTo(Reference(col_name), value)
if source_id in missing_source_ids and has_overlap_with_current:
field_pred = Or(field_pred, IsNull(Reference(col_name)))
else:
field_pred = IsNull(Reference(col_name))
predicates.append(field_pred)

per_record_exprs.append(And(*predicates) if len(predicates) > 1 else predicates[0])

per_spec_predicates[spec_id] = Or(*per_record_exprs) if len(per_record_exprs) > 1 else per_record_exprs[0]

# Open the delete snapshot and set per-spec predicates before committing.
# This mirrors Transaction.delete() but injects per_spec_predicates so that
# _compute_deletes uses the right predicate for each historical spec.
from pyiceberg.io.pyarrow import ArrowScan, _dataframe_to_data_files, _expression_to_complementary_pyarrow

with self.update_snapshot(snapshot_properties=snapshot_properties, branch=branch).delete() as delete_snapshot:
delete_snapshot._per_spec_predicates = per_spec_predicates
delete_snapshot.delete_by_predicate(exact_delete_filter)

# Handle partial-match files that need to be rewritten (copy-on-write).
if delete_snapshot.rewrites_needed is True:
bound_delete_filter = bind(self.table_metadata.schema(), exact_delete_filter, case_sensitive=True)
preserve_row_filter = _expression_to_complementary_pyarrow(bound_delete_filter, self.table_metadata.schema())

file_scan = self._scan(row_filter=exact_delete_filter)
if branch is not None:
file_scan = file_scan.use_ref(branch)

rewrite_uuid = uuid.uuid4()
rewrite_counter = itertools.count(0)
replaced_files: list[tuple[DataFile, list[DataFile]]] = []
for original_file in file_scan.plan_files():
df_orig = ArrowScan(
table_metadata=self.table_metadata,
io=self._table.io,
projected_schema=self.table_metadata.schema(),
row_filter=AlwaysTrue(),
).to_table(tasks=[original_file])
filtered_df = df_orig.filter(preserve_row_filter)
if len(filtered_df) == 0:
replaced_files.append((original_file.file, []))
elif len(df_orig) != len(filtered_df):
replaced_files.append(
(
original_file.file,
list(
_dataframe_to_data_files(
io=self._table.io,
df=filtered_df,
table_metadata=self.table_metadata,
write_uuid=rewrite_uuid,
counter=rewrite_counter,
)
),
)
)

if replaced_files:
with self.update_snapshot(
snapshot_properties=snapshot_properties, branch=branch
).overwrite() as overwrite_snapshot:
overwrite_snapshot.commit_uuid = rewrite_uuid
for original_data_file, replacement_data_files in replaced_files:
overwrite_snapshot.delete_data_file(original_data_file)
for replacement_data_file in replacement_data_files:
overwrite_snapshot.append_data_file(replacement_data_file)

with self._append_snapshot_producer(snapshot_properties, branch=branch) as append_files:
append_files.commit_uuid = append_snapshot_commit_uuid
Expand Down
20 changes: 15 additions & 5 deletions pyiceberg/table/update/snapshot.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -104,6 +104,7 @@ class _SnapshotProducer(UpdateTableMetadata[U], Generic[U]):
_target_branch: str | None
_predicate: BooleanExpression
_case_sensitive: bool
_per_spec_predicates: dict[int, BooleanExpression]

def __init__(
self,
Expand Down Expand Up @@ -134,6 +135,7 @@ def __init__(
)
self._predicate = AlwaysFalse()
self._case_sensitive = True
self._per_spec_predicates = {}

def _validate_target_branch(self, branch: str | None) -> str | None:
# if branch is none, write will be written into a staging snapshot
Expand Down Expand Up @@ -360,7 +362,8 @@ def fetch_manifest_entry(self, manifest: ManifestFile, discard_deleted: bool = T

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)
predicate = self._per_spec_predicates.get(spec_id, self._predicate)
return project(predicate)

@cached_property
def partition_filters(self) -> KeyDefaultDict[int, BooleanExpression]:
Expand Down Expand Up @@ -431,10 +434,14 @@ def _copy_with_new_status(entry: ManifestEntry, status: ManifestEntryStatus) ->
schema = table_metadata.schema()

manifest_evaluators: dict[int, Callable[[ManifestFile], bool]] = KeyDefaultDict(self._build_manifest_evaluator)
strict_metrics_evaluator = _StrictMetricsEvaluator(schema, self._predicate, case_sensitive=self._case_sensitive).eval
inclusive_metrics_evaluator = _InclusiveMetricsEvaluator(
schema, self._predicate, case_sensitive=self._case_sensitive
).eval

def _strict_metrics_for_spec(spec_id: int) -> Callable[[DataFile], bool]:
predicate = self._per_spec_predicates.get(spec_id, self._predicate)
return _StrictMetricsEvaluator(schema, predicate, case_sensitive=self._case_sensitive).eval

def _inclusive_metrics_for_spec(spec_id: int) -> Callable[[DataFile], bool]:
predicate = self._per_spec_predicates.get(spec_id, self._predicate)
return _InclusiveMetricsEvaluator(schema, predicate, case_sensitive=self._case_sensitive).eval

existing_manifests = []
total_deleted_entries = []
Expand All @@ -454,6 +461,9 @@ def _copy_with_new_status(entry: ManifestEntry, status: ManifestEntryStatus) ->
existing_manifests.append(manifest_file)
else:
# It is relevant, let's check out the content
spec_id = manifest_file.partition_spec_id
strict_metrics_evaluator = _strict_metrics_for_spec(spec_id)
inclusive_metrics_evaluator = _inclusive_metrics_for_spec(spec_id)
deleted_entries = []
existing_entries = []
for entry in manifest_file.fetch_manifest_entry(io=self._io, discard_deleted=True):
Expand Down
75 changes: 75 additions & 0 deletions tests/table/test_init.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1989,3 +1989,78 @@ def test_build_large_partition_predicate(table_v2: Table) -> None:
)

bind(table_v2.metadata.schema(), expr, case_sensitive=True)


def test_dynamic_partition_overwrite_spec_evolution(tmp_path: Any) -> None:
"""Regression test for https://github.com/apache/iceberg-python/issues/3148.

After partition spec evolution, dynamic_partition_overwrite must delete data files
written under the old spec (where the new partition field was absent / NULL) when
overwriting the matching logical partition.
"""
import tempfile

import pyarrow as pa

from pyiceberg.catalog import load_catalog
from pyiceberg.transforms import IdentityTransform
from pyiceberg.types import LongType

with tempfile.TemporaryDirectory() as warehouse:
catalog = load_catalog("test", type="sql", uri=f"sqlite:///{warehouse}/catalog.db", warehouse=f"file://{warehouse}")
catalog.create_namespace("default")

schema = Schema(
NestedField(1, "category", StringType(), required=False),
NestedField(2, "region", StringType(), required=False),
NestedField(3, "value", LongType(), required=False),
)
spec_v0 = PartitionSpec(PartitionField(source_id=1, field_id=1000, transform=IdentityTransform(), name="category"))
table = catalog.create_table("default.test_spec_evo", schema=schema, partition_spec=spec_v0)

# Write under spec-0 (region is NULL — field exists in schema but not in partition spec)
table.append(
pa.table(
{
"category": pa.array(["A", "A", "B"], type=pa.string()),
"region": pa.array([None, None, None], type=pa.string()),
"value": pa.array([1, 2, 10], type=pa.int64()),
}
)
)

# Evolve to spec-1: add region as a partition field
with table.update_spec() as u:
u.add_field("region", IdentityTransform(), "region")
table = catalog.load_table("default.test_spec_evo")

# Write under spec-1
table.append(
pa.table(
{
"category": pa.array(["A", "B"], type=pa.string()),
"region": pa.array(["us", "us"], type=pa.string()),
"value": pa.array([100, 200], type=pa.int64()),
}
)
)

# Overwrite partition {A, us} — must also delete stale spec-0 {A} files
table.dynamic_partition_overwrite(
pa.table(
{
"category": pa.array(["A"], type=pa.string()),
"region": pa.array(["us"], type=pa.string()),
"value": pa.array([999], type=pa.int64()),
}
)
)

result = table.scan().to_arrow().to_pydict()
a_values = sorted([v for c, v in zip(result["category"], result["value"], strict=True) if c == "A"])
b_values = sorted([v for c, v in zip(result["category"], result["value"], strict=True) if c == "B"])

# Spec-0 rows 1,2 (category=A, region=NULL) should be gone; only 999 remains
assert a_values == [999], f"Expected [999] but got {a_values}"
# B rows from both specs should be untouched
assert b_values == [10, 200], f"Expected [10, 200] but got {b_values}"