in python/datafusion/udf.py [0:0]
def udf(*args: Any, **kwargs: Any): # noqa: D417
"""Create a new User-Defined Function (UDF).
This class can be used both as a **function** and as a **decorator**.
Usage:
- **As a function**: Call `udf(func, input_types, return_type, volatility,
name)`.
- **As a decorator**: Use `@udf(input_types, return_type, volatility,
name)`. In this case, do **not** pass `func` explicitly.
Args:
func (Callable, optional): **Only needed when calling as a function.**
Skip this argument when using `udf` as a decorator.
input_types (list[pa.DataType]): The data types of the arguments
to `func`. This list must be of the same length as the number of
arguments.
return_type (_R): The data type of the return value from the function.
volatility (Volatility | str): See `Volatility` for allowed values.
name (Optional[str]): A descriptive name for the function.
Returns:
A user-defined function that can be used in SQL expressions,
data aggregation, or window function calls.
Example:
**Using `udf` as a function:**
```
def double_func(x):
return x * 2
double_udf = udf(double_func, [pa.int32()], pa.int32(),
"volatile", "double_it")
```
**Using `udf` as a decorator:**
```
@udf([pa.int32()], pa.int32(), "volatile", "double_it")
def double_udf(x):
return x * 2
```
"""
def _function(
func: Callable[..., _R],
input_types: list[pa.DataType],
return_type: _R,
volatility: Volatility | str,
name: Optional[str] = None,
) -> ScalarUDF:
if not callable(func):
msg = "`func` argument must be callable"
raise TypeError(msg)
if name is None:
if hasattr(func, "__qualname__"):
name = func.__qualname__.lower()
else:
name = func.__class__.__name__.lower()
return ScalarUDF(
name=name,
func=func,
input_types=input_types,
return_type=return_type,
volatility=volatility,
)
def _decorator(
input_types: list[pa.DataType],
return_type: _R,
volatility: Volatility | str,
name: Optional[str] = None,
) -> Callable:
def decorator(func: Callable):
udf_caller = ScalarUDF.udf(
func, input_types, return_type, volatility, name
)
@functools.wraps(func)
def wrapper(*args: Any, **kwargs: Any):
return udf_caller(*args, **kwargs)
return wrapper
return decorator
if args and callable(args[0]):
# Case 1: Used as a function, require the first parameter to be callable
return _function(*args, **kwargs)
# Case 2: Used as a decorator with parameters
return _decorator(*args, **kwargs)