Repository: ibis-project/ibis

Issue:
bug: Ibis UDFs don't understand optional type
### What happened?

I noticed that ibis UDFs don't support nullable type hints. Even though the UDFs do seem to be working fine with nullable input, outputs.

This makes mypy complain when it detects that a nullable type is incorrectly types. 

Additionally, it creates issues when constructing types programmatically (see my struct example) which does not allow for Union type. (For the context, we have an internal system which creates some spark types on the fly, now we want to port those to create ibis types, but Union does not work). 

Some examples:

```py
# optional input 
@ibis.udf.scalar.python
def udf(x: int | None) -> int:
     return (None if x is None else 1)
SignatureValidationError: ValueOf(int | None) has failed due to the following errors:
  `dtype`: int | None of type <class 'types.UnionType'> is not either None or coercible to a DataType

Expected signature: ValueOf(dtype: Optional[DataType] = None)


# optional output
@ibis.udf.scalar.python
def udf(x: int) -> int | None:
     return (None if x is None else 1)

TypeError: Value int | None is not a valid datatype
```

```py
>>> from typing import Optional
>>> import ibis
>>> ibis.expr.datatypes.Struct({"a": Optional[str]})
Expected signature: Struct(fields: FrozenOrderedDict[str, DataType], nullable: bool = True)
>>> ibis.expr.datatypes.Struct({"a": str | None})
Expected signature: Struct(fields: FrozenOrderedDict[str, DataType], nullable: bool = True)
```



Interestingly, both None type input and output actually does work even without the correct type hints: 

```py
# input 
@ibis.udf.scalar.python
def udf(x: str) -> str:
     return (None if x == "1" else x)


df = pd.DataFrame({"a": ["1", "2",  "3"]})
sdf = spark.createDataFrame(df)
sdf.createOrReplaceTempView("df")

con = ibis.pyspark.connect(spark)
tbl = con.table("df")

print(tbl.mutate(out=udf(_.a)).execute())

# output
@ibis.udf.scalar.python
def udf(x: int) -> int:
     return (None if x > 2 else 1)

df = pd.DataFrame({"a": [1, 2, 3]})
sdf = spark.createDataFrame(df)
sdf.createOrReplaceTempView("df")

con = ibis.pyspark.connect(spark)
tbl = con.table("df")
print(tbl.mutate(out=udf(_.a)).execute())

```


### What version of ibis are you using?

10.6.0

### What backend(s) are you using, if any?

Spark

### Relevant log output

```sh

```

### Code of Conduct

- [x] I agree to follow this project's Code of Conduct


Discussion:
I just realized that there is internal support for nullable types, it's just nullable=True by default. This is why udfs work in my examples above. While convenient, the nlulable default is non that intuitive 🤔 
This was a conscious choice. Most systems (ClickHouse being perhaps a notable exception) default to handling nulls in their type system.

When this API was designed the rationale was that since Ibis's type system defaults to nullable, UDF signature input types should also do this. Additionally, it cuts down on typing and verbosity when defining a UDF.

The cost of these things was inconsistency with the broader typing ecosystem.
Would it be ok to make the system understand Optional for the sake of mypy and for the nitty-gritty among us who want to be everywhere explicit?
I think #11506 should remove the type errors.

If you want your UDF to output a non-null dtype, you will still need to use the `signature` arg of the udf:

```python
@ibis.udf.scalar.python(signature=(("!string",), "!int64"))
def udf(s: str) -> int:
    return 5
```



Patch:
diff --git a/ibis/common/typing.py b/ibis/common/typing.py
index be67bfa1f5a9..4730bfbc1764 100644
--- a/ibis/common/typing.py
+++ b/ibis/common/typing.py
@@ -24,7 +24,7 @@
 else:
     from typing_extensions import TypeAlias
 
-    UnionType = object()
+    UnionType = type(Union[int, str])
     _ClassInfo: TypeAlias = Union[type, tuple["_ClassInfo", ...]]
 
 
diff --git a/ibis/expr/datatypes/core.py b/ibis/expr/datatypes/core.py
index b29973b7cf04..43889aadcc94 100644
--- a/ibis/expr/datatypes/core.py
+++ b/ibis/expr/datatypes/core.py
@@ -16,6 +16,7 @@
     NamedTuple,
     Optional,
     TypeVar,
+    Union,
     get_type_hints,
     overload,
 )
@@ -30,6 +31,7 @@
 from ibis.common.grounds import Concrete, Singleton
 from ibis.common.patterns import Between, Coercible, CoercionError
 from ibis.common.temporal import IntervalUnit, TimestampUnit
