spark-test commented on code in PR #59778:
URL: https://github.com/apache/spark/pull/59778#discussion_r4429903621
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala:
##########
@@ -2346,12 +2346,68 @@ class SparkConnectPlanner(
name = fun.getFunctionName,
func = function,
dataType = transformDataType(udf.getOutputType),
- pythonEvalType = udf.getEvalType,
+ pythonEvalType = correctPythonEvalType(udf.getEvalType, function),
udfDeterministic = fun.getDeterministic,
// Set only for incremental Python aggregators (see PythonAggregate).
bufferType = if (udf.hasBufferType) transformDataType(udf.getBufferType)
else null)
}
+ /**
+ * Falls back to a regular Python UDF when the Python that runs an
Arrow-optimized one cannot
+ * run it, as a Python client does before it creates the UDF, see
`_create_py_udf` in
+ * `pyspark.sql.udf`. A Spark Connect client checks its own Python rather
than the one that runs
+ * the UDF, so the server checks again. It checks once per session for each
Python executable,
+ * environment and pandas requirement, and keeps the UDF Arrow-optimized if
the check cannot
+ * run.
+ */
+ private def correctPythonEvalType(evalType: Int, function:
SimplePythonFunction): Int = {
+ if (evalType == PythonEvalType.SQL_ARROW_BATCHED_UDF) {
Review Comment:
**Non-blocking (P2):** Only `SQL_ARROW_BATCHED_UDF` is corrected here, but
Python UDTFs reach the server the same way.
`spark.sql.execution.pythonUDTF.arrow.enabled` defaults to true, both
`_create_py_udtf` implementations choose `SQL_ARROW_TABLE_UDF` after checking
only the client's own pandas and PyArrow, and
`createPythonUserDefinedTableFunction` passes `udtf.getEvalType` through for
inline and registered UDTFs. On the PR's own reproduction server
(`apache/spark:4.2.0-scala2.13-java21-python3-ubuntu`), a default `@udtf`
therefore still fails with `[PACKAGE_NOT_INSTALLED] PyArrow`, while a default
`@udf` on that server now falls back.
The same fallback applies to UDTFs: `SQL_ARROW_TABLE_UDF` to
`SQL_TABLE_UDF`, requiring pandas only under
`spark.sql.legacy.execution.pythonUDTF.pandas.conversion.enabled`, decided at
the same point as the UDF fallback. The description defers UDTFs to a
follow-up; unless that exemption is accepted, the two Arrow-by-default function
kinds behave differently on the same server.
**Recommended change:** Apply the execution-time fallback to Arrow-optimized
Python UDTFs in Spark Connect sessions. When the capability check on the
executor reports that its worker Python cannot import PyArrow (or pandas, with
the UDTF legacy pandas conversion on), evaluate the partition with the
row-based UDTF evaluation and log a warning; otherwise keep the Arrow path. Add
Spark Connect UDTF tests with PyArrow and pandas hidden from the workers'
Python, and mention UDTFs, including that their return values then convert as a
regular Python UDTF's, in the migration guide note.
**Why this works:** UDTFs then fall back under the same rule as UDFs,
decided by the Python that runs them, so a default @udtf on a server without
PyArrow returns rows instead of failing, and servers whose workers have PyArrow
keep the Arrow path.
**Scope:** Extends the execution-time Arrow fallback from Python UDFs to
Python UDTFs in Spark Connect sessions.
**Risks:** Partitions of one UDTF can use different conversions when
executors' Python installations differ. Explain output keeps showing the Arrow
UDTF operator when partitions fall back. A UDTF partition that falls back
converts return values as regular Python UDTFs do, which differs from the Arrow
path for some values, as with the client-side fallback.
**Constraints:** The row-based fallback must produce the same output
attributes as the Arrow UDTF evaluation. Pandas is required only with
spark.sql.legacy.execution.pythonUDTF.pandas.conversion.enabled.
**Exclusions:**
- UDTF analyze calls, which run on the driver and do not use the Arrow path.
**Success:** On a Spark Connect server whose workers' Python lacks PyArrow,
a default @udtf returns the rows the same UDTF returns as a regular Python
UDTF, through the row-based path, and a warning is logged. On a server whose
workers' Python has PyArrow, Arrow-optimized UDTFs run as on the base branch.
**Verification:**
- **Behavior:** With PyArrow hidden from the workers' Python, a default
@udtf in a Spark Connect session, both inline and registered, returns the rows
the same UDTF returns as a regular Python UDTF. With pandas hidden and the UDTF
legacy pandas conversion on, it falls back as well, and with that conversion
off it stays on the Arrow path.
- **Compatibility:** With the internal fallback conf off, as for classic
sessions, Arrow-optimized UDTF evaluation is unchanged.
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala:
##########
@@ -2346,12 +2346,68 @@ class SparkConnectPlanner(
name = fun.getFunctionName,
func = function,
dataType = transformDataType(udf.getOutputType),
- pythonEvalType = udf.getEvalType,
+ pythonEvalType = correctPythonEvalType(udf.getEvalType, function),
udfDeterministic = fun.getDeterministic,
// Set only for incremental Python aggregators (see PythonAggregate).
bufferType = if (udf.hasBufferType) transformDataType(udf.getBufferType)
else null)
}
+ /**
+ * Falls back to a regular Python UDF when the Python that runs an
Arrow-optimized one cannot
+ * run it, as a Python client does before it creates the UDF, see
`_create_py_udf` in
+ * `pyspark.sql.udf`. A Spark Connect client checks its own Python rather
than the one that runs
+ * the UDF, so the server checks again. It checks once per session for each
Python executable,
+ * environment and pandas requirement, and keeps the UDF Arrow-optimized if
the check cannot
+ * run.
+ */
+ private def correctPythonEvalType(evalType: Int, function:
SimplePythonFunction): Int = {
+ if (evalType == PythonEvalType.SQL_ARROW_BATCHED_UDF) {
+ val conf = session.sessionState.conf
+ // The executable and environment that the Python UDF runners launch
workers with.
+ val pythonExec =
conf.pysparkWorkerPythonExecutable.getOrElse(function.pythonExec)
+ // Arrow-optimized Python UDFs need pandas only with the legacy pandas
conversion.
+ val requirePandas = conf.legacyPandasConversionUDF
+ // An invalid environment fails the worker launch with the reason, so it
is left to that.
+ val canRunArrow =
Try(PythonWorkerEnvironment.mergeValidated(function.envVars, conf))
+ .map(_.asScala.toMap)
+ .toOption
+ .forall { envVars =>
+ sessionHolder.arrowOptimizationChecks.computeIfAbsent(
Review Comment:
**Nit (P3):** `computeIfAbsent` holds the bin's lock while this mapping
function runs, and the function starts the worker Python and waits for it: up
to `PROCESS_WAIT_TIMEOUT_MS` (10 s) in `waitFor`, then up to another 10 s in
`reader.join`. `ConcurrentHashMap` documents that other updates may block while
the mapping function runs and asks for a short computation. Operations of one
Connect session can plan concurrently, on ExecutePlan runner threads and on
AnalyzePlan calls such as `df.schema`. A second operation that plans a UDF with
the same key, or with a key that lands in the same bin, blocks entering the
monitor, and `interruptOperation` or `interruptTag` cannot cancel it, because a
thread waiting for a monitor ignores `Thread.interrupt()`. It waits about 0.5 s
typically, and about 10 s when the import hangs until the check times out.
See **Shared repair plan 1** in the review body.
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala:
##########
@@ -2346,12 +2346,68 @@ class SparkConnectPlanner(
name = fun.getFunctionName,
func = function,
dataType = transformDataType(udf.getOutputType),
- pythonEvalType = udf.getEvalType,
+ pythonEvalType = correctPythonEvalType(udf.getEvalType, function),
Review Comment:
**Non-blocking (P2):** This fixes the eval type into the UDF when Connect
plans or registers it, from a probe that `correctPythonEvalType` runs on the
driver host with the session settings as of planning. The Python UDF runners
instead resolve `spark.sql.execution.pyspark.python`, `spark.pythonWorkerEnv.*`
and the legacy pandas conf on the executors at each action, and they install
the environment at worker launch "so a cached plan cannot pin an older
environment" (`PythonUDFRunner.scala:44`).
Two cases that work today regress:
- Host: when the driver's Python at that executable path lacks PyArrow but
the executors' Python has it, for example a Connect server started in client
mode from an image without PyArrow whose executor image has it, default UDFs
move from `ArrowEvalPython` to `BatchEvalPython`. That is the slower path, and
return values whose conversion differs change silently: `@udf('boolean')`
returning `1` gives `True` with Arrow and `NULL` without, per the return-type
coercion goldens. The only signal is a server log, and the migration note does
not say that results can change. Classic PySpark makes the same driver-side
choice, but these Connect deployments run the UDFs Arrow-optimized today.
- Time: a function registered, or a DataFrame kept in the session plan
cache, while the worker Python lacked PyArrow keeps `SQL_BATCHED_UDF` after the
session points `spark.pythonWorkerEnv.PYTHONPATH` or
`spark.sql.execution.pyspark.python` at a Python that has it. In the other
direction, one planned while PyArrow was present still fails after the session
switches to a Python without it.
Local mode, where the driver runs the workers' Python, is not affected by
the host case.
See **Shared repair plan 1** in the review body.
##########
core/src/main/scala/org/apache/spark/api/python/PythonUtils.scala:
##########
@@ -196,4 +197,70 @@ private[spark] object PythonUtils extends Logging {
broadcastVars = List.empty.asJava,
accumulator = null)
}
+
+ /**
+ * Checks whether `pythonExec` can run Arrow-optimized Python UDFs, that is,
whether it has
+ * PyArrow at the version PySpark requires, and pandas too if
`requirePandas`, which the legacy
+ * pandas conversion needs. A Python client runs a similar check before it
creates such a UDF,
+ * see `_create_py_udf` in `pyspark.sql.udf`, where it always requires
pandas.
+ *
+ * It runs `pythonExec` with the environment, Python path and working
directory that
+ * `PythonWorkerFactory` gives a Python worker launched with `envVars`.
+ *
+ * @return `None` if it can, and why it cannot otherwise.
+ * @throws SparkException if the check fails otherwise, or does not finish
within `timeoutMs`.
+ */
+ private[spark] def checkArrowOptimization(
+ pythonExec: String,
+ envVars: Map[String, String],
+ requirePandas: Boolean = false,
+ timeoutMs: Long = PythonWorkerFactory.PROCESS_WAIT_TIMEOUT_MS):
Option[String] = {
+ // Exits with 3 only when PyArrow, or pandas when required, is missing or
too old, so that no
+ // other failure, such as PySpark not being importable, is taken for that.
+ val checks = (if (requirePandas) Seq("require_minimum_pandas_version")
else Nil) :+
+ "require_minimum_pyarrow_version"
+ val code =
+ s"""
+ |import sys
+ |from pyspark.sql.pandas.utils import ${checks.mkString(", ")}
+ |try:
+ | ${checks.map(_ + "()").mkString("; ")}
+ |except ImportError as e:
+ | print(e)
+ | sys.exit(3)
+ |""".stripMargin
+ val jobArtifactUUID =
+
JobArtifactSet.getCurrentJobArtifactState.map(_.uuid).getOrElse("default")
+ val workingDir = if (jobArtifactUUID != "default") {
+ val dir = new File(SparkFiles.getRootDirectory(), jobArtifactUUID)
+ dir.mkdir()
Review Comment:
**Nit (P3):** On the driver, `SparkFiles.getRootDirectory()` is
`SparkEnv.driverTmpDir`, which only `SparkEnv.stop` deletes. Closing a Connect
session does not remove `<root>/<session UUID>`: `ArtifactManager` cleanup
removes the file-server entries and the artifact path, and the executor's
session cleanup deletes only under the executor's own root, which is the
driver's root only in local mode. So on a cluster-mode server every session
that plans a default `@udf` now leaves an empty directory behind, where before
only sessions that added files or archives, or ran a Python planning worker
such as a UDTF's `analyze` or a Python data source, did. The count grows with
session churn over the server's lifetime.
See **Shared repair plan 1** in the review body.
##########
core/src/main/scala/org/apache/spark/api/python/PythonUtils.scala:
##########
@@ -196,4 +197,70 @@ private[spark] object PythonUtils extends Logging {
broadcastVars = List.empty.asJava,
accumulator = null)
}
+
+ /**
+ * Checks whether `pythonExec` can run Arrow-optimized Python UDFs, that is,
whether it has
+ * PyArrow at the version PySpark requires, and pandas too if
`requirePandas`, which the legacy
+ * pandas conversion needs. A Python client runs a similar check before it
creates such a UDF,
+ * see `_create_py_udf` in `pyspark.sql.udf`, where it always requires
pandas.
+ *
+ * It runs `pythonExec` with the environment, Python path and working
directory that
+ * `PythonWorkerFactory` gives a Python worker launched with `envVars`.
+ *
+ * @return `None` if it can, and why it cannot otherwise.
+ * @throws SparkException if the check fails otherwise, or does not finish
within `timeoutMs`.
Review Comment:
**Nit (P3):** When `pythonExec` cannot be started (a missing or
non-executable path, or a relative path that does not resolve in the working
directory), `Utils.executeCommand` throws `java.io.IOException` from
`ProcessBuilder.start` before the `try` block, and it propagates unwrapped, as
the new "checkArrowOptimization throws when the check cannot run" test asserts
for `/nonexistent/python3`. This line documents only `SparkException`, so a
`private[spark]` caller that follows it would let a launch failure escape.
`canRunArrowOptimization` catches `NonFatal`, so it is unaffected today.
**Verification:**
- **Inspection:** The documented exceptions match what the check throws for
a missing executable, a timeout and an unexpected exit code.
--
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]
