From 3506b43c0301613e3e5a52f6988f424b2c6f8ba7 Mon Sep 17 00:00:00 2001 From: Wyatt Alt Date: Wed, 28 Oct 2026 18:16:32 -0700 Subject: [PATCH] feat(python): accept input_schema or output_schema alone in @udf (#4731) Binding a nullable column needed an explicit input schema, and that forced spelling out the output schema too. Either schema can now be given alone; the other side is inferred from annotations. The udf docstring now says a bare annotation is non-nullable and `Optional[T]` binds a nullable column, and that a `bytes` input counts against the per-row input limit, with Blob V2 and `BlobReader` for large values. --- python/python/lancedb/functions.py | 63 +++++++++++-------- .../tests/test_first_class_function_slice2.py | 20 +++++- 2 files changed, 57 insertions(+), 26 deletions(-) diff --git a/python/python/lancedb/functions.py b/python/python/lancedb/functions.py index a7c85a984..ae40b4e3f 100644 --- a/python/python/lancedb/functions.py +++ b/python/python/lancedb/functions.py @@ -1109,8 +1109,6 @@ def _infer_signature( parameters = _callable_parameters(function) if method: parameters = parameters[1:] - if (input_schema is None) != (output_schema is None): - raise ValueError("input_schema and output_schema must be provided together") try: annotations = get_type_hints(function, include_extras=True) @@ -1164,14 +1162,26 @@ def _infer_signature( f"expected {expected!r}, got {actual!r}" ) inputs = tuple(input_parameter(field) for field in input_schema) - return FunctionSignature(inputs=inputs, output=_function_output(output_schema)) + else: + inputs = None + + missing = [] + if inputs is None: + missing += [p.name for p in parameters if p.name not in annotations] + if output_schema is None and "return" not in annotations: + missing.append("return") + if missing: + raise TypeError(f"missing Function annotations: {missing!r}") + if output_schema is not None: + output = _function_output(output_schema) + else: + output_type, output_nullable = _annotation_type(annotations["return"]) + if output_nullable: + raise ValueError("Function output must be non-nullable") + output = _function_output(pa.field("result", output_type, nullable=False)) + if inputs is not None: + return FunctionSignature(inputs=inputs, output=output) - missing = [ - parameter.name for parameter in parameters if parameter.name not in annotations - ] - if missing or "return" not in annotations: - names = missing + ([] if "return" in annotations else ["return"]) - raise TypeError(f"missing Function annotations: {names!r}") inputs = [] for parameter in parameters: annotation = annotations[parameter.name] @@ -1190,13 +1200,7 @@ def _infer_signature( inputs.append( input_parameter(pa.field(parameter.name, data_type, nullable=nullable)) ) - output_type, output_nullable = _annotation_type(annotations["return"]) - if output_nullable: - raise ValueError("Function output must be non-nullable") - return FunctionSignature( - inputs=tuple(inputs), - output=_function_output(pa.field("result", output_type, nullable=False)), - ) + return FunctionSignature(inputs=tuple(inputs), output=output) def _is_udf_decorator(node: ast.expr) -> bool: @@ -2061,11 +2065,20 @@ def udf( ): """Prepare a Python function or class for remote Function registration. - Input and output signatures are inferred from supported annotations. For - Arrow types annotations cannot express precisely, pass ``input_schema`` - and ``output_schema`` together. Scalar outputs must be non-nullable. Every - named-struct field may be nullable; Enterprise preserves the struct's - validity when the result is expanded into sibling columns. + Input and output signatures are inferred from supported annotations. A + parameter annotated ``T`` is non-nullable and binds only a non-nullable + column; annotate ``Optional[T]`` to bind a nullable one, and it receives + ``None`` for a null value. PyArrow fields are nullable unless declared + otherwise, so most columns need ``Optional``. For Arrow types annotations + cannot express precisely, pass ``input_schema``, ``output_schema``, or + both; the one omitted is inferred. Scalar outputs must be non-nullable. + Every named-struct field may be nullable; Enterprise preserves the + struct's validity when the result is expanded into sibling columns. + + **Large values.** A ``bytes`` parameter receives the whole value, which + is held in memory with its row and counts against the server's per-row + input limit. To pass large binary values, such as video, store them in a + Blob V2 column and annotate the parameter ``BlobReader``, which streams. **Functions and classes.** A decorated function is called once per row, or once per batch when its parameters are annotated ``pyarrow.Array`` or @@ -2125,12 +2138,12 @@ def udf( letters, digits, and underscores. input_schema : pyarrow.Schema, optional Explicit input fields in the exact order of the callable parameters. - Required for batch annotations. Must be provided together with - ``output_schema``. + Required for batch annotations. When omitted, inputs are inferred from + the parameter annotations. output_schema : pyarrow.DataType, pyarrow.Field, or pyarrow.Schema, optional Explicit scalar or named-struct output. Scalar outputs must be - non-nullable. Required for batch annotations. Must be provided together - with ``input_schema``. + non-nullable. Required for batch annotations. When omitted, the output + is inferred from the return annotation. pip : sequence of str, optional Pip requirements for the remote environment. conda : sequence of str, optional diff --git a/python/python/tests/test_first_class_function_slice2.py b/python/python/tests/test_first_class_function_slice2.py index 58afa9c7c..47c021b15 100644 --- a/python/python/tests/test_first_class_function_slice2.py +++ b/python/python/tests/test_first_class_function_slice2.py @@ -1495,7 +1495,7 @@ def test_annotation_and_explicit_schema_validation_fail_closed(): def nullable_output(value: int) -> Optional[int]: return value - with pytest.raises(ValueError, match="provided together"): + with pytest.raises(TypeError, match=r"missing Function annotations: \['return'\]"): @udf(input_schema=pa.schema([pa.field("value", pa.int64())])) def partial_schema(value): @@ -2062,3 +2062,21 @@ def test_a_function_namespace_path_must_be_a_list_of_segments(): list(db.list_functions(namespace_path="analytics")) assert state["requests"] == [] + + +def test_one_explicit_schema_infers_the_other_from_annotations(): + @udf(input_schema=pa.schema([pa.field("value", pa.int64())])) + def explicit_input(value) -> float: + return float(value or 0) + + @udf(output_schema=pa.field("result", pa.float32(), nullable=False)) + def explicit_output(value: Optional[int]): + return float(value or 0) + + for fn in (explicit_input, explicit_output): + signature = fn.registration_request.signature + assert [(p.name, p.arrow_type, p.nullable) for p in signature.inputs] == [ + ("value", "int64", True) + ] + assert explicit_input.registration_request.signature.output.arrow_type == "float64" + assert explicit_output.registration_request.signature.output.arrow_type == "float32"