User-Defined Functions#

def __call__(self, *args: Any, **kwargs: Any) -> Expression:
    self._validate_init_args()

    check_serializable(
        self.inner,
        "`@daft.udf` requires that the UDF is serializable. Please double-check that the function does not use any global variables.\n\nIf it does, please use the legacy `@daft.udf` with a class UDF instead and initialize the global in the `__init__` method.",
    )

    bound_args = self._bind_args(*args, **kwargs)
    expressions = list(bound_args.expressions().values())

    ray_options = self.ray_options.copy() if self.ray_options is not None else {}
    if "num_cpus" in ray_options:
        raise ValueError(
            "Cannot set 'num_cpus' in `ray_options`. Please use the 'num_cpus' argument in @udf instead."
        )
    if "num_gpus" in ray_options:
        raise ValueError(
            "Cannot set 'num_gpus' in `ray_options`. Please use the 'num_gpus' argument in @udf instead."
        )
    if "memory" in ray_options:
        raise ValueError(
            "Cannot set 'memory' in `ray_options`. Please use the 'memory_bytes' argument in @udf instead."
        )

    return Expression.udf(
        name=self.name,
        inner=self.wrapped_inner,
        bound_args=bound_args,
        expressions=expressions,
        return_dtype=self.return_dtype,
        init_args=self.init_args,
        resource_request=self.resource_request,
        batch_size=self.batch_size,
        concurrency=self.concurrency,
        use_process=self.use_process,
        ray_options=ray_options,
    )

def override_options(
    self,
    *,
    num_cpus: float | None = _UnsetMarker,
    num_gpus: float | None = _UnsetMarker,
    memory_bytes: int | None = _UnsetMarker,
    ray_options: dict[str, Any] | None = None,
    batch_size: int | None = _UnsetMarker,
) -> UDF:
    """Replace the resource requests for running each instance of your UDF.

    Args:
        num_cpus: Number of CPUs to allocate each running instance of your UDF. Note that this is purely used for placement (e.g. if your
            machine has 8 CPUs and you specify num_cpus=4, then Daft can run at most 2 instances of your UDF at a time).
        num_gpus: Number of GPUs to allocate each running instance of your UDF. This is used for placement and also for allocating
            the appropriate GPU to each UDF using `CUDA_VISIBLE_DEVICES`.
        memory_bytes: Amount of memory to allocate each running instance of your UDF in bytes. If your UDF is experiencing out-of-memory errors,
            this parameter can help hint Daft that each UDF requires a certain amount of heap memory for execution.
        ray_options: Ray options to pass to the UDF. see more  https://docs.ray.io/en/latest/ray-core/api/doc/ray.actor.ActorClass.options.html#ray.actor.ActorClass.options
        batch_size: Enables batching of the input into batches of at most this size. Results between batches are concatenated.

    Examples:
        For instance, if your UDF requires 4 CPUs to run, you can configure it like so:

        >>> import daft
        >>>
        >>> @daft.udf(return_dtype=daft.DataType.string())
        ... def example_udf(inputs):
        ...     # You will have access to 4 CPUs here if you configure your UDF correctly!
        ...     return inputs
        >>>
        >>> # Parametrize the UDF to run with 4 CPUs
        >>> example_udf_4CPU = example_udf.override_options(num_cpus=4)

    """
    new_resource_request = ResourceRequest() if self.resource_request is None else self.resource_request
    if num_cpus is not _UnsetMarker:
        new_resource_request = new_resource_request.with_num_cpus(num_cpus)
    if num_gpus is not _UnsetMarker:
        new_resource_request = new_resource_request.with_num_gpus(num_gpus)
    if memory_bytes is not _UnsetMarker:
        new_resource_request = new_resource_request.with_memory_bytes(memory_bytes)

    new_ray_options = ray_options if ray_options is not None else self.ray_options
    new_batch_size = self.batch_size if batch_size is _UnsetMarker else batch_size

    return dataclasses.replace(
        self, resource_request=new_resource_request, batch_size=new_batch_size, ray_options=new_ray_options
    )

