Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions mkdocs/docs/api.md
Original file line number Diff line number Diff line change
Expand Up @@ -1695,6 +1695,8 @@ scan = table.scan(
[task.file.file_path for task in scan.plan_files()]
```

When the REST catalog returns `scan-planning-mode=server` and advertises the plan endpoint, `plan_files()` / `to_arrow()` use server-side scan planning. The mode can also be returned per table in the `loadTable` response `config`, which takes precedence over the catalog-level setting, so a server can require server-side planning for some tables while others keep client-side planning. Catalogs that return async plans (`status=submitted`) are polled automatically until they reach a terminal state; see [REST Catalog configuration](configuration.md#rest-catalog).

The low level API `plan_files` methods returns a set of tasks that provide the files that might contain matching rows:

```json
Expand Down
4 changes: 4 additions & 0 deletions mkdocs/docs/configuration.md
Original file line number Diff line number Diff line change
Expand Up @@ -408,6 +408,10 @@ catalog:
| snapshot-loading-mode | refs | The snapshots to return in the body of the metadata. Setting the value to `all` would return the full set of snapshots currently valid for the table. Setting the value to `refs` would load all snapshots referenced by branches or tags. |
| `header.X-Iceberg-Access-Delegation` | `vended-credentials` | Signal to the server that the client supports delegated access via a comma-separated list of access mechanisms. The server may choose to supply access via any or none of the requested mechanisms. When using `vended-credentials`, the server provides temporary credentials to the client. When using `remote-signing`, the server signs requests on behalf of the client. (default: `vended-credentials`) |
| view-endpoints-supported | false | For backwards compatibility with older REST servers. Set to `true` if the server supports view endpoints but doesn't send the `endpoints` field in the ConfigResponse. |
| scan-planning-mode | client | When set to `server`, and the catalog advertises the plan-table-scan endpoint, `table.scan()` uses REST server-side scan planning. May be set by the client, returned by the catalog via `GET /v1/config` defaults/overrides, or returned per table in the `config` of the `loadTable` response. The `loadTable` value takes precedence over the catalog-level value, which lets a server request server-side planning only for specific tables. Async plans (`status=submitted`) are polled via `GET .../plan/{plan-id}` until terminal state (completed / failed / cancelled). |
| rest-scan-planning.poll-timeout-ms | 300000 | Maximum time to wait for an async scan plan to complete before failing (default: 5 minutes). |

When server-side planning returns `storage-credentials` on a completed plan, PyIceberg applies them to the scan-scoped FileIO (layered on top of the existing table/load-time IO properties) so planned data and delete files can be read using the creds vended by the server.

#### Headers in REST Catalog

Expand Down
11 changes: 8 additions & 3 deletions pyiceberg/catalog/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -804,8 +804,13 @@ def namespace_to_string(identifier: str | Identifier, err: type[ValueError] | ty
return ".".join(segment.strip() for segment in tuple_identifier)

@abstractmethod
def supports_server_side_planning(self) -> bool:
"""Check if the catalog supports server-side scan planning."""
def supports_server_side_planning(self, table_config: Properties = EMPTY_DICT) -> bool:
"""Check if server-side scan planning should be used.

Args:
table_config: Table configuration returned by the catalog when loading the table,
which may override the catalog-level scan planning mode for a single table.
"""

@staticmethod
def identifier_to_database(
Expand Down Expand Up @@ -907,7 +912,7 @@ def __init__(self, name: str, **properties: str):
super().__init__(name, **properties)

@override
def supports_server_side_planning(self) -> bool:
def supports_server_side_planning(self, table_config: Properties = EMPTY_DICT) -> bool:
return False

@override
Expand Down
2 changes: 1 addition & 1 deletion pyiceberg/catalog/noop.py
Original file line number Diff line number Diff line change
Expand Up @@ -98,7 +98,7 @@ def drop_table(self, identifier: str | Identifier) -> None:
raise NotImplementedError

@override
def supports_server_side_planning(self) -> bool:
def supports_server_side_planning(self, table_config: Properties = EMPTY_DICT) -> bool:
return False

@override
Expand Down
253 changes: 230 additions & 23 deletions pyiceberg/catalog/rest/__init__.py

Large diffs are not rendered by default.

15 changes: 14 additions & 1 deletion pyiceberg/catalog/rest/scan_planning.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,9 +16,10 @@
# under the License.
from __future__ import annotations

from dataclasses import dataclass, field
from datetime import date, datetime, time
from decimal import Decimal
from typing import Annotated, Generic, Literal, TypeAlias, TypeVar
from typing import TYPE_CHECKING, Annotated, Generic, Literal, TypeAlias, TypeVar
from uuid import UUID

from pydantic import Field, model_validator
Expand All @@ -28,6 +29,9 @@
from pyiceberg.manifest import FileFormat
from pyiceberg.typedef import IcebergBaseModel

if TYPE_CHECKING:
from pyiceberg.table import FileScanTask

# Primitive types that can appear in partition values and bounds
PrimitiveTypeValue: TypeAlias = bool | int | float | str | Decimal | UUID | date | time | datetime | bytes

Expand Down Expand Up @@ -207,3 +211,12 @@ class FetchScanTasksRequest(IcebergBaseModel):
"""Request body for fetching scan tasks endpoint."""

plan_task: str = Field(alias="plan-task")


@dataclass(frozen=True)
class PlannedScanResult:
"""Result of REST server-side scan planning, including optional storage credentials."""

tasks: list[FileScanTask]
storage_credentials: list[StorageCredential] = field(default_factory=list)
plan_id: str | None = None
8 changes: 8 additions & 0 deletions pyiceberg/exceptions.py
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,14 @@ class NoSuchPlanTaskError(Exception):
"""Raised when a scan plan task is not found."""


class NoSuchPlanIdError(Exception):
"""Raised when a scan plan-id is not found."""


class RemotePlanTimeoutError(Exception):
"""Raised when async remote scan planning does not complete within configured limits."""


class RESTError(Exception):
"""Raises when there is an unknown response from the REST Catalog."""

Expand Down
13 changes: 11 additions & 2 deletions pyiceberg/table/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -1437,6 +1437,7 @@ def scan(
limit=limit,
catalog=self.catalog,
table_identifier=self._identifier,
table_config=self.config,
)

def incremental_append_scan(
Expand Down Expand Up @@ -2155,6 +2156,7 @@ class TableScan(BaseScan):
snapshot_id: int | None
catalog: Catalog | None
table_identifier: Identifier | None
table_config: Properties

def __init__(
self,
Expand All @@ -2168,6 +2170,7 @@ def __init__(
limit: int | None = None,
catalog: Catalog | None = None,
table_identifier: Identifier | None = None,
table_config: Properties = EMPTY_DICT,
):
super().__init__(
table_metadata=table_metadata,
Expand All @@ -2181,6 +2184,7 @@ def __init__(
self.snapshot_id = snapshot_id
self.catalog = catalog
self.table_identifier = table_identifier
self.table_config = table_config

def snapshot(self) -> Snapshot | None:
if self.snapshot_id:
Expand Down Expand Up @@ -2425,7 +2429,7 @@ def _should_use_server_side_planning(self) -> bool:
"""Check if server-side scan planning should be used for this scan."""
if not self.catalog:
return False
return self.catalog.supports_server_side_planning()
return self.catalog.supports_server_side_planning(self.table_config)

def _plan_files_server_side(self) -> Iterable[FileScanTask]:
"""Plan files using REST server-side scan planning."""
Expand All @@ -2444,7 +2448,12 @@ def _plan_files_server_side(self) -> Iterable[FileScanTask]:
case_sensitive=self.case_sensitive,
)

return self.catalog.plan_scan(self.table_identifier, request)
result = self.catalog._plan_scan_result(self.table_identifier, request)
location = result.tasks[0].file.file_path if result.tasks else None
plan_io = self.catalog._file_io_from_plan(self.io.properties, result.storage_credentials, location)
if plan_io is not None:
self.io = plan_io
return result.tasks

def _plan_files_local(self) -> Iterable[FileScanTask]:
"""Plan files locally by reading manifests."""
Expand Down
54 changes: 54 additions & 0 deletions tests/catalog/test_rest.py
Original file line number Diff line number Diff line change
Expand Up @@ -2857,6 +2857,60 @@ def test_server_side_planning_enabled_from_server_config(self, rest_mock: Mocker

assert catalog.supports_server_side_planning() is True

def test_server_side_planning_enabled_by_table_config(self, rest_mock: Mocker) -> None:
catalog = RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN)

assert catalog.supports_server_side_planning() is False
assert catalog.supports_server_side_planning({"scan-planning-mode": ScanPlanningMode.SERVER.value}) is True

def test_server_side_planning_table_config_overrides_catalog_property(self, rest_mock: Mocker) -> None:
catalog = RestCatalog(
"rest",
uri=TEST_URI,
token=TEST_TOKEN,
**{"scan-planning-mode": ScanPlanningMode.SERVER.value},
)

assert catalog.supports_server_side_planning({"scan-planning-mode": ScanPlanningMode.CLIENT.value}) is False

def test_server_side_planning_table_config_ignored_when_endpoint_unsupported(self, requests_mock: Mocker) -> None:
requests_mock.get(
f"{TEST_URI}v1/config",
json={"defaults": {}, "overrides": {}},
status_code=200,
)
catalog = RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN)

assert catalog.supports_server_side_planning({"scan-planning-mode": ScanPlanningMode.SERVER.value}) is False

def test_server_side_planning_invalid_mode(self, rest_mock: Mocker) -> None:
catalog = RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN)

with pytest.raises(ValueError, match="Invalid scan-planning-mode: remote"):
catalog.supports_server_side_planning({"scan-planning-mode": "remote"})

def test_server_side_planning_invalid_catalog_mode_falls_back_to_default(self, rest_mock: Mocker) -> None:
catalog = RestCatalog(
"rest",
uri=TEST_URI,
token=TEST_TOKEN,
**{"scan-planning-mode": "servr"},
)

# Bad catalog config is ignored; default remains client-side planning.
assert catalog.supports_server_side_planning() is False

def test_server_side_planning_table_override_survives_invalid_catalog_mode(self, rest_mock: Mocker) -> None:
catalog = RestCatalog(
"rest",
uri=TEST_URI,
token=TEST_TOKEN,
**{"scan-planning-mode": "servr"},
)

assert catalog.supports_server_side_planning({"scan-planning-mode": ScanPlanningMode.SERVER.value}) is True
assert catalog.supports_server_side_planning({"scan-planning-mode": ScanPlanningMode.CLIENT.value}) is False

def test_supported_endpoint(self, requests_mock: Mocker) -> None:
requests_mock.get(
f"{TEST_URI}v1/config",
Expand Down
Loading
Loading