#
# 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))
[docs]
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)
