From 04e509278e5039c686ff3dbe57ff62c8617edfb1 Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Wed, 28 Oct 2026 16:37:57 -0700 Subject: [PATCH] fix(python): clarify schemas for batch UDF annotations (#4727) `@udf` documented batch annotations without explaining that `pa.Array` and `pa.RecordBatch` do not specify concrete Arrow element types or field schemas. Document the requirement for both explicit schemas (including class `__call__` methods), add an array example, and make inference errors point callers to `input_schema` and `output_schema`. Fixes #4721 Validation: - Five new regression cases pass; all four diagnostic cases fail with the original implementation. - Function packaging, first-class Function slice 1/2 tests, and UDF doctests: 168 passed, 3 failed. This includes the upstream BlobReader annotation coverage. The three existing packaging tests fail while importing temporary modules, before UDF registration; all three failures also reproduce with the unchanged source and tests at `057a39a1`. - Rebuilt the native extension in the uv-managed tests/dev environment. - Python lint and formatting checks pass (`ruff check python`, `ruff format --check python`); `cargo fmt --all` also passes without changes. - Repository-wide `ruff check .` reports 48 existing errors under `docs/web-tests`. Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> --- python/python/lancedb/functions.py | 27 ++++++++++-- .../python/tests/test_function_packaging.py | 41 +++++++++++++++++++ 2 files changed, 64 insertions(+), 4 deletions(-) diff --git a/python/python/lancedb/functions.py b/python/python/lancedb/functions.py index a3e56247a..a7c85a984 100644 --- a/python/python/lancedb/functions.py +++ b/python/python/lancedb/functions.py @@ -1000,6 +1000,11 @@ def _annotation_type(annotation: Any) -> tuple[pa.DataType, bool]: if isinstance(annotation, pa.DataType): return annotation, nullable + if annotation is pa.Array or annotation is pa.RecordBatch: + raise TypeError( + f"cannot infer Arrow types from Function annotation {annotation!r}; " + "provide input_schema and output_schema together for batch callables" + ) if annotation is bool: return pa.bool_(), nullable if annotation is int: @@ -2064,8 +2069,11 @@ def udf( **Functions and classes.** A decorated function is called once per row, or once per batch when its parameters are annotated ``pyarrow.Array`` or - it takes one ``pyarrow.RecordBatch``. A decorated class is the same - callable with state: each remote instance calls ``__init__`` once, calls + it takes one ``pyarrow.RecordBatch``. These batch annotations require both + ``input_schema`` and ``output_schema`` explicitly: the annotations do not + specify Arrow element types or fields. This also applies to a class's + ``__call__`` method. A decorated class is the same callable with state: + each remote instance calls ``__init__`` once, calls ``__call__`` for every row or batch of every input it processes, and calls ``close()``, if defined, once when it retires. Load models, open clients, and build rate limiters in ``__init__``; every batch of the instance then @@ -2117,10 +2125,12 @@ def udf( letters, digits, and underscores. input_schema : pyarrow.Schema, optional Explicit input fields in the exact order of the callable parameters. - Must be provided together with ``output_schema``. + Required for batch annotations. Must be provided together with + ``output_schema``. output_schema : pyarrow.DataType, pyarrow.Field, or pyarrow.Schema, optional Explicit scalar or named-struct output. Scalar outputs must be - non-nullable. Must be provided together with ``input_schema``. + non-nullable. Required for batch annotations. Must be provided together + with ``input_schema``. pip : sequence of str, optional Pip requirements for the remote environment. conda : sequence of str, optional @@ -2158,6 +2168,15 @@ def udf( ... return value * 2 >>> score(1.5) 3.0 + >>> import pyarrow as pa + >>> @udf( + ... input_schema=pa.schema([pa.field("value", pa.int64(), nullable=False)]), + ... output_schema=pa.int64(), + ... ) + ... def batch_identity(value: pa.Array) -> pa.Array: + ... return value + >>> batch_identity(pa.array([1, 2, 3])).to_pylist() + [1, 2, 3] >>> @udf(pip=["cupy-cuda12x"], gpu=True) ... def gpu_score(value: int) -> int: ... return value * 2 diff --git a/python/python/tests/test_function_packaging.py b/python/python/tests/test_function_packaging.py index fe86de14e..a61588777 100644 --- a/python/python/tests/test_function_packaging.py +++ b/python/python/tests/test_function_packaging.py @@ -127,6 +127,47 @@ def test_function_without_initialization_keeps_its_wire_form(): assert request["artifact"]["kind"] == "python_callable" +@pytest.mark.parametrize("annotation", ["Array", "RecordBatch"]) +@pytest.mark.parametrize("kind", ["function", "class"]) +def test_batch_annotations_require_explicit_schemas(annotation, kind, helper_package): + if kind == "class": + source = f""" + import pyarrow as pa + + class identity: + def __call__(self, value: pa.{annotation}) -> pa.{annotation}: + return value + """ + else: + source = f""" + import pyarrow as pa + + def identity(value: pa.{annotation}) -> pa.{annotation}: + return value + """ + module = helper_package("fx_batch_annotations", {"__init__.py": source}) + with pytest.raises(TypeError, match="input_schema and output_schema"): + udf(module.identity) + + +def test_array_batch_annotations_accept_explicit_schemas(): + @udf( + input_schema=pa.schema([pa.field("value", pa.int64(), nullable=False)]), + output_schema=pa.int64(), + ) + def identity(value: pa.Array) -> pa.Array: + return value + + signature = identity.registration_request.signature + assert [(p.name, p.arrow_type, p.nullable) for p in signature.inputs] == [ + ("value", "int64", False) + ] + assert signature.output.arrow_type == "int64" + assert signature.output.nullable is False + values = pa.array([1, 2, 3]) + assert identity(values).equals(values) + + @pytest.mark.parametrize("class_udf", [False, True]) @pytest.mark.parametrize("postponed", [False, True]) @pytest.mark.parametrize(