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..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 ) @@ -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 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 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.