feat(python): expose Function batch size and row timeout controls (#4807)

Python UDF registration could not set `batch_size` or
`row_timeout_seconds`, and the Rust registration model discarded these
fields before submitting SQL. Expose both optional controls and preserve
them through registration, with invalid values rejected and unchanged
canonical bytes when the controls are omitted.

The timeout applies to each Python call: one row for scalar Functions
and one batch for vectorized Functions. Initialization keeps its
existing service deadline.

Validation at `f58e55c535a7df310db14d691b329402ca4baaf5`: 6 Rust
registration-contract tests, 2 outbound-SQL tests, and all 103 Python
slice2 tests passed after rebuilding the native extension. Rust
formatting and the repository Python CI checks (Ruff 0.9.9, python/
scope) passed. A real REST + Flight SQL server integration also passed
using the rebuilt SDK: `batch_size=2` limits observed calls in a direct
SELECT over 8 rows, initial materialized-view creation, and
append/refresh to 12 rows. This CPU host test verifies SDK-to-runtime
passthrough; it does not exercise GPU initialization.

Refs ENT-3246. The original GPU initialization and video-decoding
scenario still needs verification on a current deployment; this SDK
change does not close that investigation.

The additional repository-root style checks report existing
docs/web-tests findings (nine format files and 48 lint diagnostics); all
affected files are unchanged from main. These are outside the Python CI
job scope.

Rust compatibility: callers constructing `FunctionRegistrationRequest`
must initialize the two new optional fields (`None` preserves defaults);
the request no longer implements `Eq`.

The final head `798f463c9f68ab73e28c9a1ff06b5e4ee3e5d01f` adds
creation-guide documentation and integrates main’s namespace
repr/optional OCI embedding changes. The Function registration
implementation is unchanged. All 103 Python slice2 tests passed again on
EC2 using the unchanged native extension; Rust and server-bridge
evidence remains applicable.

python/python/lancedb/functions.py

@@ -86,6 +86,27 @@ def _validate_gpu_class_name(value: Any) -> str:
     return value
 
 
+def _validate_batch_size(value: Any) -> Optional[int]:
+    if value is None:
+        return None
+    if isinstance(value, bool) or not isinstance(value, int):
+        raise ValueError("batch_size must be an integer")
+    if not 1 <= value <= 2**32 - 1:
+        raise ValueError("batch_size must be between 1 and 4294967295")
+    return value
+
+
+def _validate_row_timeout(value: Any) -> Optional[float]:
+    if value is None:
+        return None
+    if isinstance(value, bool) or not isinstance(value, (int, float)):
+        raise ValueError("row_timeout_seconds must be a number")
+    value = float(value)
+    if not math.isfinite(value) or value <= 0:
+        raise ValueError("row_timeout_seconds must be positive and finite")
+    return value
+
+
 class _FrozenDict(dict):
     def _immutable(self, *args, **kwargs):
         raise TypeError("remote canonical values are immutable")
@@ -547,8 +568,20 @@ class FunctionRegistrationRequest(_RemoteValue):
     artifact: FunctionArtifactRequest
     signature: FunctionSignature
     runtime: PythonRuntimeSpec
+    batch_size: Optional[int] = None
+    row_timeout_seconds: Optional[float] = None
     secret_bindings: tuple[SecretBinding, ...] = ()
 
+    @field_validator("batch_size", mode="before")
+    @classmethod
+    def _validate_batch_size(cls, value):
+        return _validate_batch_size(value)
+
+    @field_validator("row_timeout_seconds", mode="before")
+    @classmethod
+    def _validate_row_timeout(cls, value):
+        return _validate_row_timeout(value)
+
 
 class FunctionVersionRef(_OpenRemoteValue):
     name: str
@@ -1983,6 +2016,8 @@ class UdfDefinition:
         pip: tuple[str, ...],
         env: Mapping[str, str],
         python_version: Optional[str],
+        batch_size: Optional[int] = None,
+        row_timeout_seconds: Optional[float] = None,
         gpu: bool = False,
         gpu_class: Optional[str] = None,
         conda: tuple[str, ...] = (),
@@ -2105,6 +2140,8 @@ class UdfDefinition:
             ),
             signature=signature,
             runtime=runtime,
