Add tests and fix models

Author: geruh

pyiceberg/catalog/rest/__init__.py

@@ -424,7 +424,7 @@ def _plan_table_scan(self, identifier: str | Identifier, request: PlanTableScanR
         self._check_endpoint(Capability.V1_SUBMIT_TABLE_SCAN_PLAN)
         response = self._session.post(
             self.url(Endpoints.plan_table_scan, prefixed=True, **self._split_identifier_for_path(identifier)),
-            json=request.model_dump(by_alias=True, exclude_none=True),
+            data=request.model_dump_json(by_alias=True, exclude_none=True).encode(UTF8),
         )
         try:
             response.raise_for_status()
@@ -451,7 +451,7 @@ def _fetch_scan_tasks(self, identifier: str | Identifier, plan_task: str) -> Sca
         request = FetchScanTasksRequest(plan_task=plan_task)
         response = self._session.post(
             self.url(Endpoints.fetch_scan_tasks, prefixed=True, **self._split_identifier_for_path(identifier)),
-            json=request.model_dump(by_alias=True),
+            data=request.model_dump_json(by_alias=True).encode(UTF8),
         )
         try:
             response.raise_for_status()

pyiceberg/catalog/rest/scan_planning.py

@@ -24,7 +24,7 @@
 from pydantic import Field, model_validator
 
 from pyiceberg.catalog.rest.response import ErrorResponseMessage
-from pyiceberg.expressions import BooleanExpression
+from pyiceberg.expressions import BooleanExpression, SerializableBooleanExpression
 from pyiceberg.manifest import DataFileContent, FileFormat
 from pyiceberg.typedef import IcebergBaseModel
 
@@ -192,7 +192,7 @@ class PlanTableScanRequest(IcebergBaseModel):
 
     snapshot_id: int | None = Field(alias="snapshot-id", default=None)
     select: list[str] | None = Field(default=None)
-    filter: BooleanExpression | None = Field(default=None)
+    filter: SerializableBooleanExpression | 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)

tests/integration/test_rest_scan_planning_integration.py

@@ -15,25 +15,61 @@
 # specific language governing permissions and limitations
 # under the License.
 # pylint:disable=redefined-outer-name
+from datetime import date, datetime, time, timedelta, timezone
+from decimal import Decimal
 from typing import Any
+from uuid import uuid4
 
 import pyarrow as pa
 import pytest
 from pyspark.sql import SparkSession
 
 from pyiceberg.catalog import Catalog, load_catalog
 from pyiceberg.catalog.rest import RestCatalog
-from pyiceberg.expressions import And, BooleanExpression, EqualTo, GreaterThan, LessThan
+from pyiceberg.exceptions import NoSuchTableError
+from pyiceberg.expressions import (
+    And,
+    BooleanExpression,
+    EqualTo,
+    GreaterThan,
+    GreaterThanOrEqual,
+    In,
+    IsNull,
+    LessThan,
+    LessThanOrEqual,
+    Not,
+    NotEqualTo,
+    NotIn,
+    NotNull,
+    Or,
+    StartsWith,
+)
 from pyiceberg.partitioning import PartitionField, PartitionSpec
 from pyiceberg.schema import Schema
 from pyiceberg.table import ALWAYS_TRUE, Table
-from pyiceberg.transforms import IdentityTransform
-from pyiceberg.types import LongType, NestedField, StringType
+from pyiceberg.transforms import (
+    IdentityTransform,
+)
+from pyiceberg.types import (
+    BinaryType,
+    BooleanType,
+    DateType,
+    DecimalType,
+    DoubleType,
+    FixedType,
+    LongType,
+    NestedField,
+    StringType,
+    TimestampType,
+    TimestamptzType,
+    TimeType,
+    UUIDType,
+)
 
 
 @pytest.fixture(scope="session")
 def scan_catalog() -> Catalog:
