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(