def run_on_process(self, use_process: bool) -> UDF:
    """Override whether this UDF should run on a separate process or not.

    Examples:
        >>> import daft
        >>>
        >>> @daft.udf(return_dtype=daft.DataType.string(), num_gpus=1)
        ... class MyGpuUdf:
        ...     def __init__(self, text=" world"):
        ...         self.text = text
        ...
        ...     def __call__(self, data):
        ...         return [x + self.text for x in data]
        >>>
        >>> # New UDF that will run on a separate process
        >>> MyGpuUdf_separate_process = MyGpuUdf.run_on_process(True)
    """
    return dataclasses.replace(self, use_process=use_process)

def with_concurrency(self, concurrency: int) -> UDF:
    """Override the concurrency of this UDF, which tells Daft how many instances of your UDF to run concurrently.

    Examples:
        >>> import daft
        >>>
        >>> @daft.udf(return_dtype=daft.DataType.string(), num_gpus=1)
        ... class MyGpuUdf:
        ...     def __init__(self, text=" world"):
        ...         self.text = text
        ...
        ...     def __call__(self, data):
        ...         return [x + self.text for x in data]
        >>>
        >>> # New UDF that will have 8 concurrent running instances (will require 8 total GPUs)
        >>> MyGpuUdf_8_concurrency = MyGpuUdf.with_concurrency(8)
    """
    return dataclasses.replace(self, concurrency=concurrency)

def with_init_args(self, *args: Any, **kwargs: Any) -> UDF:
    """Replace initialization arguments for a class UDF when calling `__init__` at runtime on each instance of the UDF.

    Examples:
        >>> import daft
        >>>
        >>> @daft.udf(return_dtype=daft.DataType.string())
        ... class MyUdfWithInit:
        ...     def __init__(self, text=" world"):
        ...         self.text = text
        ...
        ...     def __call__(self, data):
        ...         return [x + self.text for x in data]
        >>>
        >>> # Create a customized version of MyUdfWithInit by overriding the init args
        >>> MyUdfWithInit_CustomInitArgs = MyUdfWithInit.with_init_args(text=" my old friend")
        >>>
        >>> df = daft.from_pydict({"foo": ["hello", "hello", "hello"]})
        >>> df = df.with_column("bar_world", MyUdfWithInit(df["foo"]))
        >>> df = df.with_column("bar_custom", MyUdfWithInit_CustomInitArgs(df["foo"]))
        >>> df.show()
        ╭────────┬─────────────┬─────────────────────╮
        │ foo    ┆ bar_world   ┆ bar_custom          │
        │ ---    ┆ ---         ┆ ---                 │
        │ String ┆ String      ┆ String              │
        ╞════════╪═════════════╪═════════════════════╡
        │ hello  ┆ hello world ┆ hello my old friend │
        ├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
        │ hello  ┆ hello world ┆ hello my old friend │
        ├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
        │ hello  ┆ hello world ┆ hello my old friend │
        ╰────────┴─────────────┴─────────────────────╯
        <BLANKLINE>
        (Showing first 3 of 3 rows)
    """
    if not isinstance(self.inner, type):
        raise ValueError("Function UDFs cannot have init args.")

    init_sig = inspect.signature(self.inner.__init__)  # type: ignore
    init_sig.bind(
        # Placeholder for `self`
        None,
        *args,
        **kwargs,
    )
    return dataclasses.replace(self, init_args=(args, kwargs))

