You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by vi...@apache.org on 2026/09/26 16:15:02 UTC
(spark) branch branch-4.x updated: [SPARK-59744][UDF] Execute scalar external UDFs through worker sessions
This is an automated email from the ASF dual-hosted git repository.
viirya pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 3735d1590bfb [SPARK-59744][UDF] Execute scalar external UDFs through worker sessions
3735d1590bfb is described below
commit 3735d1590bfb830dff40fd27ecb523c43a6c10be
Author: Haiyang Sun <ha...@databricks.com>
AuthorDate: Sat Sep 26 09:14:42 2026 -0700
[SPARK-59744][UDF] Execute scalar external UDFs through worker sessions
### What changes were proposed in this pull request?
This PR adds the scalar external UDF execution bridge:
* Execute scalar external UDFs through dispatcher-managed worker sessions.
* Project UDF arguments into row- and byte-limited Arrow batches, then join worker output back to the original rows with `HybridRowQueue`.
* Build the existing Python-shaped `Init` message through a temporary `PythonInitAdapter`, including the Arrow schemas, Python runner configuration, and complete task context.
* Mark the temporary Python UDF payload format as `experimental` while the protocol is still evolving.
* Reject named arguments until [SPARK-59745](https://issues.apache.org/jira/browse/SPARK-59745) adds metadata preservation and worker binding.
* Treat external UDF results as nullable and report output cardinality mismatches with a structured error.
* Add focused execution and planning coverage for batching, multiple partitions, empty input, malformed or mismatched output, early termination, nullability, and named-argument rejection.
This PR does not add a worker implementation. The dispatcher is intentionally unconfigured, so worker creation fails before starting a process. SPARK-59364 will replace the temporary Python-specific initialization dependency separately.
### Why are the changes needed?
The external UDF framework can plan scalar UDFs and manage worker sessions, but `ExecuteExternalUDFExec` does not yet exchange rows with those sessions. This change supplies that execution path while keeping worker-server implementation and language-agnostic initialization out of scope.
### Does this PR introduce _any_ user-facing change?
No. This extends unreleased external UDF infrastructure. Existing Python UDF execution is unchanged, and no external UDF worker is enabled by this PR.
### How was this patch tested?
Added and extended tests covering Arrow row- and byte-based batching, raw large-variable Arrow request and response types, multiple batches and partitions, empty input, ordered multi-argument data, nullable output, named-argument rejection, exact task-context values and JSON encodings, malformed responses, output cardinality mismatches, early iterator termination, failed-task session and Arrow-resource cleanup, and planning behavior.
After rebasing onto the latest `master`, ran:
```
build/sbt 'sql/testOnly org.apache.spark.sql.execution.externalUDF.ExecuteExternalUDFExecSuite'
```
All 13 tests passed.
### Was this patch authored or co-authored using generative AI tooling?
Yes
Closes #58988 from haiyangsun-db/codex/SPARK-55278-external-udf-execution.
Authored-by: Haiyang Sun <ha...@databricks.com>
Signed-off-by: Liang-Chi Hsieh <vi...@gmail.com>
(cherry picked from commit f06bbad6fcc71d6f29ca024b4562d3e243a88b0b)
Signed-off-by: Liang-Chi Hsieh <vi...@gmail.com>
---
.../src/main/resources/error/error-conditions.json | 11 +
.../expressions/ExternalUserDefinedFunction.scala | 8 +-
.../spark/sql/errors/QueryExecutionErrors.scala | 4 +
.../sql/execution/arrow/ArrowConverters.scala | 69 ++-
.../externalUDF/ExecuteExternalUDFExec.scala | 290 +++++++++-
.../execution/externalUDF/ExternalUDFExec.scala | 27 +-
.../execution/externalUDF/PlanExternalUDFs.scala | 7 +-
.../execution/externalUDF/PythonInitAdapter.scala | 104 ++++
.../externalUDF/ExecuteExternalUDFExecSuite.scala | 631 +++++++++++++++++++++
.../externalUDF/PlanExternalUDFsSuite.scala | 79 ++-
10 files changed, 1181 insertions(+), 49 deletions(-)
diff --git a/common/utils/src/main/resources/error/error-conditions.json b/common/utils/src/main/resources/error/error-conditions.json
index 2dbffd93557c..fb64db10cfa5 100644
--- a/common/utils/src/main/resources/error/error-conditions.json
+++ b/common/utils/src/main/resources/error/error-conditions.json
@@ -7087,6 +7087,12 @@
],
"sqlState" : "42939"
},
+ "RESULT_ROWS_MISMATCH" : {
+ "message" : [
+ "External UDF output row count was <output_length>; expected <input_length>. An external UDF must produce exactly one row for each input row."
+ ],
+ "sqlState" : "21000"
+ },
"ROUTINE_ALREADY_EXISTS" : {
"message" : [
"Cannot create the <newRoutineType> <routineName> because a <existingRoutineType> of that name already exists.",
@@ -9068,6 +9074,11 @@
"External UDF expressions are disabled. Set <config> to true to enable planning."
]
},
+ "EXTERNAL_UDF_IN_BARRIER_TASK" : {
+ "message" : [
+ "External UDF execution in a barrier task without callback transport."
+ ]
+ },
"EXTERNAL_UDF_IN_ON_CLAUSE" : {
"message" : [
"External UDF in the ON clause of a <joinType> JOIN. In case of an INNER JOIN consider rewriting to a CROSS JOIN with a WHERE clause."
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ExternalUserDefinedFunction.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ExternalUserDefinedFunction.scala
index 0bbaf1ee242d..e02818bd29bb 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ExternalUserDefinedFunction.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ExternalUserDefinedFunction.scala
@@ -47,7 +47,8 @@ import org.apache.spark.udf.worker.UDFWorkerSpecification
* @param children Input argument expressions.
* @param inputTypes Optional declared input types for validation.
* @param udfDeterministic Whether this UDF is deterministic.
- * @param udfNullable Whether this UDF can return null.
+ * @param udfNullable Declared return nullability. Scalar external UDF planning currently
+ * retains this metadata but treats every result as nullable.
* @param resultId Unique expression ID for this invocation.
*/
@Experimental
@@ -65,7 +66,10 @@ case class ExternalUserDefinedFunction(
override lazy val deterministic: Boolean = udfDeterministic && children.forall(_.deterministic)
- override def nullable: Boolean = udfNullable
+ // Match PythonUDF: an external worker may return null even when the function metadata declares
+ // a non-nullable result.
+ // TODO(SPARK-55278): Honor declared non-nullability once the worker protocol can enforce it.
+ override def nullable: Boolean = true
override def checkInputDataTypes(): TypeCheckResult = {
inputTypes match {
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala
index 0883f1e89e95..fed5180c9216 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala
@@ -655,6 +655,10 @@ private[sql] object QueryExecutionErrors extends QueryErrorsBase with ExecutionE
messageParameters = Map("methodName" -> methodName))
}
+ def externalUDFInBarrierTaskUnsupportedError(): SparkUnsupportedOperationException = {
+ new SparkUnsupportedOperationException("UNSUPPORTED_FEATURE.EXTERNAL_UDF_IN_BARRIER_TASK")
+ }
+
def binaryArithmeticCauseOverflowError(
eval1: Short,
symbol: String,
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowConverters.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowConverters.scala
index 464cac157b25..0441da2f17ff 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowConverters.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowConverters.scala
@@ -17,7 +17,8 @@
package org.apache.spark.sql.execution.arrow
-import java.io.{ByteArrayInputStream, ByteArrayOutputStream, FileInputStream, OutputStream}
+import java.io.{ByteArrayInputStream, ByteArrayOutputStream, FileInputStream, InputStream,
+ OutputStream}
import java.nio.channels.{Channels, ReadableByteChannel}
import scala.collection.mutable.ArrayBuffer
@@ -100,6 +101,7 @@ private[sql] object ArrowConverters extends Logging {
SQLConf.get.arrowCompressionCodec, SQLConf.get.arrowZstdCompressionLevel)
protected val unloader = new VectorUnloader(root, true, codec, true)
protected val arrowWriter = ArrowWriter.create(root)
+ protected def maxBytesPerBatch: Long = -1L
Option(context).foreach {_.addTaskCompletionListener[Unit] { _ =>
close()
@@ -115,15 +117,21 @@ private[sql] object ArrowConverters extends Logging {
Utils.tryWithSafeFinally {
var rowCount = 0L
- while (rowIter.hasNext && (maxRecordsPerBatch <= 0 || rowCount < maxRecordsPerBatch)) {
+ while (rowIter.hasNext &&
+ (maxRecordsPerBatch <= 0 || rowCount < maxRecordsPerBatch) &&
+ (maxBytesPerBatch <= 0 || rowCount == 0 ||
+ arrowWriter.sizeInBytes() < maxBytesPerBatch)) {
val row = rowIter.next()
arrowWriter.write(row)
rowCount += 1
}
arrowWriter.finish()
val batch = unloader.getRecordBatch()
- bytes = serializeBatch(batch)
- batch.close()
+ try {
+ bytes = serializeBatch(batch)
+ } finally {
+ batch.close()
+ }
} {
arrowWriter.reset()
}
@@ -137,6 +145,24 @@ private[sql] object ArrowConverters extends Logging {
}
}
+ private[sql] class SizeLimitedArrowBatchIterator(
+ rows: Iterator[InternalRow],
+ schema: StructType,
+ maxRecordsPerBatch: Long,
+ override protected val maxBytesPerBatch: Long,
+ timeZoneId: String,
+ errorOnDuplicatedFieldNames: Boolean,
+ largeVarTypes: Boolean,
+ context: TaskContext)
+ extends ArrowBatchIterator(
+ rows,
+ schema,
+ maxRecordsPerBatch,
+ timeZoneId,
+ errorOnDuplicatedFieldNames,
+ largeVarTypes,
+ context)
+
private[sql] class ArrowBatchWithSchemaIterator(
rowIter: Iterator[InternalRow],
schema: StructType,
@@ -231,6 +257,31 @@ private[sql] object ArrowConverters extends Logging {
context)
}
+ /**
+ * Maps an iterator of internal rows to serialized Arrow record batches, limiting each batch by
+ * record count and a best-effort byte target. A batch can exceed the byte target because its size
+ * is checked before appending each row, and the first row is always accepted.
+ */
+ private[sql] def toBatchIterator(
+ rowIter: Iterator[InternalRow],
+ schema: StructType,
+ maxRecordsPerBatch: Long,
+ maxBytesPerBatch: Long,
+ timeZoneId: String,
+ errorOnDuplicatedFieldNames: Boolean,
+ largeVarTypes: Boolean,
+ context: TaskContext): ArrowBatchIterator = {
+ new SizeLimitedArrowBatchIterator(
+ rowIter,
+ schema,
+ maxRecordsPerBatch,
+ maxBytesPerBatch,
+ timeZoneId,
+ errorOnDuplicatedFieldNames,
+ largeVarTypes,
+ context)
+ }
+
/**
* Convert the input rows into fully contained arrow batches.
* Different from [[toBatchIterator]], each output arrow batch starts with the schema.
@@ -514,6 +565,16 @@ private[sql] object ArrowConverters extends Logging {
new ReadChannel(Channels.newChannel(in)), allocator) // throws IOException
}
+ /**
+ * Load a serialized Arrow record batch from an input stream.
+ */
+ private[sql] def loadBatch(
+ batchInput: InputStream,
+ allocator: BufferAllocator): ArrowRecordBatch = {
+ MessageSerializer.deserializeRecordBatch(
+ new ReadChannel(Channels.newChannel(batchInput)), allocator) // throws IOException
+ }
+
private[arrow] def serializeBatch(batch: ArrowRecordBatch): Array[Byte] = {
val out = new ByteArrayOutputStream()
val writeChannel = new WriteChannel(Channels.newChannel(out))
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/ExecuteExternalUDFExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/ExecuteExternalUDFExec.scala
index 69e93a0c5911..82caa97f4df4 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/ExecuteExternalUDFExec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/ExecuteExternalUDFExec.scala
@@ -17,24 +17,54 @@
package org.apache.spark.sql.execution.externalUDF
+import java.io.ByteArrayOutputStream
+import java.nio.channels.Channels
+
+import scala.jdk.CollectionConverters._
+import scala.util.control.NonFatal
+
+import com.google.protobuf.UnsafeByteOperations
+import org.apache.arrow.memory.BufferAllocator
+import org.apache.arrow.vector.{VectorLoader, VectorSchemaRoot}
+import org.apache.arrow.vector.ipc.WriteChannel
+import org.apache.arrow.vector.ipc.message.MessageSerializer
+import org.apache.arrow.vector.util.ValueVectorUtility
+
+import org.apache.spark.{BarrierTaskContext, SparkException, TaskContext}
import org.apache.spark.annotation.Experimental
import org.apache.spark.rdd.RDD
import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.expressions.{
- Attribute,
- AttributeSet,
- ExternalUserDefinedFunction
-}
+import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeSet, Expression,
+ ExternalUserDefinedFunction, JoinedRow, MutableProjection, UnsafeProjection, UnsafeRow}
import org.apache.spark.sql.errors.QueryExecutionErrors
import org.apache.spark.sql.execution.SparkPlan
-import org.apache.spark.udf.worker.UDFWorkerSpecification
+import org.apache.spark.sql.execution.arrow.ArrowConverters
+import org.apache.spark.sql.execution.externalUDF.ExecuteExternalUDFExec._
+import org.apache.spark.sql.execution.python.HybridRowQueue
+import org.apache.spark.sql.types.{StructField, StructType}
+import org.apache.spark.sql.util.ArrowUtils
+import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch, ColumnVector}
+import org.apache.spark.udf.worker.{DataRequest, DataResponse, UDFWorkerSpecification}
/**
* :: Experimental ::
- * Physical plan node representing one scalar UDF evaluation. Worker execution is intentionally
- * unimplemented until the scalar worker path is added.
+ * Physical plan node that evaluates one scalar UDF in an external worker process.
+ * Like [[org.apache.spark.sql.execution.python.ArrowEvalPythonExec]], this node projects UDF
+ * arguments, queues original rows in [[HybridRowQueue]], and joins each response to its row.
+ * It sends projected Arrow batches through an external-UDF dispatcher session instead of the
+ * legacy PySpark runner.
+ *
+ * Each response must contain one Arrow record batch matching the output schema sent in `Init`.
+ * Record batches do not carry a schema, so the engine validates their physical buffer layout but
+ * cannot distinguish logical types with identical Arrow layouts. Arrow record-batch compression
+ * is self-describing, so responses may use any codec available from Arrow's compression factory.
*
- * @param udf UDF expression represented by this node.
+ * The dispatcher is intentionally not configured until a worker implementation is added, so
+ * worker creation currently fails before any process is started. Before enabling a production
+ * dispatcher, its transport must also add application-level flow control so a slow worker cannot
+ * buffer an entire partition in the request stream and [[HybridRowQueue]].
+ *
+ * @param udf UDF expression evaluated by the worker session.
* @param resultAttr Output attribute for the UDF expression.
* @param child Child plan providing input rows.
*/
@@ -52,11 +82,247 @@ case class ExecuteExternalUDFExec(
override def producedAttributes: AttributeSet = AttributeSet(Seq(resultAttr))
override protected def doExecute(): RDD[InternalRow] = {
- // TODO(SPARK-55278): Stream rows to and from the worker through session.process().
- throw QueryExecutionErrors.methodNotImplementedError(
- "ExecuteExternalUDFExec.doExecute")
+ val argumentExpressions: Seq[Expression] = udf.children
+ val inputSchema = StructType(
+ argumentExpressions.zipWithIndex.map { case (expression, index) =>
+ StructField(s"_$index", expression.dataType, expression.nullable)
+ })
+ val outputSchema = StructType(Seq(StructField("_0", udf.dataType, nullable = true)))
+ val timeZoneId = conf.sessionLocalTimeZone
+ val largeVarTypes = conf.arrowUseLargeVarTypes
+ val maxRecordsPerBatch = conf.arrowMaxRecordsPerBatch
+ val maxBytesPerBatch = conf.arrowMaxBytesPerBatch
+ val preparedInit = prepareInit(inputSchema, outputSchema, timeZoneId, largeVarTypes)
+
+ child.execute().mapPartitionsInternal { rows =>
+ val context = TaskContext.get()
+ ensureTaskContextSupported(context)
+ val projection = MutableProjection.create(argumentExpressions, child.output)
+ projection.initialize(context.partitionId())
+
+ val queue = HybridRowQueue(context.taskMemoryManager(), child.output.length)
+ context.addTaskCompletionListener[Unit](_ => queue.close())
+
+ var inputRowCount = 0L
+ var outputRowCount = 0L
+ val projectedRows = rows.map { row =>
+ queue.add(row.asInstanceOf[UnsafeRow])
+ inputRowCount += 1
+ projection(row)
+ }
+
+ withUDFWorkerSession(context) { session =>
+ session.init(PythonInitAdapter.build(preparedInit, context))
+
+ val batches = ArrowConverters.toBatchIterator(
+ projectedRows,
+ inputSchema,
+ maxRecordsPerBatch,
+ maxBytesPerBatch,
+ timeZoneId,
+ true,
+ largeVarTypes,
+ context)
+ val requests = batches.map { batch =>
+ DataRequest.newBuilder()
+ .setData(UnsafeByteOperations.unsafeWrap(batch))
+ .build()
+ }
+ val udfResults = new ArrowResponseRowIterator(
+ session.process(requests),
+ outputSchema,
+ timeZoneId,
+ largeVarTypes,
+ context)
+
+ val joined = new JoinedRow
+ val resultProjection = UnsafeProjection.create(output, output)
+ resultProjection.initialize(context.partitionId())
+ var cardinalityVerified = false
+
+ new Iterator[InternalRow] {
+ override def hasNext: Boolean = {
+ val hasNextResult = udfResults.hasNext
+ if (!hasNextResult && !cardinalityVerified) {
+ cardinalityVerified = true
+ if (inputRowCount != outputRowCount) {
+ throw cardinalityMismatch(inputRowCount, outputRowCount)
+ }
+ }
+ hasNextResult
+ }
+
+ override def next(): InternalRow = {
+ if (!hasNext) Iterator.empty.next()
+ if (queue.getNumElementsQueued() == 0L) {
+ val outputRowsLowerBound = outputRowCount + udfResults.bufferedRowCount
+ throw cardinalityMismatch(inputRowCount, s"at least $outputRowsLowerBound")
+ }
+ val udfResult = udfResults.next()
+ outputRowCount += 1
+ resultProjection(joined(queue.remove(), udfResult))
+ }
+ }
+ }
+ }
+ }
+
+ private def prepareInit(
+ inputSchema: StructType,
+ outputSchema: StructType,
+ timeZoneId: String,
+ largeVarTypes: Boolean): PythonInitAdapter.PreparedInit = {
+ // TODO(SPARK-59364): Make this node language-agnostic. Its dependency on the
+ // Python-specific Init builder is temporary and tracked by this Spark ticket.
+ PythonInitAdapter.prepare(
+ workerSpec,
+ udf,
+ serializeArrowSchema(inputSchema, timeZoneId, largeVarTypes),
+ serializeArrowSchema(outputSchema, timeZoneId, largeVarTypes),
+ timeZoneId,
+ PythonInitAdapter.sessionConf(conf, inputSchema))
}
override protected def withNewChildInternal(newChild: SparkPlan): ExecuteExternalUDFExec =
copy(child = newChild)
}
+
+object ExecuteExternalUDFExec {
+ private[externalUDF] def ensureTaskContextSupported(context: TaskContext): Unit = {
+ if (context.isInstanceOf[BarrierTaskContext]) {
+ throw QueryExecutionErrors.externalUDFInBarrierTaskUnsupportedError()
+ }
+ }
+
+ private def serializeArrowSchema(
+ schema: StructType,
+ timeZoneId: String,
+ largeVarTypes: Boolean): Array[Byte] = {
+ val arrowSchema = ArrowUtils.toArrowSchema(
+ schema,
+ timeZoneId,
+ true,
+ largeVarTypes)
+ val buffer = new ByteArrayOutputStream()
+ MessageSerializer.serialize(new WriteChannel(Channels.newChannel(buffer)), arrowSchema)
+ buffer.toByteArray
+ }
+
+ // Keep this condition and its parameter names aligned with verify_result_row_count in
+ // python/pyspark/eval_handlers/verification.py. PySpark and the JVM use separate error catalogs,
+ // so the condition must be registered in both.
+ private def cardinalityMismatch(inputRows: Long, outputRows: Long): SparkException = {
+ cardinalityMismatch(inputRows, outputRows.toString)
+ }
+
+ private def cardinalityMismatch(inputRows: Long, outputRows: String): SparkException = {
+ new SparkException(
+ errorClass = "RESULT_ROWS_MISMATCH",
+ messageParameters = Map(
+ "output_length" -> outputRows,
+ "input_length" -> inputRows.toString),
+ cause = null)
+ }
+
+ /** Decodes Arrow record-batch responses using the schema supplied in Init. */
+ private class ArrowResponseRowIterator(
+ responses: Iterator[DataResponse],
+ expectedSchema: StructType,
+ timeZoneId: String,
+ largeVarTypes: Boolean,
+ context: TaskContext)
+ extends Iterator[InternalRow] {
+
+ private var allocator: Option[BufferAllocator] = None
+ private var root: Option[VectorSchemaRoot] = None
+ private var columnarBatch: Option[ColumnarBatch] = None
+ private var rows: Iterator[InternalRow] = Iterator.empty
+ private var initialized = false
+ private var exhausted = false
+ private var closed = false
+ private var rowsRemainingInBatch = 0L
+
+ context.addTaskCompletionListener[Unit](_ => close())
+
+ def bufferedRowCount: Long = rowsRemainingInBatch
+
+ override def hasNext: Boolean = {
+ try {
+ val (activeAllocator, activeRoot, activeColumnarBatch) = initialize()
+ while (!rows.hasNext && !exhausted) {
+ if (responses.hasNext) {
+ val batch = ArrowConverters.loadBatch(
+ responses.next().getData.newInput(),
+ activeAllocator)
+ try {
+ new VectorLoader(activeRoot).load(batch)
+ ValueVectorUtility.validate(activeRoot)
+ } finally {
+ batch.close()
+ }
+ activeColumnarBatch.setNumRows(activeRoot.getRowCount)
+ rows = activeColumnarBatch.rowIterator().asScala
+ rowsRemainingInBatch = activeRoot.getRowCount
+ } else {
+ exhausted = true
+ close()
+ }
+ }
+ rows.hasNext
+ } catch {
+ case NonFatal(error) =>
+ close()
+ throw error
+ }
+ }
+
+ override def next(): InternalRow = {
+ if (!hasNext) Iterator.empty.next()
+ val row = rows.next()
+ rowsRemainingInBatch -= 1
+ row
+ }
+
+ private def initialize(): (BufferAllocator, VectorSchemaRoot, ColumnarBatch) = {
+ if (!initialized) {
+ initialized = true
+ val newAllocator = ArrowUtils.rootAllocator.newChildAllocator(
+ "externalUdfArrowResponse",
+ 0,
+ Long.MaxValue)
+ allocator = Some(newAllocator)
+ val newRoot = VectorSchemaRoot.create(
+ ArrowUtils.toArrowSchema(
+ expectedSchema,
+ timeZoneId,
+ true,
+ largeVarTypes),
+ newAllocator)
+ root = Some(newRoot)
+ val columns: Array[ColumnVector] =
+ newRoot.getFieldVectors.asScala.iterator.map { vector =>
+ new ArrowColumnVector(vector): ColumnVector
+ }.toArray
+ columnarBatch = Some(new ColumnarBatch(columns))
+ }
+ (allocator, root, columnarBatch) match {
+ case (Some(activeAllocator), Some(activeRoot), Some(activeColumnarBatch)) =>
+ (activeAllocator, activeRoot, activeColumnarBatch)
+ case _ =>
+ throw SparkException.internalError(
+ "The Arrow response row iterator was not initialized correctly.")
+ }
+ }
+
+ private def close(): Unit = {
+ if (!closed) {
+ closed = true
+ try {
+ root.foreach(_.close())
+ } finally {
+ allocator.foreach(_.close())
+ }
+ }
+ }
+ }
+}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/ExternalUDFExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/ExternalUDFExec.scala
index 57630ce1cba9..db88ef90d62b 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/ExternalUDFExec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/ExternalUDFExec.scala
@@ -63,9 +63,12 @@ trait ExternalUDFExec extends UnaryExecNode {
* Finalizes the session on task completion (which fires on both success and
* failure). [[WorkerSession#close]] is the single finalizer: it fetches the
* `FinishResponse` if processing completed, or cancels anything still in
- * flight and waits for the `CancelResponse`. The provided function receives
- * the session and must return the result iterator. It may use the session
- * but MUST NOT close it.
+ * flight and waits for the `CancelResponse`. For the UDF sessions used here,
+ * exhausting the data iterator completes execution and surfaces execution or
+ * finish errors. Therefore, `close` only cleans up the session and its returned
+ * termination does not determine whether the Spark task succeeds. The provided
+ * function receives the session and must return the result iterator. It may use
+ * the session but MUST NOT close it.
*/
protected def withUDFWorkerSession(
taskContext: TaskContext,
@@ -81,10 +84,9 @@ trait ExternalUDFExec extends UnaryExecNode {
// resolves to whichever terminator the stream reached:
//
// - Task completed and the result iterator was fully consumed: process()
- // sent Finish (input exhausted) and the data drained. close() then
- // typically returns the Finished termination -- but a failure during the
- // finish/cleanup phase (raised after the data drained, so it reached no one
- // through the iterator) is surfaced here too, as Failed / TransportFailed.
+ // sent Finish (input exhausted), drained the output, and surfaced any
+ // execution or finish error. close() observes the settled termination and
+ // finishes cleaning up the request side.
// - Task failed, was killed, or stopped before draining (e.g. a downstream
// LIMIT or exception): the stream has not finished, so close() sends a
// Cancel, the worker runs its cleanup, and its CancelResponse is returned
@@ -96,12 +98,11 @@ trait ExternalUDFExec extends UnaryExecNode {
// raising it (a thread interrupt may still propagate); the underlying
// failure has already surfaced through the result iterator.
//
- // The returned Termination (per-execution metrics, finish/cancel callback
- // result) is not consumed yet -- TODO [SPARK-57324] surface it once metrics
- // wiring lands.
- taskContext.addTaskCompletionListener[Unit] { _ =>
- session.close()
- }
+ // For these UDF sessions, exhausting the data iterator covers execution and
+ // surfaces its errors. close() only cleans up protocol state and releases or
+ // invalidates the worker handle, so its return value is intentionally ignored.
+ //
+ taskContext.addTaskCompletionListener[Unit](_ => session.close())
f(session)
}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/PlanExternalUDFs.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/PlanExternalUDFs.scala
index c60931b640dd..4118153647a3 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/PlanExternalUDFs.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/PlanExternalUDFs.scala
@@ -235,6 +235,11 @@ private[sql] object PlanExternalUDFs
collectEvaluableUDF(plan) match {
case None => plan
case Some(udf) =>
+ // TODO(SPARK-59745): Preserve named argument metadata in unified Python UDF execution.
+ if (udf.children.exists(_.isInstanceOf[NamedArgumentExpression])) {
+ throw QueryCompilationErrors.namedArgumentsNotSupported(
+ udf.name.getOrElse(udf.prettyName))
+ }
val childIndex = plan.children.indexWhere { child =>
udf.references.subsetOf(child.outputSet)
}
@@ -242,7 +247,7 @@ private[sql] object PlanExternalUDFs
throw QueryCompilationErrors.externalUDFWithMultipleChildrenUnsupportedError(udf)
}
- val resultAttr = AttributeReference("externalUDF", udf.dataType, udf.nullable)()
+ val resultAttr = AttributeReference("externalUDF", udf.dataType)()
val newChildren = plan.children.zipWithIndex.map { case (child, index) =>
if (index == childIndex) {
ExecuteExternalUDF(udf, resultAttr, child)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/PythonInitAdapter.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/PythonInitAdapter.scala
new file mode 100644
index 000000000000..ab76ac97c2a0
--- /dev/null
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/PythonInitAdapter.scala
@@ -0,0 +1,104 @@
+/*
+ * 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.externalUDF
+
+import scala.jdk.CollectionConverters._
+
+import com.google.protobuf.ByteString
+import org.json4s.{Formats, NoTypeHints}
+import org.json4s.jackson.Serialization
+
+import org.apache.spark.{BarrierTaskContext, TaskContext}
+import org.apache.spark.api.python.PythonEvalType
+import org.apache.spark.resource.CpuAmount
+import org.apache.spark.sql.catalyst.expressions.ExternalUserDefinedFunction
+import org.apache.spark.sql.execution.python.ArrowPythonRunner
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.udf.worker.{Init, UdfPayload, UDFWorkerDataFormat, UDFWorkerSpecification}
+
+/**
+ * Builds the temporary Python-specific Init message for scalar external UDF execution.
+ *
+ * The worker specification is deliberately an argument even though this adapter does not inspect
+ * it yet. The worker is not present, so the Init contract remains local to the Python adapter
+ * until SPARK-59364 replaces it with language-agnostic initialization.
+ */
+private[externalUDF] object PythonInitAdapter {
+ private val PYSPARK_UDF_PAYLOAD_FORMAT = "pyspark-udf-experimental"
+ private val INPUT_TYPE_CONF = "input_type"
+
+ private implicit val jsonFormats: Formats = Serialization.formats(NoTypeHints)
+
+ /** Query-invariant Init fields that can be serialized with the partition closure. */
+ final case class PreparedInit(template: Init) extends Serializable
+
+ def sessionConf(conf: SQLConf, inputSchema: StructType): Map[String, String] = {
+ (ArrowPythonRunner.getPythonRunnerConfMap(conf) - SQLConf.SESSION_LOCAL_TIMEZONE.key) +
+ (INPUT_TYPE_CONF -> inputSchema.json)
+ }
+
+ def prepare(
+ _workerSpec: UDFWorkerSpecification,
+ udf: ExternalUserDefinedFunction,
+ inputSchema: Array[Byte],
+ outputSchema: Array[Byte],
+ timeZoneId: String,
+ pythonSessionConf: Map[String, String]): PreparedInit = {
+ val udfBuilder = UdfPayload.newBuilder()
+ .setPayload(ByteString.copyFrom(udf.payload))
+ .setFormat(PYSPARK_UDF_PAYLOAD_FORMAT)
+ .setEvalType(PythonEvalType.SQL_ARROW_BATCHED_UDF.toString)
+ udf.name.foreach(udfBuilder.setName)
+
+ PreparedInit(Init.newBuilder()
+ .setProtocolVersion(1)
+ .setDataFormat(UDFWorkerDataFormat.ARROW)
+ .setUdf(udfBuilder)
+ .setInputSchema(ByteString.copyFrom(inputSchema))
+ .setOutputSchema(ByteString.copyFrom(outputSchema))
+ .setTimezone(timeZoneId)
+ .putAllSessionConf(pythonSessionConf.asJava)
+ .build())
+ }
+
+ def build(prepared: PreparedInit, context: TaskContext): Init = {
+ prepared.template.toBuilder
+ .putAllTaskContext(taskContextAsMap(context).asJava)
+ .build()
+ }
+
+ // TODO(SPARK-59364): Share this context mapping with PythonWorkerUtils when Init message
+ // construction becomes language-agnostic, while preserving the existing Python wire contract.
+ private def taskContextAsMap(context: TaskContext): Map[String, String] = {
+ val resources = context.resources().map { case (name, resource) =>
+ name -> Map("name" -> resource.name, "addresses" -> resource.addresses)
+ }
+ Map(
+ "isBarrier" -> context.isInstanceOf[BarrierTaskContext].toString,
+ "stageId" -> context.stageId().toString,
+ "partitionId" -> context.partitionId().toString,
+ "attemptNumber" -> context.attemptNumber().toString,
+ "taskAttemptId" -> context.taskAttemptId().toString,
+ "cpus" -> context.cpuAmount()
+ .setScale(0, BigDecimal.RoundingMode.CEILING).intValue.toString,
+ "cpuAmount" -> CpuAmount.toDisplayString(context.cpuAmount()),
+ "resources" -> Serialization.write(resources),
+ "localProperties" -> Serialization.write(context.getLocalProperties.asScala.toMap))
+ }
+}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/externalUDF/ExecuteExternalUDFExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/externalUDF/ExecuteExternalUDFExecSuite.scala
new file mode 100644
index 000000000000..17ebde5ffbbb
--- /dev/null
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/externalUDF/ExecuteExternalUDFExecSuite.scala
@@ -0,0 +1,631 @@
+/*
+ * 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.externalUDF
+
+import java.io.IOException
+import java.nio.channels.Channels
+import java.nio.charset.StandardCharsets
+import java.util.Properties
+
+import scala.jdk.CollectionConverters._
+
+import com.google.protobuf.ByteString
+import org.apache.arrow.vector.ipc.ReadChannel
+import org.apache.arrow.vector.ipc.message.MessageSerializer
+import org.apache.arrow.vector.types.pojo.Schema
+import org.json4s.{Formats, NoTypeHints}
+import org.json4s.jackson.Serialization
+import org.mockito.Mockito.mock
+
+import org.apache.spark.{BarrierTaskContext, SparkException, SparkUnsupportedOperationException,
+ TaskContext, TaskContextImpl}
+import org.apache.spark.api.python.PythonEvalType
+import org.apache.spark.resource.{CpuAmount, ResourceInformation}
+import org.apache.spark.sql.QueryTest
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeReference, Expression,
+ ExternalUserDefinedFunction, IsNull}
+import org.apache.spark.sql.execution.{ProjectExec, SparkPlan}
+import org.apache.spark.sql.execution.arrow.ArrowConverters
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.types.{DataType, LongType, StructField, StructType}
+import org.apache.spark.sql.util.ArrowUtils
+import org.apache.spark.udf.worker.{Cancel, DataRequest, DataResponse, Finish, FinishResponse, Init,
+ InitResponse, UDFWorkerDataFormat, UDFWorkerSpecification}
+import org.apache.spark.udf.worker.core.{Termination, WorkerHandle, WorkerLogger,
+ WorkerSecurityScope, WorkerSession}
+import org.apache.spark.util.{LongAccumulator, Utils}
+
+object ExecuteExternalUDFExecSuite {
+ private val TEST_PAYLOAD = "identity".getBytes(StandardCharsets.UTF_8)
+ private val RESPONSE_ALLOCATOR_NAME = "externalUdfArrowResponse"
+
+ private implicit val jsonFormats: Formats = Serialization.formats(NoTypeHints)
+
+ private final case class SerializedResource(name: String, addresses: Seq[String])
+
+ private sealed trait ResponseBehavior extends Serializable
+ private case object EchoResponses extends ResponseBehavior
+ private case object DropLastResponse extends ResponseBehavior
+ private case object DuplicateFirstResponse extends ResponseBehavior
+ private case object MalformedResponse extends ResponseBehavior
+ private case object CombineLongArguments extends ResponseBehavior
+ private case object ConstantLongResponses extends ResponseBehavior
+
+ private final case class TestExecution(
+ plan: SparkPlan,
+ requestCount: LongAccumulator,
+ closeCount: LongAccumulator)
+
+ private final class TestWorkerHandle extends WorkerHandle {
+ override def id: String = "test-worker"
+ override def markInvalid(): Unit = ()
+ override def releaseSession(): Unit = ()
+ }
+
+ private final class TestWorkerSession(
+ behavior: ResponseBehavior,
+ expectedInputSchema: StructType,
+ expectedOutputSchema: StructType,
+ expectedTimeZone: String,
+ expectedLargeVarTypes: Boolean,
+ requestCount: LongAccumulator,
+ closeCount: LongAccumulator)
+ extends WorkerSession(new TestWorkerHandle, WorkerLogger.NoOp) {
+
+ override protected def doInit(message: Init): InitResponse = {
+ require(message.getProtocolVersion == 1, "unexpected protocol version")
+ require(message.getDataFormat == UDFWorkerDataFormat.ARROW, "unexpected data format")
+ require(message.getUdf.getPayload.toByteArray.sameElements(TEST_PAYLOAD),
+ "unexpected UDF payload")
+ require(message.getUdf.getFormat == "pyspark-udf-experimental",
+ "unexpected UDF format")
+ require(message.getUdf.getEvalType == PythonEvalType.SQL_ARROW_BATCHED_UDF.toString,
+ "unexpected Python evaluation type")
+ require(message.getUdf.getName == "identity", "unexpected UDF name")
+ require(deserializeArrowSchema(message.getInputSchema) == ArrowUtils.toArrowSchema(
+ expectedInputSchema,
+ expectedTimeZone,
+ true,
+ expectedLargeVarTypes),
+ "unexpected input schema")
+ require(deserializeArrowSchema(message.getOutputSchema) == ArrowUtils.toArrowSchema(
+ expectedOutputSchema,
+ expectedTimeZone,
+ true,
+ expectedLargeVarTypes),
+ "unexpected output schema")
+ require(message.getTimezone == expectedTimeZone, "unexpected session timezone")
+ validateTaskContext(message, TaskContext.get())
+
+ val sessionConf = message.getSessionConfMap
+ require(sessionConf.get("input_type") == expectedInputSchema.json, "unexpected input type")
+ require(sessionConf.containsKey(SQLConf.PYSPARK_BINARY_AS_BYTES.key),
+ "Python binary representation setting is missing")
+ require(!sessionConf.containsKey(SQLConf.SESSION_LOCAL_TIMEZONE.key),
+ "timezone should use the typed Init field")
+ InitResponse.getDefaultInstance
+ }
+
+ override protected def doProcess(
+ input: Iterator[DataRequest],
+ finish: () => Finish): Iterator[DataResponse] = behavior match {
+ case EchoResponses => echoResponses(input, finish)
+ case DropLastResponse => eagerResponses(input, finish).dropRight(1).iterator
+ case DuplicateFirstResponse =>
+ val responses = eagerResponses(input, finish)
+ (responses ++ responses.take(1)).iterator
+ case MalformedResponse =>
+ eagerResponses(input, finish)
+ Iterator.single(DataResponse.newBuilder()
+ .setData(ByteString.copyFromUtf8("not an Arrow record batch"))
+ .build())
+ case CombineLongArguments => transformToLong(input, finish) { row =>
+ row.getLong(0) * 100L + row.getLong(1)
+ }
+ case ConstantLongResponses => transformToLong(input, finish)(_ => 42L)
+ }
+
+ override protected def doClose(cancel: () => Cancel): Termination = {
+ closeCount.add(1L)
+ completeTerminal(Termination.Finished(FinishResponse.getDefaultInstance))
+ settledTermination
+ }
+
+ private def echoResponses(
+ input: Iterator[DataRequest],
+ finish: () => Finish): Iterator[DataResponse] = new Iterator[DataResponse] {
+ private var finished = false
+
+ override def hasNext: Boolean = {
+ val hasNextRequest = input.hasNext
+ if (!hasNextRequest && !finished) {
+ finish()
+ finished = true
+ }
+ hasNextRequest
+ }
+
+ override def next(): DataResponse = {
+ if (!hasNext) Iterator.empty.next()
+ requestCount.add(1L)
+ response(input.next())
+ }
+ }
+
+ private def eagerResponses(
+ input: Iterator[DataRequest],
+ finish: () => Finish): Vector[DataResponse] = {
+ val requests = input.toVector
+ requestCount.add(requests.length.toLong)
+ finish()
+ requests.map(response)
+ }
+
+ private def response(request: DataRequest): DataResponse = {
+ DataResponse.newBuilder().setData(request.getData).build()
+ }
+
+ private def transformToLong(
+ input: Iterator[DataRequest],
+ finish: () => Finish)(
+ evaluate: InternalRow => Long): Iterator[DataResponse] = {
+ val requests = input.toVector
+ requestCount.add(requests.length.toLong)
+ finish()
+ val context = TaskContext.get()
+ requests.iterator.flatMap { request =>
+ val inputRows = ArrowConverters.fromBatchIterator(
+ Iterator.single(request.getData.toByteArray),
+ expectedInputSchema,
+ expectedTimeZone,
+ true,
+ expectedLargeVarTypes,
+ context)
+ val outputRows = inputRows.map { row =>
+ InternalRow(evaluate(row))
+ }
+ ArrowConverters.toBatchIterator(
+ outputRows,
+ expectedOutputSchema,
+ 0,
+ expectedTimeZone,
+ true,
+ expectedLargeVarTypes,
+ context).map { batch =>
+ DataResponse.newBuilder().setData(ByteString.copyFrom(batch)).build()
+ }
+ }
+ }
+
+ private def deserializeArrowSchema(schema: ByteString): Schema = {
+ MessageSerializer.deserializeSchema(
+ new ReadChannel(Channels.newChannel(schema.newInput())))
+ }
+ }
+
+ private def validateTaskContext(message: Init, context: TaskContext): Unit = {
+ val taskContext = message.getTaskContextMap.asScala.toMap
+ val expectedValues = Map(
+ "isBarrier" -> context.isInstanceOf[BarrierTaskContext].toString,
+ "stageId" -> context.stageId().toString,
+ "partitionId" -> context.partitionId().toString,
+ "attemptNumber" -> context.attemptNumber().toString,
+ "taskAttemptId" -> context.taskAttemptId().toString,
+ "cpus" -> context.cpuAmount()
+ .setScale(0, BigDecimal.RoundingMode.CEILING).intValue.toString,
+ "cpuAmount" -> CpuAmount.toDisplayString(context.cpuAmount()))
+ val expectedKeys = expectedValues.keySet ++ Set("resources", "localProperties")
+ require(taskContext.keySet == expectedKeys,
+ s"unexpected task context keys: ${taskContext.keySet}")
+ expectedValues.foreach { case (key, value) =>
+ require(taskContext(key) == value,
+ s"unexpected task context value for $key: ${taskContext(key)}")
+ }
+
+ val expectedResources = context.resources().map { case (name, resource) =>
+ name -> SerializedResource(resource.name, resource.addresses.toSeq)
+ }
+ require(
+ Serialization.read[Map[String, SerializedResource]](taskContext("resources")) ==
+ expectedResources,
+ "unexpected task resources")
+ require(
+ Serialization.read[Map[String, String]](taskContext("localProperties")) ==
+ context.getLocalProperties.asScala.toMap,
+ "unexpected task local properties")
+ }
+
+ private final class TestExecuteExternalUDFExec(
+ udf: ExternalUserDefinedFunction,
+ resultAttr: Attribute,
+ child: SparkPlan,
+ behavior: ResponseBehavior,
+ expectedInputSchema: StructType,
+ expectedOutputSchema: StructType,
+ expectedTimeZone: String,
+ expectedLargeVarTypes: Boolean,
+ requestCount: LongAccumulator,
+ closeCount: LongAccumulator)
+ extends ExecuteExternalUDFExec(udf, resultAttr, child) {
+
+ override protected def withUDFWorkerSession(
+ taskContext: TaskContext,
+ securityScope: Option[WorkerSecurityScope])(
+ f: WorkerSession => Iterator[InternalRow]): Iterator[InternalRow] = {
+ require(securityScope.isEmpty, "scalar external UDF execution must not request a sandbox")
+ val session = new TestWorkerSession(
+ behavior,
+ expectedInputSchema,
+ expectedOutputSchema,
+ expectedTimeZone,
+ expectedLargeVarTypes,
+ requestCount,
+ closeCount)
+ taskContext.addTaskCompletionListener[Unit](_ => session.close())
+ f(session)
+ }
+ }
+}
+
+class ExecuteExternalUDFExecSuite extends QueryTest with SharedSparkSession {
+ import ExecuteExternalUDFExecSuite._
+
+ private def testExecution(
+ behavior: ResponseBehavior,
+ rowCount: Long): TestExecution = {
+ val child = spark.range(0L, rowCount, 1L, 1).queryExecution.executedPlan
+ testExecution(
+ behavior,
+ child,
+ Seq(child.output.head),
+ LongType,
+ udfNullable = false)
+ }
+
+ private def testExecution(
+ behavior: ResponseBehavior,
+ child: SparkPlan,
+ udfChildren: Seq[Expression],
+ udfDataType: DataType,
+ udfNullable: Boolean): TestExecution = {
+ val udf = ExternalUserDefinedFunction(
+ name = Some("identity"),
+ workerSpec = UDFWorkerSpecification.getDefaultInstance,
+ payload = TEST_PAYLOAD,
+ dataType = udfDataType,
+ children = udfChildren,
+ udfDeterministic = true,
+ udfNullable = udfNullable)
+ val resultAttr = AttributeReference("externalUDF", udfDataType, udf.nullable)()
+ val requestCount = spark.sparkContext.longAccumulator("external UDF request count")
+ val closeCount = new LongAccumulator
+ closeCount.register(
+ spark.sparkContext,
+ name = Some("external UDF close count"),
+ countFailedValues = true)
+ val inputSchema = StructType(udfChildren.zipWithIndex.map {
+ case (expression, index) =>
+ StructField(s"_$index", expression.dataType, expression.nullable)
+ })
+ val outputSchema = StructType(Seq(StructField("_0", udfDataType, nullable = true)))
+ val conf = spark.sessionState.conf
+ TestExecution(
+ new TestExecuteExternalUDFExec(
+ udf,
+ resultAttr,
+ child,
+ behavior,
+ inputSchema,
+ outputSchema,
+ conf.sessionLocalTimeZone,
+ conf.arrowUseLargeVarTypes,
+ requestCount,
+ closeCount),
+ requestCount,
+ closeCount)
+ }
+
+ private def checkResultRowsMismatch(
+ error: SparkException,
+ outputLength: String,
+ inputLength: Long): Unit = {
+ checkError(
+ exception = Utils.getRootCause(error).asInstanceOf[SparkException],
+ condition = "RESULT_ROWS_MISMATCH",
+ parameters = Map(
+ "output_length" -> outputLength,
+ "input_length" -> inputLength.toString))
+ }
+
+ private def checkFailedExecutionCleanup(execution: TestExecution): Unit = {
+ assert(execution.closeCount.value === 1L)
+ val openResponseAllocators = ArrowUtils.rootAllocator.getChildAllocators.asScala
+ .filter(_.getName == RESPONSE_ALLOCATOR_NAME)
+ .toSeq
+ assert(openResponseAllocators.isEmpty,
+ s"Arrow response allocators remain open: $openResponseAllocators")
+ }
+
+ test("Python Init adapter serializes the complete task context") {
+ val localProperties = new Properties
+ localProperties.setProperty("externalUdfTestProperty", "test-value")
+ val context = new TaskContextImpl(
+ stageId = 7,
+ stageAttemptNumber = 3,
+ partitionId = 5,
+ taskAttemptId = 1234L,
+ attemptNumber = 2,
+ numPartitions = 10,
+ taskMemoryManager = null,
+ localProperties = localProperties,
+ metricsSystem = null,
+ cpuAmount = BigDecimal("1.5"),
+ resources = Map(
+ "gpu" -> new ResourceInformation("gpu", Array("0", "2"))))
+
+ val init = PythonInitAdapter.build(
+ PythonInitAdapter.PreparedInit(Init.getDefaultInstance),
+ context)
+ validateTaskContext(init, context)
+ }
+
+ test("scalar external UDF rejects barrier task contexts before worker creation") {
+ val error = intercept[SparkUnsupportedOperationException] {
+ ExecuteExternalUDFExec.ensureTaskContextSupported(mock(classOf[BarrierTaskContext]))
+ }
+ checkError(
+ exception = error,
+ condition = "UNSUPPORTED_FEATURE.EXTERNAL_UDF_IN_BARRIER_TASK",
+ parameters = Map.empty)
+ }
+
+ test("scalar external UDF exchanges multiple Arrow batches through a worker session") {
+ withSQLConf(
+ SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC",
+ SQLConf.ARROW_EXECUTION_MAX_RECORDS_PER_BATCH.key -> "2") {
+ val execution = testExecution(EchoResponses, rowCount = 5L)
+ val rows = execution.plan.executeCollect()
+
+ assert(rows.map(row => (row.getLong(0), row.getLong(1))).toSeq === Seq(
+ (0L, 0L),
+ (1L, 1L),
+ (2L, 2L),
+ (3L, 3L),
+ (4L, 4L)))
+ assert(execution.requestCount.value === 3L)
+ assert(execution.closeCount.value === 1L)
+ }
+ }
+
+ test("scalar external UDF finishes an empty input partition") {
+ val child = spark.range(0L, 0L, 1L, 1).coalesce(1).queryExecution.executedPlan
+ val execution = testExecution(
+ EchoResponses,
+ child,
+ Seq(child.output.head),
+ LongType,
+ udfNullable = false)
+
+ assert(execution.plan.executeCollect().isEmpty)
+ assert(execution.requestCount.value === 0L)
+ assert(execution.closeCount.value === 1L)
+ }
+
+ test("scalar external UDF preserves row count for zero-argument functions") {
+ withSQLConf(SQLConf.ARROW_EXECUTION_MAX_RECORDS_PER_BATCH.key -> "2") {
+ val child = spark.range(0L, 5L, 1L, 1).queryExecution.executedPlan
+ val execution = testExecution(
+ ConstantLongResponses,
+ child,
+ Seq.empty,
+ LongType,
+ udfNullable = false)
+
+ val rows = execution.plan.executeCollect()
+ assert(rows.map(row => (row.getLong(0), row.getLong(1))).toSeq ===
+ (0L until 5L).map(value => (value, 42L)))
+ assert(execution.requestCount.value === 3L)
+ assert(execution.closeCount.value === 1L)
+ }
+ }
+
+ test("scalar external UDF round-trips ordered multi-argument Arrow data") {
+ val child = spark.range(0L, 3L, 1L, 1)
+ .selectExpr("id AS left", "id + 10 AS right")
+ .queryExecution.executedPlan
+ val execution = testExecution(
+ CombineLongArguments,
+ child,
+ child.output,
+ LongType,
+ udfNullable = false)
+
+ val rows = execution.plan.executeCollect()
+ assert(rows.map(row => (row.getLong(0), row.getLong(1), row.getLong(2))).toSeq === Seq(
+ (0L, 10L, 10L),
+ (1L, 11L, 111L),
+ (2L, 12L, 212L)))
+ assert(execution.requestCount.value === 1L)
+ assert(execution.closeCount.value === 1L)
+ }
+
+ test("scalar external UDF closes sessions for multiple partitions") {
+ val child = spark.range(0L, 6L, 1L, 3).queryExecution.executedPlan
+ val execution = testExecution(
+ EchoResponses,
+ child,
+ Seq(child.output.head),
+ LongType,
+ udfNullable = false)
+
+ val rows = execution.plan.executeCollect()
+
+ assert(rows.map(row => (row.getLong(0), row.getLong(1))).toSeq ===
+ (0L until 6L).map(value => (value, value)))
+ assert(execution.requestCount.value === 3L)
+ assert(execution.closeCount.value === 3L)
+ }
+
+ test("scalar external UDF exposes nullable output to downstream consumers") {
+ val child = spark.range(0L, 2L, 1L, 1)
+ .selectExpr("IF(id = 0, CAST(NULL AS BIGINT), id) AS value")
+ .queryExecution.executedPlan
+ val input = child.output.head
+ val execution = testExecution(
+ EchoResponses,
+ child,
+ Seq(input),
+ LongType,
+ udfNullable = false)
+
+ val resultAttr = execution.plan.output.last
+ assert(resultAttr.nullable)
+ val downstream = ProjectExec(
+ Seq(resultAttr, Alias(IsNull(resultAttr), "resultIsNull")()),
+ execution.plan)
+ val rows = downstream.executeCollect()
+ assert(rows.length === 2)
+ assert(rows.head.isNullAt(0))
+ assert(rows.head.getBoolean(1))
+ assert(rows(1).getLong(0) === 1L)
+ assert(!rows(1).getBoolean(1))
+ assert(execution.requestCount.value === 1L)
+ assert(execution.closeCount.value === 1L)
+ }
+
+ test("scalar external UDF limits Arrow input batches by byte size") {
+ withSQLConf(
+ SQLConf.ARROW_EXECUTION_MAX_RECORDS_PER_BATCH.key -> "100",
+ SQLConf.ARROW_EXECUTION_MAX_BYTES_PER_BATCH.key -> "512") {
+ val child = spark.range(0L, 3L, 1L, 1)
+ .selectExpr("repeat('x', 1024) AS value")
+ .queryExecution.executedPlan
+ val input = child.output.head
+ val execution = testExecution(
+ EchoResponses,
+ child,
+ Seq(input),
+ input.dataType,
+ input.nullable)
+ val rows = execution.plan.executeCollect()
+
+ assert(rows.length === 3)
+ rows.foreach { row =>
+ assert(row.getUTF8String(0) === row.getUTF8String(1))
+ }
+ assert(execution.requestCount.value === 3L)
+ assert(execution.closeCount.value === 1L)
+ }
+ }
+
+ test("scalar external UDF uses large variable Arrow types for requests and responses") {
+ withSQLConf(SQLConf.ARROW_EXECUTION_USE_LARGE_VAR_TYPES.key -> "true") {
+ val child = spark.range(0L, 3L, 1L, 1)
+ .selectExpr("concat('value-', id) AS value")
+ .queryExecution.executedPlan
+ val input = child.output.head
+ val execution = testExecution(
+ EchoResponses,
+ child,
+ Seq(input),
+ input.dataType,
+ input.nullable)
+ val rows = execution.plan.executeCollect()
+
+ assert(rows.length === 3)
+ rows.foreach { row =>
+ assert(row.getUTF8String(0) === row.getUTF8String(1))
+ }
+ assert(execution.requestCount.value === 1L)
+ assert(execution.closeCount.value === 1L)
+ }
+ }
+
+ test("scalar external UDF decodes compressed Arrow responses") {
+ withSQLConf(
+ SQLConf.ARROW_EXECUTION_COMPRESSION_CODEC.key -> "lz4",
+ SQLConf.ARROW_EXECUTION_MAX_RECORDS_PER_BATCH.key -> "2") {
+ val child = spark.range(0L, 5L, 1L, 1)
+ .selectExpr("repeat(CAST(id AS STRING), 1024) AS value")
+ .queryExecution.executedPlan
+ val input = child.output.head
+ val execution = testExecution(
+ EchoResponses,
+ child,
+ Seq(input),
+ input.dataType,
+ input.nullable)
+ val rows = execution.plan.executeCollect()
+
+ assert(rows.length === 5)
+ rows.foreach { row =>
+ assert(row.getUTF8String(0) === row.getUTF8String(1))
+ }
+ assert(execution.requestCount.value === 3L)
+ assert(execution.closeCount.value === 1L)
+ }
+ }
+
+ test("scalar external UDF rejects fewer output rows than input rows") {
+ withSQLConf(SQLConf.ARROW_EXECUTION_MAX_RECORDS_PER_BATCH.key -> "1") {
+ val execution = testExecution(DropLastResponse, rowCount = 3L)
+ val error = intercept[SparkException] {
+ execution.plan.executeCollect()
+ }
+
+ checkResultRowsMismatch(error, outputLength = "2", inputLength = 3L)
+ checkFailedExecutionCleanup(execution)
+ }
+ }
+
+ test("scalar external UDF reports buffered extra output rows as a lower bound") {
+ withSQLConf(SQLConf.ARROW_EXECUTION_MAX_RECORDS_PER_BATCH.key -> "3") {
+ val execution = testExecution(DuplicateFirstResponse, rowCount = 3L)
+ val error = intercept[SparkException] {
+ execution.plan.executeCollect()
+ }
+
+ checkResultRowsMismatch(error, outputLength = "at least 6", inputLength = 3L)
+ checkFailedExecutionCleanup(execution)
+ }
+ }
+
+ test("scalar external UDF rejects a malformed Arrow response") {
+ val execution = testExecution(MalformedResponse, rowCount = 3L)
+ val error = intercept[SparkException] {
+ execution.plan.executeCollect()
+ }
+
+ assert(Utils.getRootCause(error).isInstanceOf[IOException])
+ checkFailedExecutionCleanup(execution)
+ }
+
+ test("scalar external UDF closes its worker session when output consumption stops early") {
+ withSQLConf(SQLConf.ARROW_EXECUTION_MAX_RECORDS_PER_BATCH.key -> "2") {
+ val execution = testExecution(EchoResponses, rowCount = 5L)
+ val rows = execution.plan.executeTake(1)
+
+ assert(rows.length === 1)
+ assert(rows.head.getLong(0) === rows.head.getLong(1))
+ assert(execution.requestCount.value === 1L)
+ assert(execution.closeCount.value === 1L)
+ }
+ }
+}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/externalUDF/PlanExternalUDFsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/externalUDF/PlanExternalUDFsSuite.scala
index 823523f0a312..5161128f1103 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/externalUDF/PlanExternalUDFsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/externalUDF/PlanExternalUDFsSuite.scala
@@ -19,13 +19,14 @@ package org.apache.spark.sql.execution.externalUDF
import java.nio.charset.StandardCharsets
-import org.apache.spark.{SparkConf, SparkUnsupportedOperationException}
+import org.apache.spark.{SparkConf, SparkException, SparkUnsupportedOperationException}
import org.apache.spark.sql.{AnalysisException, QueryTest}
import org.apache.spark.sql.catalyst.QueryPlanningTracker
import org.apache.spark.sql.catalyst.expressions.{Add, Alias, And, ArrayTransform, Attribute,
AttributeReference, AttributeSet, CreateArray, EqualTo, Expression,
- ExternalUserDefinedFunction, GreaterThan, Lag, LambdaFunction, Literal, NamedLambdaVariable,
- UnspecifiedFrame, WindowExpression, WindowSpecDefinition}
+ ExternalUserDefinedFunction, GreaterThan, IsNull, Lag, LambdaFunction, Literal,
+ NamedArgumentExpression, NamedLambdaVariable, UnspecifiedFrame, WindowExpression,
+ WindowSpecDefinition}
import org.apache.spark.sql.catalyst.expressions.aggregate.Sum
import org.apache.spark.sql.catalyst.plans.{Inner, LeftOuter}
import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, BinaryNode, ExecuteExternalUDF,
@@ -36,6 +37,7 @@ import org.apache.spark.sql.test.SharedSparkSession
import org.apache.spark.sql.types.{BooleanType, DataType, IntegerType, StructField, StructType}
import org.apache.spark.udf.worker.{DirectWorker, ProcessCallable, UDFWorkerProperties,
UDFWorkerSpecification, WorkerEnvironment}
+import org.apache.spark.util.Utils
class PlanExternalUDFsSuite extends QueryTest with SharedSparkSession {
@@ -219,6 +221,49 @@ class PlanExternalUDFsSuite extends QueryTest with SharedSparkSession {
}
}
+ test("external UDF results remain nullable through extraction and downstream expressions") {
+ val spec = workerSpec("nullable-result")
+ val range = Range(0, 10, 1, 1)
+ val function = udf("nullable-result", spec, Seq(singleOutput(range)))
+ assert(!function.udfNullable)
+ assert(function.nullable)
+ val functionProject = Project(Seq(Alias(function, "functionResult")()), range)
+ val originalResult = singleOutput(functionProject)
+ assert(originalResult.nullable)
+ val plan = Project(Seq(Alias(IsNull(originalResult), "resultIsNull")()), functionProject)
+
+ val extracted = extract(plan)
+ val evaluation = singleEvalNode(extracted)
+ assert(evaluation.resultAttr.nullable)
+ val downstreamPredicate = singleProjectExpression(extracted) match {
+ case predicate: IsNull => predicate
+ case other => fail(s"Expected an IsNull expression, found: $other")
+ }
+ val downstreamResult = downstreamPredicate.child match {
+ case attribute: Attribute => attribute
+ case other => fail(s"Expected an attribute under IsNull, found: $other")
+ }
+ assert(downstreamResult.nullable)
+
+ val optimized = optimize(extracted)
+ assert(optimized.exists(_.expressions.exists(_.exists(_.isInstanceOf[IsNull]))))
+ }
+
+ test("named arguments are rejected during logical planning") {
+ val function = udf(
+ "named-arguments",
+ workerSpec("named-arguments"),
+ Seq(NamedArgumentExpression("value", input)))
+
+ val error = intercept[AnalysisException] {
+ extract(function)
+ }
+ checkError(
+ exception = error,
+ condition = "NAMED_PARAMETERS_NOT_SUPPORTED",
+ parameters = Map("functionName" -> "`named-arguments`"))
+ }
+
test("UDFs with different worker specifications use separate nodes") {
val outerSpec = workerSpec("outer-worker")
val innerSpec = workerSpec("inner-worker")
@@ -236,7 +281,7 @@ class PlanExternalUDFsSuite extends QueryTest with SharedSparkSession {
val spec = workerSpec("safe-rendering", Map("UDF_SECRET" -> sensitiveWorkerValue))
val function = udf("safeRendering", spec, Seq(input)).copy(
payload = sensitivePayloadValue.getBytes(StandardCharsets.UTF_8))
- val resultAttr = AttributeReference("externalUDF", IntegerType, nullable = false)()
+ val resultAttr = AttributeReference("externalUDF", IntegerType, nullable = true)()
val logicalPlan = ExecuteExternalUDF(function, resultAttr, relation)
// Structured logging renders MDC values with toString, including join conditions.
val loggedCondition = EqualTo(function, Literal(1)).toString
@@ -379,6 +424,7 @@ class PlanExternalUDFsSuite extends QueryTest with SharedSparkSession {
val evaluation = singleEvalNode(extracted)
assert(evaluation.child.isInstanceOf[Aggregate])
assert(extracted.output == plan.output)
+ assert(extracted.output.head.nullable)
}
test("an external UDF used as a grouping key is evaluated before aggregate") {
@@ -486,6 +532,7 @@ class PlanExternalUDFsSuite extends QueryTest with SharedSparkSession {
assert(extracted.output.map(_.name) == Seq(input.name, "result"))
assert(extracted.output.head.semanticEquals(input))
val evaluation = singleEvalNode(extracted)
+ assert(evaluation.resultAttr.nullable)
val window = evaluation.child match {
case node: Window => node
case other => fail(s"Expected a Window node, found:\n$other")
@@ -644,7 +691,7 @@ class PlanExternalUDFsSuite extends QueryTest with SharedSparkSession {
val spec = workerSpec("limit")
val range = Range(0, 10, 1, 1)
val function = udf("limit", spec, Seq(singleOutput(range)))
- val resultAttr = AttributeReference("externalUDF", IntegerType, nullable = false)()
+ val resultAttr = AttributeReference("externalUDF", IntegerType, nullable = true)()
val plan = LocalLimit(Literal(1), ExecuteExternalUDF(function, resultAttr, range))
val optimized = optimize(plan)
@@ -656,7 +703,7 @@ class PlanExternalUDFsSuite extends QueryTest with SharedSparkSession {
val spec = workerSpec("physical-planning")
val range = Range(0, 10, 1, 1)
val function = udf("physical-planning", spec, Seq(singleOutput(range)))
- val resultAttr = AttributeReference("externalUDF", IntegerType, nullable = false)()
+ val resultAttr = AttributeReference("externalUDF", IntegerType, nullable = true)()
val logicalPlan = ExecuteExternalUDF(function, resultAttr, range)
val execution = spark.sessionState.planner.plan(logicalPlan).toSeq match {
@@ -669,11 +716,11 @@ class PlanExternalUDFsSuite extends QueryTest with SharedSparkSession {
assert(execution.workerSpec == spec)
}
- test("scalar external UDF execution reports unimplemented before worker startup") {
+ test("scalar external UDF execution reaches dispatcher creation") {
val spec = workerSpec("unimplemented-scalar")
val range = Range(0, 10, 1, 1)
val function = udf("unimplemented-scalar", spec, Seq(singleOutput(range)))
- val resultAttr = AttributeReference("externalUDF", IntegerType, nullable = false)()
+ val resultAttr = AttributeReference("externalUDF", IntegerType, nullable = true)()
val logicalPlan = ExecuteExternalUDF(function, resultAttr, range)
val execution = spark.sessionState.planner.plan(logicalPlan).toSeq match {
case Seq(node: ExecuteExternalUDFExec) => node
@@ -681,13 +728,11 @@ class PlanExternalUDFsSuite extends QueryTest with SharedSparkSession {
fail(s"Expected one ExecuteExternalUDFExec node, found:\n${other.mkString("\n")}")
}
- val exception = intercept[SparkUnsupportedOperationException] {
- execution.execute()
+ val exception = intercept[SparkException] {
+ execution.execute().collect()
}
- checkError(
- exception = exception,
- condition = "_LEGACY_ERROR_TEMP_2041",
- parameters = Map("methodName" -> "ExecuteExternalUDFExec.doExecute"))
+ assert(Utils.exceptionString(exception).contains(
+ "No UDF dispatcher factory configured. Set up a concrete factory for SPARK-55278."))
}
test("map partitions external UDF execution reports unimplemented before worker startup") {
@@ -724,7 +769,7 @@ class PlanExternalUDFsSuite extends QueryTest with SharedSparkSession {
val range = Range(0, 10, 1, 1)
val inputAttribute = singleOutput(range)
val function = udf("predicate", spec, Seq(inputAttribute))
- val resultAttr = AttributeReference("externalUDF", IntegerType, nullable = false)()
+ val resultAttr = AttributeReference("externalUDF", IntegerType, nullable = true)()
val childPredicate = GreaterThan(inputAttribute, Literal(0L))
val resultPredicate = EqualTo(resultAttr, Literal(1))
val plan = Filter(
@@ -740,8 +785,8 @@ class PlanExternalUDFsSuite extends QueryTest with SharedSparkSession {
case Filter(condition, _) => condition
case other => fail(s"Expected Filter below ExecuteExternalUDF, found:\n$other")
}
- assert(retainedPredicate.semanticEquals(resultPredicate))
- assert(pushedPredicate.semanticEquals(childPredicate))
+ assert(retainedPredicate.exists(_.semanticEquals(resultPredicate)))
+ assert(pushedPredicate.exists(_.semanticEquals(childPredicate)))
}
test("optimizer checks cartesian products after external UDF extraction") {
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org