viirya commented on code in PR #58978:
URL: https://github.com/apache/spark/pull/58978#discussion_r4333614178
##########
dev/spark-test-image/python-312/Dockerfile:
##########
@@ -15,6 +15,17 @@
# limitations under the License.
#
+# Reads the JEP and CDI versions from this checkout, just as the Spark build
does. Only the
+# version files reach the image below, so its layers change only when a
version does.
+FROM ubuntu:noble AS spark-versions
+COPY --from=root pom.xml /pom.xml
+RUN mkdir /versions && \
+ for name in jep arrow; do \
+ sed -n "s:.*<${name}.version>\(.*\)</${name}.version>.*:\1:p" /pom.xml |
head -n 1 \
+ > "/versions/${name}" && \
+ test -s "/versions/${name}"; \
Review Comment:
Fixed in 961ebc6: `test -s "/versions/${name}" || exit 1`.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/InProcessArrowEvalPythonEvaluatorFactory.scala:
##########
@@ -0,0 +1,598 @@
+/*
+ * 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 java.io.File
+import java.nio.file.Files
+import java.util.UUID
+import java.util.concurrent.TimeUnit
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicInteger}
+import java.util.concurrent.locks.ReentrantLock
+
+import scala.collection.mutable.ArrayBuffer
+import scala.jdk.CollectionConverters._
+
+import com.google.common.util.concurrent.Uninterruptibles
+import org.apache.arrow.c.{ArrowArray, ArrowSchema, BaseStruct}
+import org.apache.arrow.util.AutoCloseables
+import org.apache.arrow.vector.{BaseLargeVariableWidthVector,
BaseVariableWidthVector, DensityAwareVector, VectorSchemaRoot}
+
+import org.apache.spark.{SparkEnv, SparkException, TaskContext}
+import org.apache.spark.api.python.ChainedPythonFunctions
+import org.apache.spark.memory.MemoryConsumer
+import org.apache.spark.serializer.SerializerManager
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression,
JoinedRow, PythonUDF, UnsafeProjection, UnsafeRow}
+import org.apache.spark.sql.execution.arrow.ArrowWriter
+import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.sql.execution.python.EvalPythonExec.ArgumentMetadata
+import org.apache.spark.sql.types._
+import org.apache.spark.sql.util.ArrowUtils
+import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch,
ColumnVector}
+import org.apache.spark.util.Utils
+
+/**
+ * Evaluates scalar Python UDFs using Arrow CDI in the executor process. Only
UDF arguments
+ * are converted to Arrow. Original rows are buffered in a spillable queue and
joined with
+ * the results, unless all of them are UDF arguments that read back from Arrow
unchanged.
+ * Each batch owns its Arrow buffers so Python can safely retain input arrays.
+ *
+ * The evaluator owns its queue, so that cleanup at task completion is
coordinated with a
+ * consumer on another thread, such as a pipelined Python writer or a
TRANSFORM feed thread.
+ */
+class InProcessArrowEvalPythonEvaluatorFactory(
+ childOutput: Seq[Attribute],
+ udfs: Seq[PythonUDF],
+ output: Seq[Attribute],
+ batchSize: Int,
+ maxBytes: Long,
+ timeZoneId: String,
+ largeVarTypes: Boolean,
+ hideTraceback: Boolean,
+ simplifiedTraceback: Boolean,
+ tracebackWithLocals: Boolean,
+ fullValidation: Boolean,
+ metrics: Map[String, SQLMetric])
+ extends EvalPythonEvaluatorFactory(childOutput, udfs, output) {
+
+ private[python] def runtimeSession:
InProcessPythonRuntime.InterpreterSession =
+ InProcessPythonRuntime.currentSession
+
+ /** Unused: `evaluateJoined` always evaluates the UDFs. */
+ override protected def evaluate(
+ funcs: Seq[(ChainedPythonFunctions, Long)],
+ argMetas: Array[Array[ArgumentMetadata]],
+ rows: Iterator[InternalRow],
+ inputSchema: StructType,
+ context: TaskContext): Iterator[InternalRow] =
+ throw SparkException.internalError("In-process UDFs are evaluated with
their input rows")
+
+ override protected def evaluateJoined(
+ funcs: Seq[(ChainedPythonFunctions, Long)],
+ argMetas: Array[Array[ArgumentMetadata]],
+ rows: Iterator[InternalRow],
+ inputs: Seq[Expression],
+ inputSchema: StructType,
+ context: TaskContext): Option[Iterator[InternalRow]] = {
+ import InProcessArrowEvalPythonEvaluatorFactory.{Buffered, ReadBack,
readsBack}
+ val inputColumns = inputs.length == childOutput.length &&
inputs.zip(childOutput).forall {
+ case (a: Attribute, c) => a.exprId == c.exprId
+ case _ => false
+ }
+ // If all input columns are UDF arguments, they are written to Arrow
regardless. Read them
+ // back from the exported input vectors instead of buffering every input
row, if their
+ // values read back from Arrow exactly as written and as fast as an unsafe
row copy.
+ val joinInput = if (inputColumns && inputSchema.forall(f =>
readsBack(f.dataType))) {
+ ReadBack
+ } else if (inputColumns) {
+ Buffered(None)
+ } else {
+ // Each projected row is written to Arrow before the next input row is
pulled, so the
+ // arguments go into a reused buffer rather than being copied value by
value.
+ val projection = UnsafeProjection.create(inputs, childOutput)
+ projection.initialize(context.partitionId())
+ Buffered(Some(projection))
+ }
+ Some(evaluateBatches(funcs, argMetas, rows, inputSchema, context,
joinInput))
+ }
+
+ private[python] def evaluateBatches(
+ funcs: Seq[(ChainedPythonFunctions, Long)],
+ argMetas: Array[Array[ArgumentMetadata]],
+ rows: Iterator[InternalRow],
+ inputSchema: StructType,
+ context: TaskContext,
+ joinInput: InProcessArrowEvalPythonEvaluatorFactory.JoinInput):
Iterator[InternalRow] = {
+ import InProcessArrowEvalPythonEvaluatorFactory.{Buffered, ReadBack}
+ ArrowUtils.failDuplicatedFieldNames(inputSchema)
+ val functions = funcs.map { case (chain, _) =>
+ if (chain.funcs.size != 1) {
+ throw SparkException.internalError(
+ "In-process UDF chains must use separate evaluation nodes")
+ }
+ chain.funcs.head
+ }
+ val inputOrdinals = argMetas.map(_.map(_.offset))
+ def checkCancellation(): Unit = context.killTaskIfInterrupted()
+
+ val expectedFields = udfs.map { udf =>
+ ArrowUtils.toArrowField("result", udf.dataType, true, timeZoneId,
largeVarTypes)
+ }
+ val arrowSchema = ArrowUtils.toArrowSchema(inputSchema, timeZoneId,
largeVarTypes)
+ // Capture before consuming input: an old task must never join a later
context's session.
+ val runtime = runtimeSession
+ // Rows are copied out of the queue and Arrow vectors before they are
returned, so they
+ // remain valid after task completion releases those, on whichever thread
consumes them.
+ val resultProj = UnsafeProjection.create(output, output)
+ // Spill files go into a directory of the queue's own, created with the
first disk queue, so
+ // that task completion can delete them when it cannot close the queue.
+ @volatile var spillDir: File = null
Review Comment:
Done in 961ebc6: `spillDir` and the abandoned flag are private fields of
`InProcessRowQueue`, which overrides `close()` to delete the directory, and
`releaseTaskMemory` is now `queue.close()`.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/InProcessPythonUDFSuite.scala:
##########
@@ -0,0 +1,540 @@
+/*
+ * 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 java.io.File
+import java.util.Properties
+import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue, TimeUnit}
+import java.util.concurrent.atomic.AtomicReference
+
+import scala.jdk.CollectionConverters._
+
+import org.apache.spark.{SparkEnv, SparkException, TaskContextImpl}
+import org.apache.spark.api.python.PythonEvalType
+import org.apache.spark.internal.config.{BUFFER_PAGESIZE, PLUGINS}
+import org.apache.spark.memory.{TaskMemoryManager, TestMemoryConsumer,
TestMemoryManager}
+import org.apache.spark.sql.{AnalysisException, Column, QueryTest}
+import org.apache.spark.sql.api.python.PythonSQLUtils
+import org.apache.spark.sql.catalyst.expressions.PythonUDF
+import org.apache.spark.sql.catalyst.plans.logical.{Aggregate,
ArrowEvalPython, Filter, LocalLimit}
+import org.apache.spark.sql.execution.{GlobalLimitExec, ProjectExec, SortExec}
+import org.apache.spark.sql.functions._
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.types.{CharType, LongType, StructType, VarcharType}
+import org.apache.spark.util.Utils
+
+/**
+ * Planning regressions, and evaluator tests that need no Python; runtime
coverage lives in
+ * the PySpark integration suite.
+ */
+class InProcessPythonUDFSuite extends QueryTest with SharedSparkSession {
+
+ import InProcessEvaluatorTestUtils._
+ import testImplicits._
+
+ private val plugin =
"org.apache.spark.sql.execution.python.InProcessPythonPlugin"
+
+ override def beforeEach(): Unit = {
+ super.beforeEach()
+ // These tests plan queries without loading a native interpreter.
Advertise the plugin
+ // after context creation; actual plugin initialization is covered by
integration tests.
+ SparkEnv.get.conf.set(PLUGINS, Seq(plugin))
+ }
+
+ override def afterEach(): Unit = {
+ try { SparkEnv.get.conf.remove(PLUGINS) } finally { super.afterEach() }
+ }
+
+ private def makeUDF(
+ name: String,
+ input: Column,
+ deterministic: Boolean = true): Column = {
+ // Each call creates fresh bytes, as Py4J does. Semantic equality must
compare their contents.
+ InProcessPythonUDFBuilder.create(name, Array[Byte](1, 2), LongType.json,
"3.11")
+ .apply(Seq(input).asJava, deterministic)
+ }
+
+ test("in-process UDFs use PythonUDF and ArrowEvalPython planning contracts")
{
+ val df = spark.range(10)
+ val doubled = makeUDF("double", df("id"))
+ val expr = doubled.expr.asInstanceOf[PythonUDF]
+ assert(expr.evalType == PythonEvalType.SQL_SCALAR_ARROW_INPROCESS_UDF)
+ assert(expr.expensive)
+ assert(expr.semanticEquals(makeUDF("double", df("id")).expr))
+
+ val query = df.select(doubled)
+ val eval = query.queryExecution.optimizedPlan.collect { case p:
ArrowEvalPython => p }
+ assert(eval.size == 1)
+ assert(eval.head.evalType == PythonEvalType.SQL_SCALAR_ARROW_INPROCESS_UDF)
+ val physical = query.queryExecution.executedPlan.collect {
+ case p: InProcessArrowEvalPythonExec => p
+ }
+ assert(physical.size == 1)
+ assert(physical.head.producedAttributes ==
+ (physical.head.outputSet -- physical.head.child.outputSet))
+ assert(physical.head.missingInput.isEmpty)
+ // There are no Python workers, so no boot time and descriptions of
in-process execution.
+ val metrics = physical.head.metrics
+ assert(metrics.keySet == Set("pythonDataSent", "pythonDataReceived",
"pythonInitTime",
+ "pythonTotalTime", "pythonProcessingTime", "pythonNumRowsReceived"))
+ assert(!metrics.values.exists(_.name.exists(_.contains("worker"))))
+ // Times are summed in nanoseconds, so sub-millisecond registrations and
batches count.
+ assert(Seq("pythonInitTime", "pythonTotalTime", "pythonProcessingTime")
+ .forall(metrics(_).metricType == "nsTiming"))
+ }
+
+ test("a committed write that invalidates a cached in-process plan does not
fail") {
+ withTempPath { dir =>
+ val path = dir.getCanonicalPath
+ spark.range(3).write.parquet(path)
+ val cached = spark.read.parquet(path).select(makeUDF("identity",
col("id")))
+ cached.cache()
+ try {
+
assert(spark.sharedState.cacheManager.lookupCachedData(cached).nonEmpty)
+ // Re-caching plans the entry in this session, which rejects
in-process UDFs.
+ withSQLConf(SQLConf.PYTHON_UDF_PROFILER.key -> "perf") {
+ spark.range(3, 5).write.mode("append").parquet(path)
+ }
+ assert(spark.sharedState.cacheManager.lookupCachedData(cached).isEmpty)
+ assert(spark.read.parquet(path).count() == 5)
+ } finally {
+ cached.unpersist()
+ }
+ }
+ }
+
+ test("unsupported configuration added after column creation fails before
task submission") {
+ val column = makeUDF("identity", col("id"))
+ for (partitionEvaluator <- Seq("true", "false")) {
+ withSQLConf(
+ SQLConf.USE_PARTITION_EVALUATOR.key -> partitionEvaluator,
+ SQLConf.PYTHON_UDF_PROFILER.key -> "perf") {
+ val error = intercept[SparkException] {
+ spark.range(1).select(column).queryExecution.executedPlan
+ }
+ checkError(
+ exception = error,
+ condition = "INVALID_SPARK_CONFIG.UNSUPPORTED_IN_PROCESS_PYTHON_UDF",
+ parameters = Map("config" -> SQLConf.PYTHON_UDF_PROFILER.key))
+ }
+ }
+ }
+
+ test("legacy Python profilers are rejected") {
+ for (key <- Seq("spark.python.profile", "spark.python.profile.memory")) {
+ SparkEnv.get.conf.set(key, "true")
+ try {
+ checkError(
+ exception = intercept[SparkException] {
+
InProcessPythonUDFBuilder.checkConfiguration(spark.sessionState.conf)
+ },
+ condition = "INVALID_SPARK_CONFIG.UNSUPPORTED_IN_PROCESS_PYTHON_UDF",
+ parameters = Map("config" -> key))
+ } finally {
+ SparkEnv.get.conf.remove(key)
+ }
+ // As in PySpark, other values leave the profiler off, without failing
to parse.
+ for (value <- Seq("True", "0", "off")) {
Review Comment:
Done in 961ebc6: both blocks use `withSparkEnvConfs`.
##########
python/pyspark/sql/tests/test_inprocess_udf.py:
##########
@@ -0,0 +1,1992 @@
+#
+# 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.
+#
+
+"""End-to-end tests for in-process Python UDFs.
+
+Run with python/run-tests like other SQL tests. JEP paths are discovered from
the
+selected Python environment before the Spark JVM starts. ARROW_C_DATA_JAR must
+point to the provided Arrow CDI JAR. Set INPROCESS_TESTS=1
+to require the suite (missing dependencies then fail), or 0 to disable it.
+Otherwise, the suite runs when JEP, PyArrow and the CDI JAR are available.
+"""
+
+import os
+import shutil
+import tempfile
+import unittest
+import zipfile
+from importlib.util import find_spec
+from pathlib import Path
+from unittest.mock import patch
+
+from pyspark.testing.sqlutils import ReusedSQLTestCase
+from pyspark.testing.utils import eventually
+
+_jep_spec = find_spec("jep")
+_cdi_jar = os.environ.get("ARROW_C_DATA_JAR")
+_test_mode = os.environ.get("INPROCESS_TESTS")
+_run_inprocess = _test_mode == "1" or (
+ _test_mode != "0"
+ and _jep_spec is not None
+ and find_spec("pyarrow") is not None
Review Comment:
Fixed in 961ebc6: the suite uses `have_pyarrow`.
--
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]
