Repository: ibis-project/ibis

Issue:
feat: Actual type hints (not representative) for UDFs
### Is your feature request related to a problem?

Sorry that I'm making these issues on UDFs when I'd imagine you are about to release them :(

This is what I currently do:

```python
import pyarrow.compute as pc

@ibis.udf.scalar.pyarrow
def add_one(x: int) -> int:
    return pc.add(x, 1)
```

In the function, `x` isn't a python int, but is actually a pyarrow chunked array of type int. This both is confusing to me as the reader of this function, and it makes my editor sad.

### Describe the solution you'd like

Could we annotate in such a way that standard python tooling eg mypy is happy? Maybe using typing.Annotated

```python
import ibis
from pyarrow.lib import IntegerArray
from typing import Annotated
import pyarrow.compute as pc


@ibis.udf.scalar.pyarrow
def add_one(x: Annotated[IntegerArray, int]) -> Annotated[IntegerArray, int]:
    return pc.add(x, 1)

t = ibis.examples.penguins.fetch()
result = add_one(t.year)
```

### What version of ibis are you running?

main

### What backend(s) are you using, if any?

NA

### Code of Conduct

- [X] I agree to follow this project's Code of Conduct


Discussion:
Initial UDF support accepts type hints that are meant to be convenient for describing the element type of the inputs, not comprehensively compatible with mypy.

Is there some way we can have both?

I suppose we could add support for the `Annotated` versions of things for people who want better mypy compatibility.

I think this would ship after 7.0.
I think accepting either `int` or `Annotated[pyarrow_type, int]` would just be extra code, but should work fine? Perhaps you've seen this, but probably the most common lib that uses Annotated is [FastAPI](https://fastapi.tiangolo.com/tutorial/query-params-str-validations), perhaps there are some lessons there if we go down there.

But, thinking about this more, I might actually prefer if telling ibis about the input and output types didn't use python type hints at all. It is cute how we use the type hints language feature that python has, but were not quite using them accurately, and I would rather have a slightly less succinct API than a succinct but misleading one. Apologies if these were already considered with the UDF API, and I can read up on any discussion on them.

```python
# don't love the __returns=...
# but I'm spitballing here
@ibis.udf.scalar.pyarrow(x=int, __returns=int)
def add_one(x: IntegerArray) -> IntegerArray:
    return pc.add(x, 1)

# python type hints aren't required. They aren't wrong now for mypy!
# Same number of characters as the current API, I don't see the downside?
@ibis.udf.scalar.pyarrow(x=int, __returns=int)
def add_one(x):
    return pc.add(x, 1)

# switch to pandas implementation
# and the python type hints change (as they should)
# but the ibis hints don't change (as they shouldn't)
@ibis.udf.scalar.pandas(x=int, __returns=int)
def add_one(x: pd.Series) -> pd.Series:
    return x + 1

# For pyarrow types where the element type is known,
# we don't need to tell ibis at all, and we can maintain our own
# mapping of pyarrow -> ibis types
# For pandas UDFs, where the element type isn't included in the
# pd.Series class, then we always need to help ibis out.
@ibis.udf.scalar.pyarrow
def add_one(x: IntegerArray) -> IntegerArray:
    return pc.add(x, 1)

# can directly plug in 3rd party functions
# without needing to do a `def` block
from otherlib import process_series
ibis.udf.scalar.pandas(x=int, __returns=int)(process_series)

# or lambdas (OK this is silly I know)
# (is the name important to ibis???)
ibis.udf.scalar.pandas(x=int, __returns=int)(name="to_codes", lambda x: x.factorize[0])
```
I feel very strongly about this, is there anything I can do to move this forward before the next release (after which it will be harder to change the API)? Thanks so much, I appreciate the extra effort towards re-thinking about what you must imagine is an already-solved problem.
We are planning to release on Monday. UDF support is pretty new and is currently marked experimental, which means it's subject to breakage across non-major releases.

I don't think any changes we'd make to accommodate more mypy-accurate type hinting need to block the Monday release.
Another option might be

```
@ibis.udf.scalar.pyarrow(x=int, y=int, z=int).returns(float)
```
> (is the name important to ibis???)

Yes. We have to generate SQL for any calls to a UDF, so there needs to be name to use for the function call.
> @ibis.udf.scalar.pyarrow(x=int, y=int, z=int).returns(float)

I like that! Perhaps really embrace the builder pattern to be consistent between args and return types, to support possible per-implementor options such as [pyarrow's set_cpu_count()](https://arrow.apache.org/docs/python/api/misc.html), and to add a way to define the name separately if you use a lambda, or you want to override passed in callable's `__name__`. All of these are optional in the simple case.

```
@ibis.udf.scalar.pyarrow(n_cpus=2).consumes(x=int, y=int, z=int).returns(float).name("my_func")
def _private_name(...):
    ...
```

Resolving patch:
diff --git a/ibis/expr/operations/udf.py b/ibis/expr/operations/udf.py
--- a/ibis/expr/operations/udf.py
+++ b/ibis/expr/operations/udf.py
@@ -78,33 +78,47 @@ def _make_node(
         input_type: InputType,
         name: str | None = None,
         schema: str | None = None,
+        signature: tuple[tuple, Any] | None = None,
         **kwargs,
     ) -> type[S]:
         """Construct a scalar user-defined function that is built-in to the backend."""
 
-        annotations = typing.get_type_hints(fn)
-        if (return_annotation := annotations.pop("return", None)) is None:
-            raise exc.MissingReturnAnnotationError(fn)
+        if signature is None:
+            annotations = typing.get_type_hints(fn)
+            if (return_annotation := annotations.pop("return", None)) is None:
+                raise exc.MissingReturnAnnotationError(fn)
+            fields = {
+                arg_name: Argument(
+                    pattern=rlz.ValueOf(annotations.get(arg_name)),
+                    default=param.default,
+                    typehint=annotations.get(arg_name, Any),
+                )
+                for arg_name, param in inspect.signature(fn).parameters.items()
+            }
+        else:
+            arg_types, return_annotation = signature
+            arg_names = list(inspect.signature(fn).parameters)
+            fields = {
+                arg_name: Argument(pattern=rlz.ValueOf(typ), typehint=typ)
+                for arg_name, typ in zip(arg_names, arg_types)
+            }
 
         func_name = name if name is not None else fn.__name__
 
-        fields = {
-            arg_name: Argument(
-                pattern=rlz.ValueOf(annotations.get(arg_name)), default=param.default
-            )
-            for arg_name, param in inspect.signature(fn).parameters.items()
-        } | {
-            "dtype": dt.dtype(return_annotation),
-            "__input_type__": input_type,
-            # must wrap `fn` in a `property` otherwise `fn` is assumed to be a
-            # method
-            "__func__": property(fget=lambda _, fn=fn: fn),
-            "__config__": FrozenDict(kwargs),
-            "__udf_namespace__": schema,
-            "__module__": fn.__module__,
-            "__func_name__": func_name,
-            "__full_name__": ".".join(filter(None, (schema, func_name))),
-        }
+        fields.update(
+            {
+                "dtype": dt.dtype(return_annotation),
+                "__input_type__": input_type,
+                # must wrap `fn` in a `property` otherwise `fn` is assumed to be a
+                # method
+                "__func__": property(fget=lambda _, fn=fn: fn),
+                "__config__": FrozenDict(kwargs),
+                "__udf_namespace__": schema,
+                "__module__": fn.__module__,
+                "__func_name__": func_name,
+                "__full_name__": ".".join(filter(None, (schema, func_name))),
+            }
+        )
 
         return type(fn.__name__, (cls._base,), fields)
 
@@ -144,23 +158,30 @@ def builtin(
         *,
         name: str | None = None,
         schema: str | None = None,
+        signature: tuple[tuple[Any, ...], Any] | None = None,
         **kwargs: Any,
     ) -> Callable[[Callable], Callable[..., ir.Value]]:
         ...
 
     @util.experimental
     @classmethod
-    def builtin(cls, fn=None, *, name=None, schema=None, **kwargs):
+    def builtin(cls, fn=None, *, name=None, schema=None, signature=None, **kwargs):
         """Construct a scalar user-defined function that is built-in to the backend.
 
         Parameters
         ----------
         fn
-            The The function to wrap.
+            The function to wrap.
         name
             The name of the UDF in the backend if different from the function name.
         schema
             The schema in which the builtin function resides.
+        signature
+            An optional signature to use for the UDF. If present, should be a
+            tuple containing a tuple of argument types and a return type. For
+            example, a function taking an int and a float and returning a
+            string would be `((int, float), str)`. If not present, the argument
+            types will be derived from the wrapped function.
         kwargs
             Additional backend-specific configuration arguments for the UDF.
 
@@ -181,6 +202,7 @@ def builtin(cls, fn=None, *, name=None, schema=None, **kwargs):
             fn,
             name=name,
             schema=schema,
+            signature=signature,
             **kwargs,
         )
 
@@ -196,13 +218,14 @@ def python(
         *,
         name: str | None = None,
         schema: str | None = None,
+        signature: tuple[tuple[Any, ...], Any] | None = None,
         **kwargs: Any,
     ) -> Callable[[Callable], Callable[..., ir.Value]]:
         ...
 
     @util.experimental
     @classmethod
-    def python(cls, fn=None, *, name=None, schema=None, **kwargs):
+    def python(cls, fn=None, *, name=None, schema=None, signature=None, **kwargs):
         """Construct a **non-vectorized** scalar user-defined function that accepts Python scalar values as inputs.
 
         ::: {.callout-warning collapse="true"}
@@ -221,11 +244,17 @@ def python(cls, fn=None, *, name=None, schema=None, **kwargs):
         Parameters
         ----------
         fn
-            The The function to wrap.
+            The function to wrap.
         name
             The name of the UDF in the backend if different from the function name.
         schema
             The schema in which to create the UDF.
+        signature
+            An optional signature to use for the UDF. If present, should be a
+            tuple containing a tuple of argument types and a return type. For
+            example, a function taking an int and a float and returning a
+            string would be `((int, float), str)`. If not present, the argument
+            types will be derived from the wrapped function.
         kwargs
             Additional backend-specific configuration arguments for the UDF.
 
@@ -251,6 +280,7 @@ def python(cls, fn=None, *, name=None, schema=None, **kwargs):
             fn,
             name=name,
             schema=schema,
+            signature=signature,
             **kwargs,
         )
 
@@ -266,23 +296,30 @@ def pandas(
         *,
         name: str | None = None,
         schema: str | None = None,
+        signature: tuple[tuple[Any, ...], Any] | None = None,
         **kwargs: Any,
     ) -> Callable[[Callable], Callable[..., ir.Value]]:
         ...
 
     @util.experimental
     @classmethod
-    def pandas(cls, fn=None, *, name=None, schema=None, **kwargs):
+    def pandas(cls, fn=None, *, name=None, schema=None, signature=None, **kwargs):
         """Construct a **vectorized** scalar user-defined function that accepts pandas Series' as inputs.
 
         Parameters
         ----------
         fn
-            The The function to wrap.
+            The function to wrap.
         name
             The name of the UDF in the backend if different from the function name.
         schema
             The schema in which to create the UDF.
+        signature
+            An optional signature to use for the UDF. If present, should be a
+            tuple containing a tuple of argument types and a return type. For
+            example, a function taking an int and a float and returning a
+            string would be `((int, float), str)`. If not present, the argument
+            types will be derived from the wrapped function.
         kwargs
             Additional backend-specific configuration arguments for the UDF.
 
@@ -310,6 +347,7 @@ def pandas(cls, fn=None, *, name=None, schema=None, **kwargs):
             fn,
             name=name,
             schema=schema,
+            signature=signature,
             **kwargs,
         )
 
@@ -325,23 +363,30 @@ def pyarrow(
         *,
         name: str | None = None,
         schema: str | None = None,
+        signature: tuple[tuple[Any, ...], Any] | None = None,
         **kwargs: Any,
     ) -> Callable[[Callable], Callable[..., ir.Value]]:
         ...
 
     @util.experimental
     @classmethod
-    def pyarrow(cls, fn=None, *, name=None, schema=None, **kwargs):
+    def pyarrow(cls, fn=None, *, name=None, schema=None, signature=None, **kwargs):
         """Construct a **vectorized** scalar user-defined function that accepts PyArrow Arrays as input.
 
         Parameters
         ----------
         fn
-            The The function to wrap.
+            The function to wrap.
         name
             The name of the UDF in the backend if different from the function name.
         schema
             The schema in which to create the UDF.
+        signature
+            An optional signature to use for the UDF. If present, should be a
+            tuple containing a tuple of argument types and a return type. For
+            example, a function taking an int and a float and returning a
+            string would be `((int, float), str)`. If not present, the argument
+            types will be derived from the wrapped function.
         kwargs
             Additional backend-specific configuration arguments for the UDF.
 
@@ -368,6 +413,7 @@ def pyarrow(cls, fn=None, *, name=None, schema=None, **kwargs):
             fn,
             name=name,
             schema=schema,
+            signature=signature,
             **kwargs,
         )
 
