From 29103de2b5ea228683eec90284dbefc38a08eee4 Mon Sep 17 00:00:00 2001 From: rich7420 Date: Sun, 25 Oct 2026 21:21:01 +0800 Subject: [PATCH 1/2] Fix time-travel row filters to use the snapshot schema Bind scan planning and Arrow filters to the same full schema as projection. Keep current-schema binding for deletes and unpinned partition inspection. --- pyiceberg/io/pyarrow.py | 4 +- pyiceberg/table/__init__.py | 38 ++++++++++++++---- pyiceberg/table/inspect.py | 2 +- tests/io/test_pyarrow.py | 27 +++++++------ tests/table/test_init.py | 78 +++++++++++++++++++++++++++++++++++++ tests/table/test_inspect.py | 20 ++++++++++ 6 files changed, 149 insertions(+), 20 deletions(-) diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py index 0cab26309b..b5337cba78 100644 --- a/pyiceberg/io/pyarrow.py +++ b/pyiceberg/io/pyarrow.py @@ -1849,11 +1849,13 @@ def __init__( limit: int | None = None, *, dictionary_columns: tuple[str, ...] = (), + row_filter_schema: Schema | None = None, ) -> None: self._table_metadata = table_metadata self._io = io self._projected_schema = projected_schema - self._bound_row_filter = bind(table_metadata.schema(), row_filter, case_sensitive=case_sensitive) + filter_schema = row_filter_schema if row_filter_schema is not None else table_metadata.schema() + self._bound_row_filter = bind(filter_schema, row_filter, case_sensitive=case_sensitive) self._case_sensitive = case_sensitive self._limit = limit self._downcast_ns_timestamp_to_us = Config().get_bool(DOWNCAST_NS_TIMESTAMP_TO_US_ON_WRITE) diff --git a/pyiceberg/table/__init__.py b/pyiceberg/table/__init__.py index 7da0e59d26..df5f3ab338 100644 --- a/pyiceberg/table/__init__.py +++ b/pyiceberg/table/__init__.py @@ -791,7 +791,17 @@ def delete( file_scan = self._scan(row_filter=delete_filter, case_sensitive=case_sensitive) if branch is not None: file_scan = file_scan.use_ref(branch) - files = file_scan.plan_files() + snapshot = file_scan.snapshot() + files = ( + ManifestGroupPlanner( + table_metadata=self.table_metadata, + io=self._table.io, + row_filter=delete_filter, + case_sensitive=case_sensitive, + ).plan_files(snapshot.manifests(self._table.io)) + if snapshot is not None + else [] + ) commit_uuid = uuid.uuid4() counter = itertools.count(0) @@ -2087,6 +2097,9 @@ def __init__( self.options = options self.limit = limit + def _scan_schema(self) -> Schema: + return self.table_metadata.schema() + @abstractmethod def plan_files(self) -> Iterable[ScanTask]: ... @@ -2205,7 +2218,7 @@ def snapshot(self) -> Snapshot | None: return self.table_metadata.snapshot_by_id(self.snapshot_id) return self.table_metadata.current_snapshot() - def projection(self) -> Schema: + def _scan_schema(self) -> Schema: current_schema = self.table_metadata.schema() if self.snapshot_id is not None: snapshot = self.table_metadata.snapshot_by_id(self.snapshot_id) @@ -2219,6 +2232,10 @@ def projection(self) -> Schema: warnings.warn(f"Metadata does not contain schema with id: {snapshot.schema_id}", stacklevel=2) else: raise ValueError(f"Snapshot not found: {self.snapshot_id}") + return current_schema + + def projection(self) -> Schema: + current_schema = self._scan_schema() if "*" in self.selected_fields: return current_schema @@ -2384,6 +2401,7 @@ def _to_arrow_via_file_scan_tasks( scan.case_sensitive, scan.limit, dictionary_columns=dictionary_columns, + row_filter_schema=scan._scan_schema(), ).to_table(tasks) @@ -2404,6 +2422,7 @@ def _to_arrow_batch_reader_via_file_scan_tasks( scan.case_sensitive, scan.limit, dictionary_columns=dictionary_columns, + row_filter_schema=scan._scan_schema(), ).to_record_batches(tasks) if dictionary_columns: @@ -2433,6 +2452,7 @@ def _manifest_planner(self) -> ManifestGroupPlanner: row_filter=self.row_filter, case_sensitive=self.case_sensitive, options=self.options, + schema=self._scan_schema(), ) @cached_property @@ -2565,6 +2585,7 @@ def count(self) -> int: projected_schema=self.projection(), row_filter=self.row_filter, case_sensitive=self.case_sensitive, + row_filter_schema=self._scan_schema(), ) tbl = arrow_scan.to_table([task]) res += len(tbl) @@ -2767,6 +2788,7 @@ class ManifestGroupPlanner: row_filter: BooleanExpression case_sensitive: bool options: Properties + schema: Schema def __init__( self, @@ -2775,12 +2797,14 @@ def __init__( row_filter: str | BooleanExpression = ALWAYS_TRUE, case_sensitive: bool = True, options: Properties = EMPTY_DICT, + schema: Schema | None = None, ): self.table_metadata = table_metadata self.io = io self.row_filter = _parse_row_filter(row_filter) self.case_sensitive = case_sensitive self.options = options + self.schema = schema if schema is not None else table_metadata.schema() @cached_property def partition_filters(self) -> KeyDefaultDict[int, BooleanExpression]: @@ -2873,16 +2897,16 @@ def plan_files( ] def _build_partition_projection(self, spec_id: int) -> BooleanExpression: - project = inclusive_projection(self.table_metadata.schema(), self.table_metadata.specs()[spec_id], self.case_sensitive) + project = inclusive_projection(self.schema, self.table_metadata.specs()[spec_id], self.case_sensitive) return project(self.row_filter) def _build_manifest_evaluator(self, spec_id: int) -> Callable[[ManifestFile], bool]: spec = self.table_metadata.specs()[spec_id] - return manifest_evaluator(spec, self.table_metadata.schema(), self.partition_filters[spec_id], self.case_sensitive) + return manifest_evaluator(spec, self.schema, self.partition_filters[spec_id], self.case_sensitive) def _build_partition_evaluator(self, spec_id: int) -> Callable[[DataFile], bool]: spec = self.table_metadata.specs()[spec_id] - partition_type = spec.partition_type(self.table_metadata.schema()) + partition_type = spec.partition_type(self.schema) partition_schema = Schema(*partition_type.fields) partition_expr = self.partition_filters[spec_id] evaluator = expression_evaluator(partition_schema, partition_expr, self.case_sensitive) @@ -2892,7 +2916,7 @@ def _build_partition_evaluator(self, spec_id: int) -> Callable[[DataFile], bool] return lambda data_file: evaluator(data_file.partition) def _build_metrics_evaluator(self) -> Callable[[DataFile], bool]: - schema = self.table_metadata.schema() + schema = self.schema include_empty_files = strtobool(self.options.get("include_empty_files", "false")) # Metrics evaluators keep file-specific state local to each call, so one @@ -2916,7 +2940,7 @@ def _build_residual_evaluator(self, spec_id: int) -> Callable[[DataFile], Residu spec=spec, expr=self.row_filter, case_sensitive=self.case_sensitive, - schema=self.table_metadata.schema(), + schema=self.schema, ) @staticmethod diff --git a/pyiceberg/table/inspect.py b/pyiceberg/table/inspect.py index e24e251fa4..447125d80f 100644 --- a/pyiceberg/table/inspect.py +++ b/pyiceberg/table/inspect.py @@ -403,7 +403,7 @@ def partitions( io=self.tbl.io, row_filter=row_filter, case_sensitive=case_sensitive, - snapshot_id=snapshot.snapshot_id, + snapshot_id=snapshot_id, ) partitions_map: dict[tuple[str, Any], Any] = {} diff --git a/tests/io/test_pyarrow.py b/tests/io/test_pyarrow.py index 0ba8117ef3..b7a787d69d 100644 --- a/tests/io/test_pyarrow.py +++ b/tests/io/test_pyarrow.py @@ -1423,6 +1423,20 @@ def test_identity_transform_column_projection(tmp_path: str, catalog: InMemoryCa ).to_arrow() == pa.table({"other_field": ["bar"]}) assert table.scan(row_filter="partition_id = 1 AND other_field = 'bar'", selected_fields=("other_field",)).count() == 1 + snapshot = table.current_snapshot() + assert snapshot is not None + with table.update_schema() as update: + update.rename_column("partition_id", "renamed_partition_id") + + scan = table.scan( + snapshot_id=snapshot.snapshot_id, + row_filter="partition_id = 1 AND other_field = 'bar'", + selected_fields=("other_field",), + ) + assert scan.to_arrow() == pa.table({"other_field": ["bar"]}) + assert scan.to_arrow_batch_reader().read_all().to_pydict() == {"other_field": ["bar"]} + assert scan.count() == 1 + @pytest.mark.parametrize( "partition_field_type, arrow_partition_type, partition_value, row_filter", @@ -5959,10 +5973,9 @@ def test_to_arrow_batch_reader_preserves_dictionary_columns(tmpdir: str) -> None # RecordBatchReader.cast(target_schema) was called with a plain schema, silently converting # dictionary arrays back to their value type so to_arrow_batch_reader(dictionary_columns=...) # .read_all() returned plain strings instead of dictionary-encoded arrays. - from pyiceberg.expressions import AlwaysTrue from pyiceberg.io.pyarrow import PyArrowFileIO from pyiceberg.partitioning import PartitionSpec - from pyiceberg.table import FileScanTask, _to_arrow_batch_reader_via_file_scan_tasks + from pyiceberg.table import DataScan, FileScanTask, _to_arrow_batch_reader_via_file_scan_tasks from pyiceberg.table.metadata import TableMetadataV2 arrow_schema = pa.schema( @@ -5990,17 +6003,9 @@ def test_to_arrow_batch_reader_preserves_dictionary_columns(tmpdir: str) -> None partition_specs=[PartitionSpec()], ) - class _MockScan: - def __init__(self) -> None: - self.table_metadata = table_metadata - self.io = PyArrowFileIO() - self.row_filter = AlwaysTrue() - self.case_sensitive = True - self.limit = None - tasks = [FileScanTask(data_file)] result = _to_arrow_batch_reader_via_file_scan_tasks( - _MockScan(), # type: ignore[arg-type] + DataScan(table_metadata=table_metadata, io=PyArrowFileIO()), iceberg_schema, tasks, dictionary_columns=("label",), diff --git a/tests/table/test_init.py b/tests/table/test_init.py index f6b984f1a3..c8274f6cb6 100644 --- a/tests/table/test_init.py +++ b/tests/table/test_init.py @@ -21,6 +21,7 @@ from pathlib import Path from typing import Any +import pyarrow as pa import pytest from pydantic import BaseModel, TypeAdapter, ValidationError from pytest_lazy_fixtures import lf @@ -370,6 +371,83 @@ def test_table_scan_projection_unknown_column(table_v2: Table) -> None: assert "Could not find column: 'a'" in str(exc_info.value) +def test_time_travel_row_filter_after_drop_and_readd(catalog: Catalog, arrow_table_simple: pa.Table) -> None: + catalog.create_namespace("default") + table = catalog.create_table("default.readded_column", schema=arrow_table_simple.schema) + table.append(arrow_table_simple) + snapshot = table.current_snapshot() + assert snapshot is not None + table.manage_snapshots().create_tag(snapshot.snapshot_id, "before_schema_change").commit() + + with table.update_schema() as update: + update.delete_column("bar") + with table.update_schema() as update: + update.add_column("bar", IntegerType()) + + for scan in ( + table.scan(snapshot_id=snapshot.snapshot_id, row_filter="bar >= 2"), + table.scan(row_filter="BAR >= 2", case_sensitive=False).use_ref("before_schema_change"), + ): + for selected_fields in (("*",), ("foo",)): + projected_scan = scan.select(*selected_fields) + expected = arrow_table_simple.slice(1) + if selected_fields != ("*",): + expected = expected.select(selected_fields) + assert projected_scan.to_arrow().to_pydict() == expected.to_pydict() + assert projected_scan.to_arrow_batch_reader().read_all().to_pydict() == expected.to_pydict() + assert projected_scan.count() == 2 + + for current_scan in (table.scan(row_filter="bar >= 2"), table.incremental_append_scan(row_filter="bar >= 2")): + assert current_scan.to_arrow().num_rows == 0 + assert current_scan.to_arrow_batch_reader().read_all().num_rows == 0 + assert table.scan(row_filter="bar >= 2").count() == 0 + assert table.scan(row_filter="bar IS NULL").count() == 3 + + +def test_time_travel_row_filter_after_partition_column_rename(catalog: Catalog, arrow_table_simple: pa.Table) -> None: + catalog.create_namespace("default") + table = catalog.create_table("default.renamed_partition_column", schema=arrow_table_simple.schema) + with table.update_spec() as update: + update.add_identity("baz") + table.append(arrow_table_simple) + snapshot = table.current_snapshot() + assert snapshot is not None + with table.update_schema() as update: + update.rename_column("baz", "enabled") + update.add_column("added", IntegerType()) + + scan = table.scan(snapshot_id=snapshot.snapshot_id, row_filter="baz = false", selected_fields=("foo",)) + tasks = list(scan.plan_files()) + assert len(tasks) == 1 + assert tasks[0].residual == AlwaysTrue() + assert scan.to_arrow().to_pydict() == {"foo": ["b"]} + assert scan.to_arrow_batch_reader().read_all().to_pydict() == {"foo": ["b"]} + assert scan.count() == 1 + + for column, row_filter in (("enabled", "enabled = false"), ("added", "added = 1")): + with pytest.raises(ValueError, match=f"Could not find field with name {column}"): + table.scan(snapshot_id=snapshot.snapshot_id, row_filter=row_filter).to_arrow() + assert table.scan(row_filter="enabled = false", selected_fields=("foo",)).to_arrow().to_pydict() == {"foo": ["b"]} + + +def test_delete_after_column_rename_uses_current_schema(catalog: Catalog, arrow_table_simple: pa.Table) -> None: + catalog.create_namespace("default") + for branch in (MAIN_BRANCH, "audit"): + table = catalog.create_table(f"default.delete_after_rename_{branch}", schema=arrow_table_simple.schema) + table.append(arrow_table_simple) + snapshot = table.current_snapshot() + assert snapshot is not None + if branch != MAIN_BRANCH: + table.manage_snapshots().create_branch(snapshot.snapshot_id, branch).commit() + with table.update_schema() as update: + update.rename_column("bar", "score") + + table.delete("score >= 2", branch=branch) + assert table.scan().use_ref(branch).to_arrow().to_pydict() == {"foo": ["a"], "score": [1], "baz": [True]} + if branch != MAIN_BRANCH: + assert table.scan().count() == 3 + + def test_data_scan_plan_files_no_current_snapshot(example_table_metadata_no_snapshot_v1: dict[str, Any]) -> None: table = Table( identifier=("default", "test_no_snapshot"), diff --git a/tests/table/test_inspect.py b/tests/table/test_inspect.py index 2404eadc74..cc8b68051f 100644 --- a/tests/table/test_inspect.py +++ b/tests/table/test_inspect.py @@ -76,6 +76,26 @@ def test_inspect_entries_and_files_render_null_bound(catalog: InMemoryCatalog) - assert files_metrics["upper_bound"] is None +def test_partitions_row_filter_after_column_rename(catalog: InMemoryCatalog) -> None: + schema = Schema(NestedField(1, "category", StringType())) + table = catalog.create_table( + "default.renamed_partition", + schema=schema, + partition_spec=PartitionSpec(PartitionField(1, 1000, IdentityTransform(), "category")), + ) + table.append(pa.table({"category": ["a", "b", "a"]})) + snapshot = table.current_snapshot() + assert snapshot is not None + with table.update_schema() as update: + update.rename_column("category", "cat") + + current = table.inspect.partitions(row_filter="cat = 'a'") + historical = table.inspect.partitions(snapshot_id=snapshot.snapshot_id, row_filter="category = 'a'") + expected = [{"partition": {"category": "a"}, "record_count": 2}] + assert current.select(["partition", "record_count"]).to_pylist() == expected + assert historical.select(["partition", "record_count"]).to_pylist() == expected + + @pytest.mark.parametrize("newest_first", [False, True]) def test_partitions_last_updated_uses_latest_snapshot_regardless_of_order(newest_first: bool) -> None: # Manifest entries are visited in manifest order, which is not chronological, so the From 1775101c7b0e8b699c1a7f0599bac72d5ddde8f1 Mon Sep 17 00:00:00 2001 From: rich7420 Date: Mon, 26 Oct 2026 13:38:27 +0800 Subject: [PATCH 2/2] Keep transaction scans on the current schema Use the transaction schema for both filtering and projection so schema changes do not bind upsert joins to an older branch snapshot. Share this scan behavior with deletes and cover added or replaced join columns, branch writes, staged writes, and updates after a staged schema change. --- pyiceberg/table/__init__.py | 23 +++++---------- tests/table/test_upsert.py | 58 +++++++++++++++++++++++++++++++++++++ 2 files changed, 66 insertions(+), 15 deletions(-) diff --git a/pyiceberg/table/__init__.py b/pyiceberg/table/__init__.py index df5f3ab338..2d7eb65e9a 100644 --- a/pyiceberg/table/__init__.py +++ b/pyiceberg/table/__init__.py @@ -326,7 +326,7 @@ def _apply( def _scan(self, row_filter: str | BooleanExpression = ALWAYS_TRUE, case_sensitive: bool = True) -> DataScan: """Minimal data scan of the table with the current state of the transaction.""" - return DataScan( + return _TransactionScan( table_metadata=self.table_metadata, io=self._table.io, row_filter=row_filter, case_sensitive=case_sensitive ) @@ -791,17 +791,7 @@ def delete( file_scan = self._scan(row_filter=delete_filter, case_sensitive=case_sensitive) if branch is not None: file_scan = file_scan.use_ref(branch) - snapshot = file_scan.snapshot() - files = ( - ManifestGroupPlanner( - table_metadata=self.table_metadata, - io=self._table.io, - row_filter=delete_filter, - case_sensitive=case_sensitive, - ).plan_files(snapshot.manifests(self._table.io)) - if snapshot is not None - else [] - ) + files = file_scan.plan_files() commit_uuid = uuid.uuid4() counter = itertools.count(0) @@ -946,9 +936,7 @@ def upsert( # We must use Transaction.table_metadata for the scan. This includes all uncommitted - but relevant - changes. - matched_iceberg_record_batches_scan = DataScan( - table_metadata=self.table_metadata, - io=self._table.io, + matched_iceberg_record_batches_scan = self._scan( row_filter=matched_predicate, case_sensitive=case_sensitive, ) @@ -2592,6 +2580,11 @@ def count(self) -> int: return res +class _TransactionScan(DataScan): + def _scan_schema(self) -> Schema: + return self.table_metadata.schema() + + IAS = TypeVar("IAS", bound="IncrementalAppendScan", covariant=True) diff --git a/tests/table/test_upsert.py b/tests/table/test_upsert.py index 78ddbc7c5c..dc80ff6057 100644 --- a/tests/table/test_upsert.py +++ b/tests/table/test_upsert.py @@ -427,6 +427,64 @@ def test_upsert_into_empty_table(catalog: Catalog) -> None: assert upd.rows_inserted == 4 +def test_upsert_uses_current_schema_after_join_column_change(catalog: Catalog) -> None: + for drop_existing in (False, True): + for branch in ("main", "audit", None): + original = pa.table({"value": [1, 2]}) + if drop_existing: + original = original.append_column("key", pa.array(["old", "keep"])) + table = catalog.create_table(f"default.changed_join_column_{drop_existing}_{branch}", schema=original.schema) + table.append(original) + snapshot = table.current_snapshot() + assert snapshot is not None + if branch == "audit": + table.manage_snapshots().create_branch(snapshot.snapshot_id, branch).commit() + if drop_existing: + with table.update_schema() as update: + update.delete_column("key") + with table.update_schema() as update: + update.add_column("key", StringType()) + + result = table.upsert(pa.table({"value": [3], "key": ["old"]}), join_cols=["key"], branch=branch) + assert_upsert_result(result, expected_updated=0, expected_inserted=1) + scan = table.scan(snapshot_id=table.snapshots()[-1].snapshot_id) if branch is None else table.scan().use_ref(branch) + assert sorted(scan.to_arrow().to_pylist(), key=lambda row: row["value"]) == [ + {"value": 1, "key": None}, + {"value": 2, "key": None}, + {"value": 3, "key": "old"}, + ] + if branch != "main": + assert table.metadata.current_snapshot_id == snapshot.snapshot_id + assert table.scan().count() == 2 + + +def test_upsert_projects_staged_schema_for_matched_rows(catalog: Catalog) -> None: + for branch in ("main", "audit"): + original = pa.table({"id": [1, 2], "value": [10, 20]}) + table = catalog.create_table(f"default.staged_upsert_schema_{branch}", schema=original.schema) + table.append(original) + snapshot = table.current_snapshot() + assert snapshot is not None + if branch == "audit": + table.manage_snapshots().create_branch(snapshot.snapshot_id, branch).commit() + + with table.transaction() as transaction: + with transaction.update_schema() as update: + update.add_column("note", StringType()) + result = transaction.upsert( + pa.table({"id": [1], "value": [11], "note": ["updated"]}), join_cols=["id"], branch=branch + ) + assert_upsert_result(result, expected_updated=1, expected_inserted=0) + + assert sorted(table.scan().use_ref(branch).to_arrow().to_pylist(), key=lambda row: row["id"]) == [ + {"id": 1, "value": 11, "note": "updated"}, + {"id": 2, "value": 20, "note": None}, + ] + if branch != "main": + assert table.metadata.current_snapshot_id == snapshot.snapshot_id + assert table.scan().to_arrow()["value"].to_pylist() == [10, 20] + + def test_create_match_filter_single_condition() -> None: """ Test create_match_filter with a composite key where the source yields exactly one unique key.