+            batch_size=batch_size,
+            row_timeout_seconds=row_timeout_seconds,
         )
         # A class's own __dict__ holds its methods; copying it would shadow
         # this wrapper's attributes.
@@ -2199,6 +2236,8 @@ def udf(
     pip: tuple[str, ...] | list[str] = (),
     env: Optional[Mapping[str, str]] = None,
     python_version: Optional[str] = None,
+    batch_size: Optional[int] = None,
+    row_timeout_seconds: Optional[float] = None,
     gpu: bool = False,
     gpu_class: Optional[str] = None,
     conda: tuple[str, ...] | list[str] = (),
@@ -2218,6 +2257,8 @@ def udf(
     pip: tuple[str, ...] | list[str] = (),
     env: Optional[Mapping[str, str]] = None,
     python_version: Optional[str] = None,
+    batch_size: Optional[int] = None,
+    row_timeout_seconds: Optional[float] = None,
     gpu: bool = False,
     gpu_class: Optional[str] = None,
     conda: tuple[str, ...] | list[str] = (),
@@ -2339,6 +2380,15 @@ def udf(
         values that differ between bindings.
     python_version : str, optional
         Remote Python major/minor version. Defaults to the client version.
+    batch_size : int, optional
+        Maximum number of input rows supplied in one execution batch. Must be
+        between 1 and 4294967295. Remote execution may use smaller batches.
+        When omitted, the service chooses the batch size adaptively.
+    row_timeout_seconds : float, optional
+        Positive, finite deadline in seconds for each Python call: one row for
+        scalar Functions and one batch for vectorized Functions. This does not
+        govern class initialization. When omitted, the service applies its
+        default per-call budget.
     gpu : bool, default False
         Whether every remote execution requires a GPU. The execution platform
         selects one compatible GPU for each worker. The requirement is part of
@@ -2457,6 +2507,8 @@ def udf(
             pip=tuple(pip),
             env={} if env is None else env,
             python_version=python_version,
+            batch_size=batch_size,
+            row_timeout_seconds=row_timeout_seconds,
             gpu=gpu,
             gpu_class=gpu_class,
             conda=tuple(conda),

python/python/tests/test_e2e_remote_db.py

@@ -10,7 +10,7 @@ import lancedb
 import numpy as np
 import pyarrow as pa
 import pytest
-from lancedb import LanceDBConnection
+from lancedb import LanceDBConnection, udf
 from lancedb.index import FTS
 
 # TODO: setup integ test mark and script
@@ -74,6 +74,106 @@ def test_json_fts_against_server():
         conn.drop_table(name)
 
 
+@pytest.mark.skipif(
+    not os.environ.get("LANCEDB_URI"),
+    reason="Set LANCEDB_URI to a server with Python Functions enabled",
+)
+def test_function_execution_controls_through_sql_and_materialized_view():
+    """Exercise registered batch controls through the real remote runtime.
+
+    The Function returns the input batch length for every row. A cap of two must
+    therefore produce only one- or two-element observations in both a direct
+    SELECT and a materialized view refreshed after an append.
+    """
+    suffix = uuid4().hex[:12]
+    conn = lancedb.connect(
+        uri=os.environ["LANCEDB_URI"],
+        api_key=os.environ.get("LANCEDB_API_KEY"),
+        region=os.environ.get("LANCEDB_REGION", "us-east-1"),
+        host_override=os.environ.get("LANCEDB_HOST_OVERRIDE"),
+        sql_host_override=os.environ.get("LANCEDB_SQL_HOST_OVERRIDE"),
+    )
+    function_name = f"batch_probe_{suffix}"
+    source_name = f"batch_probe_source_{suffix}"
+    view_name = f"batch_probe_view_{suffix}"
+    registered_version = None
+
+    @udf(
+        name=function_name,
+        input_schema=pa.schema([pa.field("value", pa.int64(), nullable=False)]),
+        output_schema=pa.int64(),
+        batch_size=2,
+        row_timeout_seconds=15,
+    )
+    def observed_batch(value: pa.Array) -> pa.Array:
+        size = len(value)
+        return pa.array([size] * size, type=pa.int64())
+
+    def row_batch_sizes(table_name):
+        rows = (
+            conn.execute_query(
+                f"SELECT {function_name}(value) AS observed FROM {table_name}"
+            )
+            .read_all()
+            .to_pylist()
+        )
+        return [row["observed"] for row in rows]
+
+    def drop_table_if_present(table_name):
+        try:
+            conn.drop_table(table_name)
+        except Exception:
+            pass
+
+    try:
+        table = conn.create_table(
+            source_name,
+            pa.table(
+                {"value": pa.array(range(8), type=pa.int64())},
+                schema=pa.schema([pa.field("value", pa.int64(), nullable=False)]),
+            ),
+        )
+        registered_version = conn.create_function(observed_batch, namespace_path=None)
+
+        direct = row_batch_sizes(source_name)
+        assert len(direct) == 8
+        assert all(1 <= size <= 2 for size in direct)
+
+        view = conn.create_materialized_view(
+            view_name,
+            source=source_name,
+            select=[("observed", f"{function_name}(value)")],
+        )
+        initial_view = [
+            row["observed"]
+            for row in view.table.search().select(["observed"]).to_list()
+        ]
+        assert len(initial_view) == 8
+        assert all(1 <= size <= 2 for size in initial_view)
+
+        table.add(
+            pa.table(
+                {"value": pa.array(range(8, 12), type=pa.int64())},
+                schema=pa.schema([pa.field("value", pa.int64(), nullable=False)]),
+            )
+        )
+        view.refresh()
+        refreshed = [
+            row["observed"]
+            for row in view.table.search().select(["observed"]).to_list()
+        ]
+        assert len(refreshed) == 12
+        assert all(1 <= size <= 2 for size in refreshed)
+    finally:
+        drop_table_if_present(view_name)
+        drop_table_if_present(source_name)
+        if registered_version is not None:
+            try:
+                conn.drop_function(function_name, version=registered_version.version)
+            except Exception:
+                pass
+
+
 @pytest.mark.skip(reason="Need to set up a local server")
 def test_against_local_server():
     conn = LanceDBConnection("lancedb+http://localhost:10024")

python/python/tests/test_first_class_function_slice2.py

@@ -92,6 +92,57 @@ def test_scalar_udf_matches_shared_registration_golden_and_remains_callable():
     }
 
 
+def test_udf_execution_controls_match_shared_registration_golden():
+    controlled = udf(
+        pip=["numpy>=2"],
+        env={"MODE": "test"},
+        python_version="3.12",
+        batch_size=2,
+        row_timeout_seconds=0.75,
+    )(normalize_score.__wrapped__)
+    assert (
+        controlled.registration_request.to_canonical_json()
+        == (FIXTURES / "remote_function_controlled_registration_request.canonical.json")
+        .read_text()
+        .strip()
+    )
+
+
+def test_udf_execution_controls_apply_to_vectorized_functions():
+    @udf(
+        input_schema=pa.schema([pa.field("value", pa.int64(), nullable=False)]),
+        output_schema=pa.int64(),
+        batch_size=3,
+        row_timeout_seconds=1.25,
+    )
+    def identity(value: pa.Array) -> pa.Array:
+        return value
+
+    request = json.loads(identity.registration_request.to_canonical_json())
+    assert request["batch_size"] == 3
+    assert request["row_timeout_seconds"] == 1.25
+
+
+@pytest.mark.parametrize(
+    ("options", "field"),
+    [
+        ({"batch_size": 0}, "batch_size"),
+        ({"batch_size": 2**32}, "batch_size"),
+        ({"batch_size": True}, "batch_size"),
+        ({"row_timeout_seconds": 0}, "row_timeout_seconds"),
+        ({"row_timeout_seconds": -0.25}, "row_timeout_seconds"),
+        ({"row_timeout_seconds": float("inf")}, "row_timeout_seconds"),
+        ({"row_timeout_seconds": True}, "row_timeout_seconds"),
+    ],
+)
+def test_udf_execution_controls_require_their_positive_ranges(options, field):
+    def identity(value: int) -> int:
+        return value
+
+    with pytest.raises(ValueError, match=field):
+        udf(**options)(identity)
+
+
 @pytest.mark.parametrize(
     "name", ["examples.image_info", "image-info", "image_info.", "image_info-"]
 )