+from ibis.common.typing import UnionType
 
 if TYPE_CHECKING:
     import numpy as np
@@ -245,6 +247,18 @@ def from_string(cls, value: str, nullable: bool = True) -> Self:
 
     @classmethod
     def from_typehint(cls, typ, nullable=True) -> Self:
+        # All types are nullable by default
+        # (eg `ibis.dtype(int)` results in `Int64(nullable=True)`),
+        # so if we are handed Union[None, T], simplify it to T.
+        #
+        # Depending on if you use `Union[A, B]` vs `A | B`, the type differs
+        #
+        # with Optional, you get __origin__ and with | syntax you get a UnionType
+        if getattr(typ, "__origin__", None) is Union or isinstance(typ, UnionType):
+            not_nones = [t for t in get_args(typ) if t is not type(None)]
+            if len(not_nones) == 1:
+                typ = not_nones[0]
+
         origin_type = get_origin(typ)
 
         if origin_type is None:


Test patch:
diff --git a/ibis/expr/datatypes/tests/test_core.py b/ibis/expr/datatypes/tests/test_core.py
index 73927395367a..900f93065002 100644
--- a/ibis/expr/datatypes/tests/test_core.py
+++ b/ibis/expr/datatypes/tests/test_core.py
@@ -2,9 +2,10 @@
 
 import datetime  # noqa: TC003
 import decimal  # noqa: TC003
+import sys
 import uuid  # noqa: TC003
 from dataclasses import dataclass
-from typing import Annotated, NamedTuple
+from typing import Annotated, NamedTuple, Optional, Union
 
 import pytest
 from pytest import param
@@ -56,14 +57,33 @@ def test_dtype(spec, expected):
 @pytest.mark.parametrize(
     ("args", "expected"),
     [
-        ((int,), dt.Int64(nullable=True)),
-        ((int, False), dt.Int64(nullable=False)),
-        (("!int",), dt.Int64(nullable=False)),
-        (("!int", True), dt.Int64(nullable=False)),  # "!" overrides `nullable`
+        (lambda: (int,), dt.Int64(nullable=True)),
+        (lambda: (int, False), dt.Int64(nullable=False)),
+        (lambda: (Union[int, None],), dt.Int64(nullable=True)),
+        (lambda: (Union[int, None], False), dt.Int64(nullable=False)),
+        (lambda: (Optional[int],), dt.Int64(nullable=True)),
+        (lambda: (Optional[int], False), dt.Int64(nullable=False)),
+        pytest.param(
+            lambda: (int | None,),
+            dt.Int64(nullable=True),
+            marks=pytest.mark.xfail(sys.version_info < (3, 10), reason="python 3.9"),
+        ),
+        pytest.param(
+            lambda: (int | None, False),
+            dt.Int64(nullable=False),
+            marks=pytest.mark.xfail(sys.version_info < (3, 10), reason="python 3.9"),
+        ),
+        (lambda: ("!int",), dt.Int64(nullable=False)),
+        (lambda: ("!int", True), dt.Int64(nullable=False)),  # "!" overrides `nullable`
     ],
 )
 def test_nullable_dtype(args, expected):
-    assert dt.dtype(*args) == expected
+    assert dt.dtype(*args()) == expected
+
+
+def test_bogus_union():
+    with pytest.raises(TypeError):
+        dt.dtype(Union[int, str])
 
 
 @pytest.mark.parametrize(
diff --git a/ibis/tests/expr/test_udf.py b/ibis/tests/expr/test_udf.py
index e6ed02f511d6..976b734cb038 100644
--- a/ibis/tests/expr/test_udf.py
+++ b/ibis/tests/expr/test_udf.py
@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from typing import Optional
+
 import pytest
 
 import ibis
@@ -152,3 +154,22 @@ def version() -> str: ...
     expr = version()
     assert expr.type().is_string()
     assert type(expr.op().shape) is ibis.expr.datashape.Scalar
+
+
+def test_None_type_hint(table):
+    @ibis.udf.scalar.python
+    def claims_to_return_nullable(x: Optional[int]) -> Optional[int]: ...
+
+    expr = claims_to_return_nullable(table.a)
+    assert isinstance(expr, ir.IntegerColumn)
+    assert expr.type().nullable
+
+    # Even if the signature is annotated with a non-nullable type,
+    # we still interpret the output as nullable, consistent with
+    # how `ibis.dtype(int)` results in a nullable type.
+    @ibis.udf.scalar.python
+    def claims_to_return_nonnullable(x: int) -> int: ...
+
+    expr = claims_to_return_nonnullable(table.a)
+    assert isinstance(expr, ir.IntegerColumn)
+    assert expr.type().nullable