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`.

<!-- lance-gatekeeper-fix:v1 agent=2c6567ee3f128d3a0c3b9c68e36eca38
generation=1 -->

Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com>

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

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(