From 8025f9fa88ea40457086371945863033ceb4bf45 Mon Sep 17 00:00:00 2001 From: linzch3 <20380442+linzch3@users.noreply.github.com> Date: Sat, 14 Nov 2026 19:19:50 +0800 Subject: [PATCH] REST: Apply the scan filter when the server omits the residual Use the original request filter for file scan tasks without a server residual so count() applies the same filter as Arrow scans. --- pyiceberg/catalog/rest/__init__.py | 11 +++-- pyiceberg/table/__init__.py | 5 ++- tests/catalog/test_scan_planning_models.py | 51 +++++++++++++++++++++- 3 files changed, 60 insertions(+), 7 deletions(-) diff --git a/pyiceberg/catalog/rest/__init__.py b/pyiceberg/catalog/rest/__init__.py index 9ae6f489c2..94bf984a3d 100644 --- a/pyiceberg/catalog/rest/__init__.py +++ b/pyiceberg/catalog/rest/__init__.py @@ -71,6 +71,7 @@ UnauthorizedError, ViewAlreadyExistsError, ) +from pyiceberg.expressions import AlwaysTrue, BooleanExpression from pyiceberg.io import ( AWS_ACCESS_KEY_ID, AWS_PROFILE_NAME, @@ -860,13 +861,15 @@ def _poll_until_completed(self, identifier: str | Identifier, plan_id: str) -> P sleep_ms = min(sleep_ms * REST_SCAN_PLANNING_POLL_SCALE_FACTOR, REST_SCAN_PLANNING_POLL_MAX_SLEEP_MS) retries += 1 - def _expand_plan_tasks(self, identifier: str | Identifier, response: PlanCompleted) -> list[FileScanTask]: + def _expand_plan_tasks( + self, identifier: str | Identifier, response: PlanCompleted, fallback_residual: BooleanExpression + ) -> list[FileScanTask]: """Expand a completed plan response into FileScanTask objects, including pagination.""" tasks: list[FileScanTask] = [] # Collect tasks from initial response for task in response.file_scan_tasks: - tasks.append(FileScanTask.from_rest_response(task, response.delete_files)) + tasks.append(FileScanTask.from_rest_response(task, response.delete_files, fallback_residual=fallback_residual)) # Fetch and collect from additional batches pending_tasks = deque(response.plan_tasks) @@ -874,7 +877,7 @@ def _expand_plan_tasks(self, identifier: str | Identifier, response: PlanComplet plan_task = pending_tasks.popleft() batch = self._fetch_scan_tasks(identifier, plan_task) for task in batch.file_scan_tasks: - tasks.append(FileScanTask.from_rest_response(task, batch.delete_files)) + tasks.append(FileScanTask.from_rest_response(task, batch.delete_files, fallback_residual=fallback_residual)) pending_tasks.extend(batch.plan_tasks) return tasks @@ -901,7 +904,7 @@ def _plan_scan_result(self, identifier: str | Identifier, request: PlanTableScan if not isinstance(response, PlanCompleted): raise RuntimeError(f"Invalid planStatus for response: {type(response).__name__}") - tasks = self._expand_plan_tasks(identifier, response) + tasks = self._expand_plan_tasks(identifier, response, request.filter if request.filter is not None else AlwaysTrue()) return PlannedScanResult( tasks=tasks, storage_credentials=list(response.storage_credentials or []), diff --git a/pyiceberg/table/__init__.py b/pyiceberg/table/__init__.py index 9ef7519865..4d65084cbf 100644 --- a/pyiceberg/table/__init__.py +++ b/pyiceberg/table/__init__.py @@ -2343,12 +2343,15 @@ def __init__( def from_rest_response( rest_task: RESTFileScanTask, delete_files: list[RESTDeleteFile], + *, + fallback_residual: BooleanExpression = ALWAYS_TRUE, ) -> FileScanTask: """Convert a RESTFileScanTask to a FileScanTask. Args: rest_task: The REST file scan task. delete_files: The list of delete files from the ScanTasks response. + fallback_residual: The scan filter to use when the server omits the residual. Returns: A FileScanTask with the converted data and delete files. @@ -2371,7 +2374,7 @@ def from_rest_response( return FileScanTask( data_file=data_file, delete_files=resolved_deletes, - residual=rest_task.residual_filter if rest_task.residual_filter else ALWAYS_TRUE, + residual=rest_task.residual_filter if rest_task.residual_filter is not None else fallback_residual, ) diff --git a/tests/catalog/test_scan_planning_models.py b/tests/catalog/test_scan_planning_models.py index 03df70deba..dde8e86f6a 100644 --- a/tests/catalog/test_scan_planning_models.py +++ b/tests/catalog/test_scan_planning_models.py @@ -14,8 +14,11 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from pathlib import Path from typing import Any +import pyarrow as pa +import pyarrow.parquet as pq import pytest from pydantic import TypeAdapter, ValidationError from requests_mock import Mocker @@ -39,6 +42,7 @@ ValueMap, ) from pyiceberg.expressions import AlwaysTrue, EqualTo, Reference +from pyiceberg.io.pyarrow import schema_to_pyarrow from pyiceberg.manifest import FileFormat from pyiceberg.table import FileScanTask @@ -462,6 +466,7 @@ def test_plan_scan_completed_single_batch(rest_scan_catalog: RestCatalog, reques assert len(tasks) == 1 assert tasks[0].file.file_path == "s3://bucket/tbl/data/file1.parquet" + assert tasks[0].residual == AlwaysTrue() def test_plan_scan_with_pagination(rest_scan_catalog: RestCatalog, requests_mock: Mocker) -> None: @@ -474,7 +479,7 @@ def test_plan_scan_with_pagination(rest_scan_catalog: RestCatalog, requests_mock "status": "completed", "plan-id": "plan-123", "delete-files": [], - "file-scan-tasks": [{"data-file": file_one}], + "file-scan-tasks": [{"data-file": file_one, "residual-filter": True}], "plan-tasks": ["token-batch-2"], }, status_code=200, @@ -490,13 +495,16 @@ def test_plan_scan_with_pagination(rest_scan_catalog: RestCatalog, requests_mock status_code=200, ) - request = PlanTableScanRequest() + row_filter = EqualTo("x", 1) + request = PlanTableScanRequest(filter=row_filter) tasks = rest_scan_catalog.plan_scan(("db", "tbl"), request) assert len(tasks) == 2 assert tasks[0].file.file_path == "s3://bucket/tbl/data/file1.parquet" assert tasks[1].file.file_path == "s3://bucket/tbl/data/file2.parquet" + assert tasks[0].residual == AlwaysTrue() + assert tasks[1].residual == row_filter def test_plan_scan_with_delete_files(rest_scan_catalog: RestCatalog, requests_mock: Mocker) -> None: @@ -840,6 +848,45 @@ def test_scan_uses_server_side_planning_from_load_table_config( assert plan_mock.call_count == 1 +@pytest.mark.parametrize("paginated", [False, True]) +def test_count_applies_scan_filter_when_rest_residual_is_missing( + rest_scan_catalog: RestCatalog, + requests_mock: Mocker, + example_table_metadata_v2: dict[str, Any], + tmp_path: Path, + paginated: bool, +) -> None: + _mock_load_table(requests_mock, example_table_metadata_v2, {}) + table = rest_scan_catalog.load_table(("db", "tbl")) + data_file = tmp_path / "file.parquet" + pq.write_table( + pa.table({"x": [1, 2, 1], "y": [4, 5, 6], "z": [7, 8, 9]}, schema=schema_to_pyarrow(table.schema())), data_file + ) + file_task = { + "data-file": _rest_data_file(file_path=str(data_file), file_size_in_bytes=data_file.stat().st_size, record_count=3) + } + + requests_mock.post( + f"{TEST_URI}v1/namespaces/db/tables/tbl/plan", + json={ + "status": "completed", + "plan-id": "plan-123", + "delete-files": [], + "file-scan-tasks": [] if paginated else [file_task], + "plan-tasks": ["next"] if paginated else [], + }, + ) + if paginated: + requests_mock.post( + f"{TEST_URI}v1/namespaces/db/tables/tbl/tasks", + json={"delete-files": [], "file-scan-tasks": [file_task], "plan-tasks": []}, + ) + + scan = table.scan(row_filter="x != 1") + assert scan.to_arrow().num_rows == 1 + assert scan.count() == 1 + + def test_scan_keeps_client_side_planning_without_load_table_config( rest_client_planning_catalog: RestCatalog, requests_mock: Mocker, example_table_metadata_v2: dict[str, Any] ) -> None: