diff --git a/pyiceberg/table/__init__.py b/pyiceberg/table/__init__.py index 63b87d290e..018556bdf0 100644 --- a/pyiceberg/table/__init__.py +++ b/pyiceberg/table/__init__.py @@ -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 diff --git a/pyiceberg/table/update/snapshot.py b/pyiceberg/table/update/snapshot.py index 7931edacdd..7172a6fa3e 100644 --- a/pyiceberg/table/update/snapshot.py +++ b/pyiceberg/table/update/snapshot.py @@ -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, @@ -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 @@ -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]: @@ -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 = [] @@ -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): diff --git a/tests/table/test_init.py b/tests/table/test_init.py index 1670e62587..e60dc58956 100644 --- a/tests/table/test_init.py +++ b/tests/table/test_init.py @@ -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}"