def __call__(
    self,
    fn: Callable[P, T] | None = None,
    *,
    return_dtype: DataTypeLike | None = None,
    unnest: bool = False,
    cpus: float | None = None,
    gpus: float = 0,
    use_process: bool | None = None,
    max_concurrency: int | None = None,
    max_retries: int | None = None,
    on_error: Literal["raise", "log", "ignore"] | None = None,
    ray_options: dict[str, Any] | None = None,
) -> Callable[[Callable[P, T]], Func[P, T, None]] | Func[P, T, None]:
    """Decorator to convert a Python function into a Daft user-defined function.

    Args:
        return_dtype: The data type that this function should return or yield. If not specified, it is derived from the function's return type hint.
        unnest: Whether to unnest/flatten out return type fields into columns. Return dtype must be `DataType.struct(..)` when this is set to true.
        use_process: Whether to run each instance of the function in a separate process. If unset, Daft will automatically choose based on runtime performance.
        max_concurrency: The maximum number of concurrent coroutines for async functions. Only valid for async functions; raises an error if used with synchronous functions.

    Daft function variants:
    - **Row-wise** (1 row in, 1 row out) - the default variant
    - **Async row-wise** (1 row in, 1 row out) - created by decorating a Python async function
    - **Generator** (1 row in, N rows out) - created by decorating a Python generator function

    Decorated functions accept both their original argument types and Daft Expressions.
    When any arguments are Expressions, they return a Daft Expression that can be used in DataFrame operations.
    When called without Expression arguments, they execute immediately and the behavior is the same as if the function was not decorated.

    Examples:
        Basic Example

        >>> import daft
        >>> @daft.func
        ... def my_sum(a: int, b: int) -> int:
        ...     return a + b
        >>>
        >>> df = daft.from_pydict({"x": [1, 2, 3], "y": [4, 5, 6]})
        >>> df.select(my_sum(df["x"], df["y"])).collect()
        ╭───────╮
        │ x     │
        │ ---   │
        │ Int64 │
        ╞═══════╡
        │ 5     │
        ├╌╌╌╌╌╌╌┤
        │ 7     │
        ├╌╌╌╌╌╌╌┤
        │ 9     │
        ╰───────╯
        <BLANKLINE>
        (Showing first 3 of 3 rows)

        Calling the decorator directly on an existing function

        >>> import daft
        >>> def tokenize(text: str) -> list[int]:
        ...     vocab = {char: i for i, char in enumerate(text)}
        ...     return [vocab[char] for char in text]
        >>>
        >>> daft_tokenize = daft.func(tokenize)  # creates a new function rather than modifying `tokenize`
        >>> df = daft.from_pydict({"text": ["hello", "world", "daft"]})
        >>> df.select(daft_tokenize(df["text"])).collect()
        ╭─────────────────╮
        │ text            │
        │ ---             │
        │ List[Int64]     │
        ╞═════════════════╡
        │ [0, 1, 3, 3, 4] │
        ├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
        │ [0, 1, 2, 3, 4] │
        ├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
        │ [0, 1, 2, 3]    │
        ╰─────────────────╯
        <BLANKLINE>
        (Showing first 3 of 3 rows)

        Manually specifying the return type

        >>> import daft
        >>> @daft.func(return_dtype=daft.DataType.int32())
        ... def my_sum(a: int, b: int):
        ...     return a + b
        >>>
        >>> df = daft.from_pydict({"x": [1, 2, 3], "y": [4, 5, 6]})
        >>> df.select(my_sum(df["x"], df["y"])).collect()
        ╭───────╮
        │ x     │
        │ ---   │
        │ Int32 │
        ╞═══════╡
        │ 5     │
        ├╌╌╌╌╌╌╌┤
        │ 7     │
        ├╌╌╌╌╌╌╌┤
        │ 9     │
        ╰───────╯
        <BLANKLINE>
        (Showing first 3 of 3 rows)

        Decorating an async function

        >>> import daft
        >>> import asyncio
        >>> @daft.func
        ... async def my_sum(a: int, b: int) -> int:
        ...     await asyncio.sleep(1)
        ...     return a + b
        >>>
        >>> df = daft.from_pydict({"x": [1], "y": [2]})
        >>> df.select(my_sum(df["x"], df["y"])).collect()
        ╭───────╮
        │ x     │
        │ ---   │
        │ Int64 │
        ╞═══════╡
        │ 3     │
        ╰───────╯
        <BLANKLINE>
        (Showing first 1 of 1 rows)

        Decorating a generator function

        >>> import daft
        >>> from typing import Iterator
        >>> @daft.func
        ... def my_gen_func(to_repeat: str, n: int) -> Iterator[str]:
        ...     for _ in range(n):
        ...         yield to_repeat
        >>>
        >>> df = daft.from_pydict({"id": [0, 1, 2], "value": ["pip", "install", "daft"], "occurrences": [0, 2, 4]})
        >>> df = df.select("id", my_gen_func(df["value"], df["occurrences"]))
        >>> df.collect()  # other output columns are repeated to match generator output length
        ╭───────┬─────────╮
        │ id    ┆ value   │
        │ ---   ┆ ---     │
        │ Int64 ┆ String  │
        ╞═══════╪═════════╡
        │ 0     ┆ None    │
        ├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
        │ 1     ┆ install │
        ├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
        │ 1     ┆ install │
        ├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
        │ 2     ┆ daft    │
        ├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
        │ 2     ┆ daft    │
        ├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
        │ 2     ┆ daft    │
        ├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
        │ 2     ┆ daft    │
        ╰───────┴─────────╯
        <BLANKLINE>
        (Showing first 7 of 7 rows)

        Unnesting multiple return fields

        >>> import daft
        >>> from daft import DataType
        >>> @daft.func(
        ...     return_dtype=DataType.struct({"int": DataType.int64(), "str": DataType.string()}), unnest=True
        ... )
        ... def my_multi_return(val: int):
        ...     return {"int": val * 2, "str": str(val) * 2}
        >>> df = daft.from_pydict({"x": [1, 2, 3]})
        >>> df.select(my_multi_return(df["x"])).collect()
        ╭───────┬────────╮
        │ int   ┆ str    │
        │ ---   ┆ ---    │
        │ Int64 ┆ String │
        ╞═══════╪════════╡
        │ 2     ┆ 11     │
        ├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
        │ 4     ┆ 22     │
        ├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
        │ 6     ┆ 33     │
        ╰───────┴────────╯
        <BLANKLINE>
        (Showing first 3 of 3 rows)
    """

    def partial_func(fn: Callable[P, T]) -> Func[P, T, None]:
        return Func._from_func(
            fn,
            return_dtype,
            unnest,
            cpus,
            gpus,
            use_process,
            False,
            None,
            max_concurrency=max_concurrency,
            max_retries=max_retries,
            on_error=on_error,
            ray_options=ray_options,
        )

    return partial_func if fn is None else partial_func(fn)

