Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 5 additions & 4 deletions python/datafusion/user_defined.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@
if TYPE_CHECKING:
from _typeshed import CapsuleType as _PyCapsule

_R = TypeVar("_R", bound=pa.DataType)
_R = TypeVar("_R", bound=pa.Array)
from collections.abc import Callable, Sequence


Expand Down Expand Up @@ -125,7 +125,7 @@ def __init__(
name: str,
func: Callable[..., _R],
input_fields: list[pa.Field],
return_field: _R,
return_field: pa.Field,
volatility: Volatility | str,
) -> None:
"""Instantiate a scalar user-defined function (UDF).
Expand Down Expand Up @@ -202,7 +202,8 @@ def udf(*args: Any, **kwargs: Any): # noqa: D417
input_fields (list[pa.Field | pa.DataType]): The data types or Fields
of the arguments to ``func``. This list must be of the same length
as the number of arguments.
return_field (_R): The field of the return value from the function.
return_field (pa.DataType | pa.Field): The field of the return value
from the function.
volatility (Volatility | str): See `Volatility` for allowed values.
name (Optional[str]): A descriptive name for the function.

Expand Down Expand Up @@ -264,7 +265,7 @@ def _function(

def _decorator(
input_fields: Sequence[pa.DataType | pa.Field] | pa.DataType | pa.Field,
return_field: _R,
return_field: pa.DataType | pa.Field,
volatility: Volatility | str,
name: str | None = None,
) -> Callable:
Expand Down
Loading