feat: Add models for rest scan planning

Author: geruh

pyiceberg/catalog/rest/scan_planning.py

@@ -0,0 +1,208 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from datetime import date, datetime, time
+from decimal import Decimal
+from typing import Annotated, Generic, Literal, TypeAlias, TypeVar
+from uuid import UUID
+
+from pydantic import Field, model_validator
+
+from pyiceberg.catalog.rest.response import ErrorResponseMessage
+from pyiceberg.expressions import BooleanExpression
+from pyiceberg.manifest import FileFormat
+from pyiceberg.typedef import IcebergBaseModel
+
+# Primitive types that can appear in partition values and bounds
+PrimitiveTypeValue: TypeAlias = bool | int | float | str | Decimal | UUID | date | time | datetime | bytes
+
+V = TypeVar("V")
+
+
+class KeyValueMap(IcebergBaseModel, Generic[V]):
+    """Map serialized as parallel key/value arrays for column statistics."""
+
+    keys: list[int] = Field(default_factory=list)
+    values: list[V] = Field(default_factory=list)
+
+    @model_validator(mode="after")
+    def _validate_lengths_match(self) -> KeyValueMap[V]:
+        if len(self.keys) != len(self.values):
+            raise ValueError(f"keys and values must have same length: {len(self.keys)} != {len(self.values)}")
+        return self
+
+    def to_dict(self) -> dict[int, V]:
+        """Convert to dictionary mapping field ID to value."""
+        return dict(zip(self.keys, self.values, strict=True))
+
+
+class CountMap(KeyValueMap[int]):
+    """Map of field IDs to counts."""
+
+
+class ValueMap(KeyValueMap[PrimitiveTypeValue]):
+    """Map of field IDs to primitive values (for lower/upper bounds)."""
+
+
+class StorageCredential(IcebergBaseModel):
+    """Storage credential for accessing content files."""
+
+    prefix: str = Field(description="Storage location prefix this credential applies to")
+    config: dict[str, str] = Field(default_factory=dict)
+
+
+class RESTContentFile(IcebergBaseModel):
+    """Base model for data and delete files from REST API."""
+
+    spec_id: int = Field(alias="spec-id")
+    partition: list[PrimitiveTypeValue] = Field(default_factory=list)
+    content: Literal["data", "position-deletes", "equality-deletes"]
+    file_path: str = Field(alias="file-path")
+    file_format: FileFormat = Field(alias="file-format")
+    file_size_in_bytes: int = Field(alias="file-size-in-bytes")
+    record_count: int = Field(alias="record-count")
+    key_metadata: str | None = Field(alias="key-metadata", default=None)
+    split_offsets: list[int] | None = Field(alias="split-offsets", default=None)
+    sort_order_id: int | None = Field(alias="sort-order-id", default=None)
+
+
+class RESTDataFile(RESTContentFile):
+    """Data file from REST API."""
+
+    content: Literal["data"] = Field(default="data")
+    first_row_id: int | None = Field(alias="first-row-id", default=None)
+    column_sizes: CountMap | None = Field(alias="column-sizes", default=None)
+    value_counts: CountMap | None = Field(alias="value-counts", default=None)
+    null_value_counts: CountMap | None = Field(alias="null-value-counts", default=None)
+    nan_value_counts: CountMap | None = Field(alias="nan-value-counts", default=None)
+    lower_bounds: ValueMap | None = Field(alias="lower-bounds", default=None)
+    upper_bounds: ValueMap | None = Field(alias="upper-bounds", default=None)
+
+
+class RESTPositionDeleteFile(RESTContentFile):
+    """Position delete file from REST API."""
+
+    content: Literal["position-deletes"] = Field(default="position-deletes")
+    referenced_data_file: str | None = Field(alias="referenced-data-file", default=None)
+    content_offset: int | None = Field(alias="content-offset", default=None)
+    content_size_in_bytes: int | None = Field(alias="content-size-in-bytes", default=None)
+
+
+class RESTEqualityDeleteFile(RESTContentFile):
+    """Equality delete file from REST API."""
+
+    content: Literal["equality-deletes"] = Field(default="equality-deletes")
+    equality_ids: list[int] | None = Field(alias="equality-ids", default=None)
+
+
+# Discriminated union for delete files
+RESTDeleteFile = Annotated[
+    RESTPositionDeleteFile | RESTEqualityDeleteFile,
+    Field(discriminator="content"),
+]
+
+
+class RESTFileScanTask(IcebergBaseModel):
+    """A file scan task from the REST server."""
+
+    data_file: RESTDataFile = Field(alias="data-file")
+    delete_file_references: list[int] | None = Field(alias="delete-file-references", default=None)
+    residual_filter: BooleanExpression | None = Field(alias="residual-filter", default=None)
+
+
+class ScanTasks(IcebergBaseModel):
+    """Container for scan tasks returned by the server."""
+
+    delete_files: list[RESTDeleteFile] = Field(alias="delete-files", default_factory=list)
+    file_scan_tasks: list[RESTFileScanTask] = Field(alias="file-scan-tasks", default_factory=list)
+    plan_tasks: list[str] = Field(alias="plan-tasks", default_factory=list)
+
+    @model_validator(mode="after")
+    def _validate_delete_file_references(self) -> ScanTasks:
+        # validate delete file references are in bounds
+        max_idx = len(self.delete_files) - 1
+        for task in self.file_scan_tasks:
+            for idx in task.delete_file_references or []:
+                if idx < 0 or idx > max_idx:
+                    raise ValueError(f"Invalid delete file reference: {idx} (valid range: 0-{max_idx})")
+
+        if self.delete_files and not self.file_scan_tasks:
+            raise ValueError("Invalid response: deleteFiles should only be returned with fileScanTasks that reference them")
+
+        return self
+
+
+class PlanCompleted(ScanTasks):
+    """Completed scan plan result."""
+
+    status: Literal["completed"] = "completed"
+    plan_id: str | None = Field(alias="plan-id", default=None)
+    storage_credentials: list[StorageCredential] | None = Field(alias="storage-credentials", default=None)
+
+
+class PlanSubmitted(IcebergBaseModel):
+    """Scan plan submitted, poll for completion."""
+
+    status: Literal["submitted"] = "submitted"
+    plan_id: str | None = Field(alias="plan-id", default=None)
+
+
+class PlanCancelled(IcebergBaseModel):
+    """Planning was cancelled."""
+
+    status: Literal["cancelled"] = "cancelled"
+
+
+class PlanFailed(IcebergBaseModel):
+    """Planning failed with error."""
+
+    status: Literal["failed"] = "failed"
+    error: ErrorResponseMessage
+
+
+PlanningResponse = Annotated[
+    PlanCompleted | PlanSubmitted | PlanCancelled | PlanFailed,
+    Field(discriminator="status"),
+]
+
+
+class PlanTableScanRequest(IcebergBaseModel):
+    """Request body for planning a REST scan."""
+
+    snapshot_id: int | None = Field(alias="snapshot-id", default=None)
+    select: list[str] | None = Field(default=None)
+    filter: BooleanExpression | None = Field(default=None)
+    case_sensitive: bool = Field(alias="case-sensitive", default=True)
+    use_snapshot_schema: bool = Field(alias="use-snapshot-schema", default=False)
+    start_snapshot_id: int | None = Field(alias="start-snapshot-id", default=None)
+    end_snapshot_id: int | None = Field(alias="end-snapshot-id", default=None)
+    stats_fields: list[str] | None = Field(alias="stats-fields", default=None)
+
+    @model_validator(mode="after")
+    def _validate_snapshot_fields(self) -> PlanTableScanRequest:
+        if self.start_snapshot_id is not None and self.end_snapshot_id is None:
+            raise ValueError("end-snapshot-id is required when start-snapshot-id is specified")
+        if self.snapshot_id is not None and self.start_snapshot_id is not None:
+            raise ValueError("Cannot specify both snapshot-id and start-snapshot-id")
+        return self
+
+
+class FetchScanTasksRequest(IcebergBaseModel):
+    """Request body for fetching scan tasks endpoint."""
+
+    plan_task: str = Field(alias="plan-task")