def batch(
    self,
    *,
    return_dtype: DataTypeLike,
    unnest: bool = False,
    cpus: float | None = None,
    gpus: float = 0,
    use_process: bool | None = None,
    max_concurrency: int | None = None,
    batch_size: int | None = None,
    max_retries: int | None = None,
    on_error: Literal["raise", "log", "ignore"] | None = None,
    ray_options: dict[str, Any] | None = None,
) -> Callable[[Callable[P, T]], Func[P, T, None]]:
    """Decorator to convert a Python function into a Daft user-defined batch function.

    Args:
        return_dtype: The data type that this function should return.
        unnest: Whether to unnest/flatten out return type fields into columns. Return dtype must be `DataType.struct(..)` when this is set to true.
        use_process: Whether to run each instance of the function in a separate process. If unset, Daft will automatically choose based on runtime performance.
        max_concurrency: The maximum number of concurrent coroutines for async functions. Only valid for async functions; raises an error if used with synchronous functions.
        batch_size: The max number of rows in each input batch.

    Batch functions receive `daft.Series` arguments, and return a `daft.Series`, `list`, `numpy.ndarray`, or `pyarrow.Array`.
    You can also call them with scalar arguments, which will be passed in without modification.
    When called without Expression arguments, they execute immediately and the behavior is the same as if the function was not decorated.

    Examples:
        Basic Usage

        >>> import daft
        >>> from daft import DataType, Series
        >>>
        >>> @daft.func.batch(return_dtype=DataType.int64())
        ... def my_sum(a: Series, b: Series) -> Series:
        ...     import pyarrow.compute as pc
        ...
        ...     a = a.to_arrow()
        ...     b = b.to_arrow()
        ...     result = pc.add(a, b)
        ...     return result
        >>>
        >>> df = daft.from_pydict({"x": [1, 2, 3], "y": [4, 5, 6]})
        >>> df.select(my_sum(df["x"], df["y"])).collect()
        ╭───────╮
        │ x     │
        │ ---   │
        │ Int64 │
        ╞═══════╡
        │ 5     │
        ├╌╌╌╌╌╌╌┤
        │ 7     │
        ├╌╌╌╌╌╌╌┤
        │ 9     │
        ╰───────╯
        <BLANKLINE>
        (Showing first 3 of 3 rows)

        Mixing Series and Scalar Arguments

        >>> @daft.func.batch(return_dtype=daft.DataType.int64())
        ... def my_sum_with_scalar(a: daft.Series, b: int) -> daft.Series:
        ...     import pyarrow.compute as pc
        ...
        ...     a = a.to_arrow()
        ...     result = pc.add(a, b)
        ...     return result
        >>>
        >>> df = daft.from_pydict({"x": [1, 2, 3]})
        >>> df.select(my_sum_with_scalar(df["x"], 4)).collect()
        ╭───────╮
        │ x     │
        │ ---   │
        │ Int64 │
        ╞═══════╡
        │ 5     │
        ├╌╌╌╌╌╌╌┤
        │ 6     │
        ├╌╌╌╌╌╌╌┤
        │ 7     │
        ╰───────╯
        <BLANKLINE>
        (Showing first 3 of 3 rows)
    """

    def partial_func(fn: Callable[P, T]) -> Func[P, T, None]:
        return Func._from_func(
            fn,
            return_dtype,
            unnest,
            cpus,
            gpus,
            use_process,
            True,
            batch_size,
            max_concurrency=max_concurrency,
            max_retries=max_retries,
            on_error=on_error,
            ray_options=ray_options,
        )

    return partial_func

