Fix time-travel row filters to use the snapshot schema

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)

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
         )
 
@@ -936,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,
         )
@@ -2087,6 +2085,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 +2206,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 +2220,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 +2389,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 +2410,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 +2440,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,12 +2573,18 @@ 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)
         return res
 
 
+class _TransactionScan(DataScan):
+    def _scan_schema(self) -> Schema:
+        return self.table_metadata.schema()
+
+
 IAS = TypeVar("IAS", bound="IncrementalAppendScan", covariant=True)
 
 
@@ -2767,6 +2781,7 @@ class ManifestGroupPlanner:
     row_filter: BooleanExpression
     case_sensitive: bool
     options: Properties
+    schema: Schema
 
     def __init__(
         self,
@@ -2775,12 +2790,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 +2890,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 +2909,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 +2933,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

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] = {}

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",),

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"),

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