dongjoon-hyun commented on code in PR #58978:
URL: https://github.com/apache/spark/pull/58978#discussion_r4256066923
##########
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:
**[Low, performance] This covers an in-process parent, but not an in-process
argument.**
Thanks for the fix. The mirror image of the case in the new test is still
kept in Python: `hasOnlyPythonUDFInputs` (`PythonUDF.scala` L438) is
`arguments.forall(_.isInstanceOf[PythonUDF])`, so it counts an in-process
argument as a Python UDF input, and L1349 keeps a transpilable UDF in Python
when its parent is a worker UDF.
For example, with transpilation on, `w(t(inproc(x)))`, where `w` and `t` are
worker UDFs and `t` is transpilable: `w` passes `parentIsUdf = true`,
`t.hasOnlyPythonUDFInputs` is true, so `t` takes `keepPython`.
`ExtractPythonUDFs` then extracts `inproc` alone (L203, L208), and `t` runs in
Python for every row inside the worker node of `w`, although nothing chains
with its input. Transpiled, the same node would compute `t` in the JVM.
Could in-process arguments be excluded there as well, e.g. in
`hasOnlyPythonUDFInputs` or at L1349, with a `w(t(inproc(x)))` case in
`ConvertToCatalystSuite`?
##########
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:
**[Low] The JEP version is pinned in three different ways.**
- This module pom hard-codes `4.3.2`. JEP is the only third-party dependency
in `sql/core/pom.xml` with an explicit version, while `arrow-c-data` went into
the root `dependencyManagement`.
- The CI image pins `jep==4.3.2` on its own
(`dev/spark-test-image/python-312/Dockerfile` L74), although it reads
`arrow.version` from the pom "just as the Spark build does" (L78).
- The guide asks for `jep>=4.3.2` (`sql-pyspark-inprocess-udf.md` L337,
L368).
The last one matters at runtime: the guide puts `$JEP_DIR/*` from the pip
package on the system classpath (L380), so the JAR that runs is whatever `pip
install "jep>=4.3.2"` resolves, while Spark is compiled and tested against
4.3.2. `ManagedSharedInterpreter` overrides JEP's protected
`configureInterpreter` hook, and the runtime calls
`MainInterpreter.setInitParams(PyConfig.isolated()...)`, so a later JEP release
that changes these internals would fail at plugin init with a linkage error
that CI never exercised.
Could the version be a `jep.version` property in the root
`dependencyManagement`, read by the Dockerfile like `arrow.version`, and could
the guide name a bounded range, e.g. 4.3.x?
##########
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:
**[Low, cleanup] Each pipeline heuristic now needs its own exemption for
258.**
Adding 258 here and then exempting it rule by rule spreads `evalType !=
SQL_SCALAR_ARROW_INPROCESS_UDF` across `ExtractPythonUDFs` (L203),
`CollapseProject` (`Optimizer.scala` L1688) and `ConvertToCatalyst` (L1415).
`hasOnlyPythonUDFInputs` still lacks it (see my comment at `Optimizer.scala`
L1415), which shows the cost: every rule that reasons about worker chaining,
including future ones, has to remember the exemption, or it silently treats
in-process UDFs as chainable.
Could a shared predicate, e.g. `PythonUDF.chainsInWorker(e)` (a scalar
Python UDF that is not in-process), be used by these heuristics instead, so
that the exemption lives in one place?
##########
docs/configuration.md:
##########
@@ -287,6 +287,21 @@ of the most common options to set are:
</td>
<td>0.7.0</td>
</tr>
+<tr>
+ <td><code>spark.python.inProcess.sitePackages</code></td>
Review Comment:
**[Low, docs] This opt-in setting is listed among the most common options.**
This row is in the "Application Properties" table, which the page introduces
as "some of the most common options to set", right before
`spark.executor.pyspark.memory`. The other `spark.python.*` settings, e.g.
`spark.python.profile`, `spark.python.worker.memory` and
`spark.python.worker.reuse`, are documented under "Runtime Environment"
(L934-L970).
An opt-in setting that needs an executor plugin and native dependencies fits
better next to them. Could this row move there, so that the common table stays
focused?
##########
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:
**[Low, cleanup] This copies the metric construction of `PythonSQLMetrics`.**
This block repeats `PythonSQLMetrics.pythonMetrics` (size, timing and other
description maps, each mapped to `SQLMetrics.create*Metric`), and only the
descriptions differ. Any later change to how Python operator metrics are built,
e.g. a new metric kind or the `additionalMetrics` hook, then has to be made in
both places.
Could the trait expose the description maps as overridable members, e.g.
`protected def timingMetricsDesc: Map[String, String] =
PythonSQLMetrics.pythonTimingMetricsDesc`, so that this node mixes in
`PythonSQLMetrics` again and overrides only its descriptions?
##########
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:
**[Low, API] `inprocess_udf(f, LongType())` silently returns the decorator.**
`inprocess_udf` only works as a decorator factory. The function form that
`udf`, `pandas_udf` and `arrow_udf` accept (`arrow_udf(f=None, returnType=None,
functionType=None)`) binds `returnType=f` and `deterministic=LongType()` here,
and returns the inner `decorator` without an error.
For example, a user who follows the migration section, which says that
`inprocess_udf` and `pandas_udf` "have nearly identical call-site syntax"
(guide L635), may rewrite `double = pandas_udf(lambda x: x * 2, LongType())` as
`double = inprocess_udf(lambda x: pc.multiply(x, 2), LongType())`. Then
`df.select(double("id"))` calls `decorator("id")`, i.e.
`InProcessUDFWrapper("id", <lambda>, LongType())`, which fails with
`NOT_EXPECTED_TYPE` saying that `returnType` is a function. Similarly,
`deterministic` is never type-checked, so `inprocess_udf("long",
deterministic="no")` reaches Py4J as a `String` and fails because
`apply(ArrayList, String)` does not exist.
Since this is a new 4.4.0 API, the signature is hard to change after the
release. Could it either accept `f` first like `arrow_udf`, or make
`deterministic` keyword-only and validate both arguments in `inprocess_udf`
itself, so that a misuse fails at once?
##########
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:
**[Low, metrics] This description suggests Python time, but
`pythonTotalTime` is the wall time of the whole consumer.**
`pythonInitTime` and `pythonProcessingTime` are timed on the interpreter
thread (`timedOnInterpreterThread`), but `pythonTotalTime` is the time from the
consumer's first `hasNext` (`resources.start()`, factory L226) until `close()`
at the end of input or at task completion (L203). The in-process consumer pulls
its input synchronously (`rows.hasNext` / `rows.next()` in `pullRow`), so this
includes the child operators, e.g. a scan, a join or a nested in-process node,
and the downstream work between `next()` calls.
For example, a 10-minute scan that feeds a UDF whose Python runs for 1
second shows about 10 minutes for "time to run in-process Python UDFs", next to
1 second for "time to execute Python code", which points users at Python. Could
the description say what is measured, e.g. "total time of in-process UDF
evaluation, including reading its input"?
--
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]