-    return load_catalog(
+    catalog = load_catalog(
         "local",
         **{
             "type": "rest",
@@ -44,13 +80,15 @@ def scan_catalog() -> Catalog:
             "rest-scan-planning-enabled": "true",
         },
     )
+    catalog.create_namespace_if_not_exists("default")
+    return catalog
 
 
 def recreate_table(catalog: Catalog, identifier: str, **kwargs: Any) -> Table:
     """Drop table if exists and create a new one."""
     try:
         catalog.drop_table(identifier)
-    except Exception:
+    except NoSuchTableError:
         pass
     return catalog.create_table(identifier, **kwargs)
 
@@ -186,3 +224,124 @@ def test_rest_scan_with_partitioning(scan_catalog: RestCatalog, session_catalog:
         )
     finally:
         scan_catalog.drop_table(identifier)
+
+
+@pytest.mark.integration
+def test_rest_scan_primitive_types(scan_catalog: RestCatalog, session_catalog: Catalog) -> None:
+    identifier = "default.test_primitives"
+
+    schema = Schema(
+        NestedField(1, "bool_col", BooleanType()),
+        NestedField(2, "long_col", LongType()),
+        NestedField(3, "double_col", DoubleType()),
+        NestedField(4, "decimal_col", DecimalType(10, 2)),
+        NestedField(5, "string_col", StringType()),
+        NestedField(6, "date_col", DateType()),
+        NestedField(7, "time_col", TimeType()),
+        NestedField(8, "timestamp_col", TimestampType()),
+        NestedField(9, "timestamptz_col", TimestamptzType()),
+        NestedField(10, "uuid_col", UUIDType()),
+        NestedField(11, "fixed_col", FixedType(16)),
+        NestedField(12, "binary_col", BinaryType()),
+    )
+
+    table = recreate_table(scan_catalog, identifier, schema=schema)
+
+    now = datetime.now()
+    now_tz = datetime.now(tz=timezone.utc)
+    today = date.today()
+    uuid1, uuid2, uuid3 = uuid4(), uuid4(), uuid4()
+
+    arrow_table = pa.Table.from_pydict(
+        {
+            "bool_col": [True, False, True],
+            "long_col": [100, 200, 300],
+            "double_col": [1.11, 2.22, 3.33],
+            "decimal_col": [Decimal("1.23"), Decimal("4.56"), Decimal("7.89")],
+            "string_col": ["a", "b", "c"],
+            "date_col": [today, today - timedelta(days=1), today - timedelta(days=2)],
+            "time_col": [time(8, 30, 0), time(12, 0, 0), time(18, 45, 30)],
+            "timestamp_col": [now, now - timedelta(hours=1), now - timedelta(hours=2)],
+            "timestamptz_col": [now_tz, now_tz - timedelta(hours=1), now_tz - timedelta(hours=2)],
+            "uuid_col": [uuid1.bytes, uuid2.bytes, uuid3.bytes],
+            "fixed_col": [b"0123456789abcdef", b"abcdef0123456789", b"fedcba9876543210"],
+            "binary_col": [b"hello", b"world", b"test"],
+        },
+        schema=schema.as_arrow(),
+    )
+    table.append(arrow_table)
+
+    try:
+        _assert_remote_scan_matches_local_scan(table, session_catalog, identifier)
+    finally:
+        scan_catalog.drop_table(identifier)
+
+
+@pytest.mark.integration
+def test_rest_scan_complex_filters(scan_catalog: RestCatalog, session_catalog: Catalog) -> None:
+    identifier = "default.test_complex_filters"
+
+    schema = Schema(
+        NestedField(1, "id", LongType()),
+        NestedField(2, "name", StringType()),
+        NestedField(3, "value", LongType()),
+        NestedField(4, "optional", StringType(), required=False),
+    )
+
+    table = recreate_table(scan_catalog, identifier, schema=schema)
+
+    table.append(
+        pa.Table.from_pydict(
+            {
+                "id": list(range(1, 21)),
+                "name": [f"item_{i}" for i in range(1, 21)],
+                "value": [i * 100 for i in range(1, 21)],
+                "optional": [None if i % 3 == 0 else f"opt_{i}" for i in range(1, 21)],
+            }
+        )
+    )
+
+    try:
+        filters = [
+            EqualTo("id", 10),
+            NotEqualTo("id", 10),
+            GreaterThan("value", 1000),
+            GreaterThanOrEqual("value", 1000),
+            LessThan("value", 500),
+            LessThanOrEqual("value", 500),
+            In("id", [1, 5, 10, 15]),
+            NotIn("id", [1, 5, 10, 15]),
+            IsNull("optional"),
+            NotNull("optional"),
+            StartsWith("name", "item_1"),
+            And(GreaterThan("id", 5), LessThan("id", 15)),
+            Or(EqualTo("id", 1), EqualTo("id", 20)),
+            Not(EqualTo("id", 10)),
+        ]
+
+        for filter_expr in filters:
+            _assert_remote_scan_matches_local_scan(table, session_catalog, identifier, row_filter=filter_expr)
+    finally:
+        scan_catalog.drop_table(identifier)
+
+
+@pytest.mark.integration
+def test_rest_scan_empty_table(scan_catalog: RestCatalog, session_catalog: Catalog) -> None:
+    identifier = "default.test_empty_table"
+
+    schema = Schema(
+        NestedField(1, "id", LongType()),
+        NestedField(2, "data", StringType()),
+    )
+
+    table = recreate_table(scan_catalog, identifier, schema=schema)
+
+    try:
+        rest_tasks = list(table.scan().plan_files())
+        local_table = session_catalog.load_table(identifier)
+        local_tasks = list(local_table.scan().plan_files())
+
+        assert len(rest_tasks) == 0
+        assert len(local_tasks) == 0
+    finally:
+        scan_catalog.drop_table(identifier)