python/pyspark/inprocess/udf.py

#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

"""
Python API for in-process UDF registration.

Usage::

    import pyarrow.compute as pc
    from pyspark.inprocess import inprocess_udf
    from pyspark.sql.types import LongType

    @inprocess_udf(returnType=LongType())
    def double(x):
        # x is a pa.Array; return a pa.Array
        return pc.multiply(x, 2)

    df.select(double(df.value)).show()
"""

import io
import sys
from functools import update_wrapper
from inspect import signature
from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Union

from pyspark import Accumulator, Broadcast, cloudpickle
from pyspark.errors import PySparkNotImplementedError, PySparkTypeError, PySparkValueError
from pyspark.sql.column import Column
from pyspark.sql.types import DataType, _parse_datatype_string
from pyspark.util import PythonEvalType

if TYPE_CHECKING:
    from py4j.java_gateway import JavaObject


class _InProcessPickler(cloudpickle.CloudPickler):
    def reducer_override(self, obj: Any) -> Any:
        if isinstance(obj, (Broadcast, Accumulator)):
            raise PySparkNotImplementedError(
                errorClass="NOT_IMPLEMENTED",
                messageParameters={
                    "feature": "Spark broadcasts or accumulators in in-process UDFs"
                },
            )
        return super().reducer_override(obj)


def _serialize_udf(func: Callable) -> bytes:
    buffer = io.BytesIO()
    _InProcessPickler(buffer).dump(func)
    return buffer.getvalue()


def _check_return_type_argument(returnType: Any) -> Union[DataType, str]:
    if not isinstance(returnType, (DataType, str)):
        raise PySparkTypeError(
            errorClass="NOT_EXPECTED_TYPE",
            messageParameters={
                "expected_type": "DataType or str",
                "arg_name": "returnType",
                "arg_type": type(returnType).__name__,
            },
        )
    return returnType


class InProcessUDFWrapper:
    """
    Wraps a Python function as an in-process UDF.

    Returned by ``@inprocess_udf``. Calling an instance with Spark ``Column``
    arguments creates a ``Column`` expression backed by ``PythonUDF``
    on the JVM side.
    """

    def __init__(
        self, func: Callable, returnType: Union[DataType, str], deterministic: bool = True
    ) -> None:
        if not callable(func):
            raise PySparkTypeError(
                errorClass="NOT_EXPECTED_TYPE",
                messageParameters={
                    "expected_type": "function",
                    "arg_name": "f",
                    "arg_type": type(func).__name__,
                },
            )
        _check_return_type_argument(returnType)
        self._return_type = returnType
        self._parsed_return_type: Optional[DataType] = None
        self.evalType = PythonEvalType.SQL_SCALAR_ARROW_INPROCESS_UDF
        self._deterministic: bool = deterministic
        # As for other Python UDFs, e.g. a functools.partial or a callable object.
        self._name: str = getattr(func, "__name__", type(func).__name__)

        from pyspark.sql.pandas.utils import require_minimum_pyarrow_version

        require_minimum_pyarrow_version()
        if not signature(func).parameters:
            raise PySparkValueError(
                errorClass="INVALID_PANDAS_UDF",
                messageParameters={"detail": "0-arg inprocess_udfs are not supported."},
            )
        self._func = func
        # The JVM-side function, created once so that calls do not resend the command.
        self._jfunc: Optional["JavaObject"] = None
        update_wrapper(self, func, updated=())
        self.__name__ = self._name

    @property
    def func(self) -> Callable:
        return self._func

    @property
    def returnType(self) -> DataType:
        if self._parsed_return_type is None:
            parsed = (
                _parse_datatype_string(self._return_type)
                if isinstance(self._return_type, str)
                else self._return_type
            )
            from pyspark.sql.udf import UserDefinedFunction

            UserDefinedFunction._check_return_type(parsed, PythonEvalType.SQL_SCALAR_ARROW_UDF)
            from pyspark.sql.pandas.types import to_arrow_type

            to_arrow_type(parsed, timezone="UTC", error_on_duplicated_field_names_in_struct=True)
            self._parsed_return_type = parsed
        return self._parsed_return_type

    @property
    def deterministic(self) -> bool:
        return self._deterministic

    def asNondeterministic(self) -> "InProcessUDFWrapper":
        self._deterministic = False
        return self

    def __getstate__(self) -> Dict[str, Any]:
        # The JVM-side function holds the py4j gateway, which cannot be pickled or copied. A
        # copy, possibly in another Python, serializes the function and creates its own.
        state = self.__dict__.copy()
        state["_jfunc"] = None
        return state

    def _serialize(self) -> bytes:
        # Validate first, including driver-only UDT definitions. The command is not kept: the
        # JVM-side function, created once, holds it.
        self.returnType
        return _serialize_udf(self._func)

    def __call__(self, *cols: Union[Column, str], **kwargs: Union[Column, str]) -> Column:
        """
        Create a ``Column`` expression invoking this UDF with the given columns.

        Args:
            *cols: Spark ``Column`` objects (e.g. ``df.value``, ``col("x")``)

        Returns:
            pyspark.sql.Column
        """
        from pyspark.sql.classic.column import _to_java_column
        from pyspark.sql.utils import get_active_spark_context, is_remote

        if is_remote():
            raise PySparkNotImplementedError(
                errorClass="NOT_IMPLEMENTED",
                messageParameters={"feature": "In-process Python UDFs in Spark Connect"},
            )
        sc = get_active_spark_context()

        jvm = sc._jvm
        assert jvm is not None

        # Convert Python Column objects to JVM Column objects
        if not cols and not kwargs:
            raise PySparkValueError(
                errorClass="INVALID_PANDAS_UDF",
                messageParameters={"detail": "An inprocess_udf requires at least one argument."},
            )
        jcols = [_to_java_column(c) for c in cols]
        jcols.extend(
            jvm.PythonSQLUtils.namedArgumentExpression(name, _to_java_column(value))
            for name, value in kwargs.items()
        )

        # Build a Java ArrayList (py4j vararg spread doesn't work with Arrays.asList)
        jlist = jvm.java.util.ArrayList()
        for jcol in jcols:
            jlist.add(jcol)

        if self._jfunc is None:
            # Use the existing PythonUDF planning contracts with an in-process eval type.
            builder = jvm.org.apache.spark.sql.execution.python.InProcessPythonUDFBuilder
            self._jfunc = builder.create(
                self._name,
                self._serialize(),
                self.returnType.json(),
                "%d.%d" % sys.version_info[:2],
            )
        return Column(self._jfunc.apply(jlist, self._deterministic))


