viirya commented on code in PR #58978: URL: https://github.com/apache/spark/pull/58978#discussion_r4291057499
########## sql/core/pom.xml: ########## @@ -35,6 +35,17 @@ </properties> <dependencies> + <!-- + jep (Java Embedded Python): embeds CPython in the JVM via JNI for in-process Python UDFs. + Marked as "provided": users must install jep on each executor separately. + See: https://github.com/ninia/jep + --> + <dependency> + <groupId>black.ninia</groupId> + <artifactId>jep</artifactId> + <version>4.3.2</version> Review Comment: Done in 01b61ff as suggested: the root pom has a `jep.version` property and a `provided` entry in `dependencyManagement`, `sql/core` no longer names a version, and the CI image reads `jep.version` from the pom the same way it reads `arrow.version`. The guide now asks for `jep>=4.3.2,<4.4` and says that Spark is built and tested against 4.3.2 and uses JEP internals that a later minor version may change. ########## python/pyspark/inprocess/udf.py: ########## @@ -0,0 +1,264 @@ +# +# 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() + + +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 isinstance(returnType, (DataType, str)): + raise PySparkTypeError( + errorClass="NOT_EXPECTED_TYPE", + messageParameters={ + "expected_type": "DataType or str", + "arg_name": "returnType", + "arg_type": type(returnType).__name__, + }, + ) + self._return_type = returnType + self._parsed_return_type: Optional[DataType] = None + self.evalType = PythonEvalType.SQL_SCALAR_ARROW_INPROCESS_UDF + self._deterministic: bool = deterministic + self._name: str = getattr(func, "__name__", "inprocess_udf") + + 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 + self._serialized: Optional[bytes] = None + # The JVM-side function, created once so that calls do not resend the command. + self._jfunc: Optional["JavaObject"] = None + update_wrapper(self, func, updated=()) + + @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, and + # the serialized command is labeled with the version of the process that sends it, so + # a copy, possibly in another Python, serializes the function and creates its own. + state = self.__dict__.copy() + state["_jfunc"] = None + state["_serialized"] = None + return state + + def _serialize(self) -> bytes: + if self._serialized is None: + # Validate before caching the command, including driver-only UDT definitions. + self.returnType + self._serialized = _serialize_udf(self._func) + return self._serialized + + 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(returnType: Union[DataType, str], deterministic: bool = True) -> Callable: Review Comment: Done in 01b61ff with the `arrow_udf` shape: `inprocess_udf(f=None, returnType=None, *, deterministic=True)`. A DataType or a string as `f`, or no `f`, returns the decorator, as before, and a function returns the wrapper, so `inprocess_udf(lambda x: pc.multiply(x, 2), LongType())` works. `deterministic` is keyword-only and must be a `bool`, and every argument is validated when `inprocess_udf` is called: a non-callable `f`, a missing or invalid return type, a non-bool `deterministic`, or a return type given both positionally and as `returnType` raises a `PySparkTypeError` at once. The runtime test covers each of these. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala: ########## @@ -1407,9 +1408,13 @@ object ConvertToCatalyst extends Rule[LogicalPlan] { // preserve the UDF batch pipeline (e.g. an outer UDF that could not be // transpiled wrapping one that could). // + // An in-process UDF never chains with worker UDFs, so it has no pipeline to preserve. + // // A conditional branch is not treated like a lambda, on purpose: the column stands, which // makes the argument eager, and the interpreted UDF is eager there too. See `transpile.py`. - expression.mapChildren(applyExpr(_, parentIsUdf = isScalarPythonUDF(expression), + val chains = isScalarPythonUDF(expression) && expression.asInstanceOf[PythonUDF].evalType != Review Comment: Fixed in 01b61ff: `hasOnlyPythonUDFInputs` now counts only arguments that chain in a worker (see the next thread), so `t` in `w(t(inproc(x)))` is transpiled. "apply transpiles a UDF whose only Python argument is an in-process UDF" in `ConvertToCatalystSuite` covers that shape and fails with the previous predicate. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/python/InProcessArrowEvalPythonExec.scala: ########## @@ -0,0 +1,79 @@ +/* + * 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. + */ + +package org.apache.spark.sql.execution.python + +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Attribute, PythonUDF} +import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} +import org.apache.spark.sql.internal.SQLConf + +/** Row-based CDI execution, sharing the standard Python UDF evaluator contracts. */ +case class InProcessArrowEvalPythonExec( + udfs: Seq[PythonUDF], + resultAttrs: Seq[Attribute], + child: SparkPlan) extends EvalPythonExec { + + // The keys of `PythonSQLMetrics`, which the evaluator updates, without `pythonBootTime`: + // there are no workers to start. + private val pythonMetrics: Map[String, SQLMetric] = + InProcessArrowEvalPythonExec.sizeMetricsDesc.map { case (k, v) => + k -> SQLMetrics.createSizeMetric(sparkContext, v) + } ++ InProcessArrowEvalPythonExec.timingMetricsDesc.map { case (k, v) => + k -> SQLMetrics.createTimingMetric(sparkContext, v) + } ++ InProcessArrowEvalPythonExec.otherMetricsDesc.map { case (k, v) => + k -> SQLMetrics.createMetric(sparkContext, v) + } + + override lazy val metrics: Map[String, SQLMetric] = pythonMetrics + + override protected def doExecute(): RDD[InternalRow] = { + InProcessPythonUDFBuilder.checkConfiguration(conf) + super.doExecute() + } + + // The evaluator consumes each input row before pulling the next one, copying it into its + // queue and writing its arguments to Arrow, so the rows need not be copied first. + override protected def evaluatorInput: RDD[InternalRow] = child.execute() + + override protected def evaluatorFactory: EvalPythonEvaluatorFactory = { + new InProcessArrowEvalPythonEvaluatorFactory( + child.output, udfs, output, conf.arrowMaxRecordsPerBatch, conf.arrowMaxBytesPerBatch, + conf.sessionLocalTimeZone, conf.arrowUseLargeVarTypes, conf.pysparkHideTraceback, + conf.pysparkSimplifiedTraceback, conf.pysparkTracebackWithLocals, + conf.getConf(SQLConf.PYTHON_UDF_IN_PROCESS_FULL_VALIDATION), pythonMetrics) + } + + override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan = + copy(child = newChild) +} + +object InProcessArrowEvalPythonExec { + private[python] val sizeMetricsDesc: Map[String, String] = Map( + "pythonDataSent" -> "data exported to in-process Python", + "pythonDataReceived" -> "data returned from in-process Python") + + private[python] val timingMetricsDesc: Map[String, String] = Map( + "pythonInitTime" -> "time to register UDFs in in-process Python", + "pythonTotalTime" -> "time to run in-process Python UDFs", Review Comment: Updated in 01b61ff with your wording: "total time of in-process UDF evaluation, including reading its input". ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/PythonUDF.scala: ########## @@ -51,7 +51,8 @@ object PythonUDF { PythonEvalType.SQL_SCALAR_PANDAS_UDF, PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF, PythonEvalType.SQL_SCALAR_ARROW_UDF, - PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF + PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF, + PythonEvalType.SQL_SCALAR_ARROW_INPROCESS_UDF Review Comment: Done in 01b61ff: `PythonUDF.chainsInWorker(e)` is a scalar Python UDF that is not in-process, and `ExtractPythonUDFs`, `CollapseProject`, `ConvertToCatalyst` and `hasOnlyPythonUDFInputs` use it instead of their own exemptions. Its Scaladoc asks rules that keep or build worker chains to use it rather than `isScalarPythonUDF`. Making it equal to `isScalarPythonUDF` fails both transpiler tests and "nested calls and collapsed projects produce separate evaluation nodes". ########## sql/core/src/main/scala/org/apache/spark/sql/execution/python/InProcessArrowEvalPythonExec.scala: ########## @@ -0,0 +1,79 @@ +/* + * 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. + */ + +package org.apache.spark.sql.execution.python + +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Attribute, PythonUDF} +import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} +import org.apache.spark.sql.internal.SQLConf + +/** Row-based CDI execution, sharing the standard Python UDF evaluator contracts. */ +case class InProcessArrowEvalPythonExec( + udfs: Seq[PythonUDF], + resultAttrs: Seq[Attribute], + child: SparkPlan) extends EvalPythonExec { + + // The keys of `PythonSQLMetrics`, which the evaluator updates, without `pythonBootTime`: + // there are no workers to start. + private val pythonMetrics: Map[String, SQLMetric] = Review Comment: Done in 01b61ff as suggested: `PythonSQLMetrics` builds its metrics from `protected def sizeMetricsDesc`, `timingMetricsDesc` and `otherMetricsDesc`, which default to the shared maps, and `InProcessArrowEvalPythonExec` mixes it in again and overrides only those three with constants from its companion object, since the trait reads them during construction. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