def __call__(
    self,
    method: Callable[P, T] | None = None,
    *,
    return_dtype: DataTypeLike | None = None,
    unnest: bool = False,
    max_retries: int | None = None,
    on_error: Literal["raise", "log", "ignore"] | None = None,
) -> Callable[P, T] | Callable[[Callable[P, T]], Callable[P, T]]:
    """Decorator to convert a Python method into a Daft user-defined function. This should be used in a class that is decorated with `@daft.cls`.

    Args:
        return_dtype: The data type that this function should return or yield. If not specified, it is derived from the method's return type hint.
        unnest: Whether to unnest/flatten out return type fields into columns. Return dtype must be `DataType.struct(..)` when this is set to true. Defaults to false.

    Similar to `@daft.func`, `@daft.method` supports three variants: row-wise, async row-wise, and generator. See `@daft.func` for more details.

    Decorated methods accept both their original argument types and Daft Expressions.
    When any arguments are Expressions, they return a Daft Expression that can be used in DataFrame operations.
    When called without Expression arguments, methods execute immediately, first initializing a local instance of the class if it does not already exist.

    See `@daft.func` and `@daft.cls` for more details.
    """

    def partial_method(m: Callable[P, T]) -> Callable[P, T]:
        return mark_cls_method(m, return_dtype, unnest, False, None, max_retries, on_error)

    return partial_method if method is None else partial_method(method)

def batch(
    self,
    method: Callable[P, T] | None = None,
    *,
    return_dtype: DataTypeLike | None = None,
    unnest: bool = False,
    batch_size: int | None = None,
    max_retries: int | None = None,
    on_error: Literal["raise", "log", "ignore"] | None = None,
) -> Callable[P, T] | Callable[[Callable[P, T]], Callable[P, T]]:
    """Decorator to convert a Python method into a Daft user-defined batch function. This should be used in a class that is decorated with `@daft.cls`.

    Args:
        return_dtype: The data type that this function should return.
        unnest: Whether to unnest/flatten out return type fields into columns. Return dtype must be `DataType.struct(..)` when this is set to true.
        batch_size: The max number of rows in each input batch.

    Batch methods receive `daft.Series` arguments, and return a `daft.Series`, `list`, `numpy.ndarray`, or `pyarrow.Array`.
    You can also call them with scalar arguments, which will be passed in without modification.
    When called without Expression arguments, they execute immediately, first initializing a local instance of the class if it does not already exist.

    See `@daft.func.batch` and `@daft.cls` for more details.
    """

    def partial_method(m: Callable[P, T]) -> Callable[P, T]:
        return mark_cls_method(m, return_dtype, unnest, True, batch_size, max_retries, on_error)

    return partial_method if method is None else partial_method(method)