def inprocess_udf(
    f: Optional[Union[Callable, DataType, str]] = None,
    returnType: Optional[Union[DataType, str]] = None,
    *,
    deterministic: bool = True,
) -> Any:
    """
    Creates an in-process UDF, as a decorator or from a function, like ``arrow_udf``.

    .. versionadded:: 4.4.0

    The function receives one ``pa.Array`` per input column and must
    return a single ``pa.Array`` of the declared ``returnType``.

    The result must have the same length as the input batch and its Arrow type
    must match the declared Spark type, including nested fields. Timezone-aware
    timestamps are relabeled to the session timezone without changing their instants;
    string/binary offset widths are converted to match ``useLargeVarTypes``, and
    ``large_list``, ``fixed_size_list``, ``string_view``, ``binary_view``,
    ``fixed_size_binary`` and dictionary representations are cast to the declared type.
    Other value types must match exactly. Nested nullability may be widened, but actual
    nulls cannot be returned in non-nullable fields. Sliced results are copied when
    required by Arrow Java.

    Spark broadcasts, accumulators, ``SparkFiles``, ``--py-files``,
    ``spark.submit.pyFiles``, ``SparkContext.addPyFile`` and
    ``SparkSession.addArtifacts(..., pyfile=True)`` are unsupported.
    Install dependencies on executors before starting Spark. The driver's Python
    major.minor version must match the embedded interpreter.

    Parameters
    ----------
    f : function, :class:`pyspark.sql.types.DataType` or str, optional
        The function to wrap. When used as a decorator, the return type may be passed here
        instead, as in ``@inprocess_udf(LongType())``.
    returnType : :class:`pyspark.sql.types.DataType` or str
        The return type of the UDF, as a DataType or a DDL-formatted type string.
    deterministic : bool, optional
        Keyword-only. Whether this UDF produces the same output for the same input. Set to
        ``False`` for UDFs that use randomness or external state, so the optimizer does not
        deduplicate or reorder calls to this UDF. Default: ``True``.

    Returns
    -------
    :class:`InProcessUDFWrapper` or function
        The wrapper if ``f`` is a function, otherwise a decorator that wraps a function as
        one.

    Examples
    --------
    .. code-block:: python

        import pyarrow.compute as pc

        double = inprocess_udf(lambda x: pc.multiply(x, 2), LongType())

        @inprocess_udf(returnType=LongType())
        def double(x):
            import pyarrow.compute as pc
            return pc.multiply(x, 2)

        @inprocess_udf(returnType=LongType(), deterministic=False)
        def random_noise(x):
            import pyarrow as pa, numpy as np
            return pa.array(np.random.randint(0, 100, len(x)), type=pa.int64())
    """

    if not isinstance(deterministic, bool):
        raise PySparkTypeError(
            errorClass="NOT_EXPECTED_TYPE",
            messageParameters={
                "expected_type": "bool",
                "arg_name": "deterministic",
                "arg_type": type(deterministic).__name__,
            },
        )
    if f is None or isinstance(f, (DataType, str)):
        # @inprocess_udf(returnType) or @inprocess_udf(returnType=returnType)
        if f is not None and returnType is not None:
            raise PySparkTypeError(
                errorClass="CANNOT_SET_TOGETHER",
                messageParameters={"arg_list": "a return type in f and returnType"},
            )
        # Fails at once rather than when the decorator is applied.
        return_type = _check_return_type_argument(f if f is not None else returnType)

        def decorator(func: Callable) -> InProcessUDFWrapper:
            return InProcessUDFWrapper(func, return_type, deterministic)

        return decorator
    if not callable(f):
        raise PySparkTypeError(
            errorClass="NOT_EXPECTED_TYPE",
            messageParameters={
                "expected_type": "function, DataType or str",
                "arg_name": "f",
                "arg_type": type(f).__name__,
            },
        )
    return InProcessUDFWrapper(f, _check_return_type_argument(returnType), deterministic)