@@ -389,23 +435,30 @@ def builtin(
         *,
         name: str | None = None,
         schema: str | None = None,
+        signature: tuple[tuple[Any, ...], Any] | None = None,
         **kwargs: Any,
     ) -> Callable[[Callable], Callable[..., ir.Value]]:
         ...
 
     @util.experimental
     @classmethod
-    def builtin(cls, fn=None, *, name=None, schema=None, **kwargs):
+    def builtin(cls, fn=None, *, name=None, schema=None, signature=None, **kwargs):
         """Construct an aggregate user-defined function that is built-in to the backend.
 
         Parameters
         ----------
         fn
-            The The function to wrap.
+            The function to wrap.
         name
             The name of the UDF in the backend if different from the function name.
         schema
             The schema in which the builtin function resides.
+        signature
+            An optional signature to use for the UDF. If present, should be a
+            tuple containing a tuple of argument types and a return type. For
+            example, a function taking an int and a float and returning a
+            string would be `((int, float), str)`. If not present, the argument
+            types will be derived from the wrapped function.
         kwargs
             Additional backend-specific configuration arguments for the UDF.
 
@@ -427,5 +480,6 @@ def builtin(cls, fn=None, *, name=None, schema=None, **kwargs):
             fn,
             name=name,
             schema=schema,
+            signature=signature,
             **kwargs,
         )


Test patch:
diff --git a/ibis/tests/expr/test_udf.py b/ibis/tests/expr/test_udf.py
--- a/ibis/tests/expr/test_udf.py
+++ b/ibis/tests/expr/test_udf.py
@@ -83,6 +83,50 @@ def test_vectorized_udf_operations(table, klass, output_type):
         )
 
 
+@pytest.mark.parametrize(
+    "dec",
+    [
+        pytest.param(ibis.udf.scalar.builtin, id="scalar-builtin"),
+        pytest.param(ibis.udf.scalar.pandas, id="scalar-pandas"),
+        pytest.param(ibis.udf.scalar.pyarrow, id="scalar-pyarrow"),
+        pytest.param(ibis.udf.scalar.python, id="scalar-python"),
+        pytest.param(ibis.udf.agg.builtin, id="agg-builtin"),
+    ],
+)
+def test_udf_from_annotations(dec, table):
+    @dec
+    def myfunc(x: int, y: str) -> float:
+        ...
+
+    assert myfunc(table.a, table.b).type().is_floating()
+
+    with pytest.raises(ValidationError):
+        # Wrong arg types
+        myfunc(table.b, table.a)
+
+
+@pytest.mark.parametrize(
+    "dec",
+    [
+        pytest.param(ibis.udf.scalar.builtin, id="scalar-builtin"),
+        pytest.param(ibis.udf.scalar.pandas, id="scalar-pandas"),
+        pytest.param(ibis.udf.scalar.pyarrow, id="scalar-pyarrow"),
+        pytest.param(ibis.udf.scalar.python, id="scalar-python"),
+        pytest.param(ibis.udf.agg.builtin, id="agg-builtin"),
+    ],
+)
+def test_udf_from_sig(dec, table):
+    @dec(signature=((int, str), float))
+    def myfunc(x, y):
+        ...
+
+    assert myfunc(table.a, table.b).type().is_floating()
+
+    with pytest.raises(ValidationError):
+        # Wrong arg types
+        myfunc(table.b, table.a)
+
+
 @pytest.mark.parametrize(
     "dec",
     [