cloud-fan commented on code in PR #58988:
URL: https://github.com/apache/spark/pull/58988#discussion_r4083853255
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/ExecuteExternalUDFExec.scala:
##########
@@ -47,16 +70,253 @@ case class ExecuteExternalUDFExec(
override def workerSpec: UDFWorkerSpecification = udf.workerSpec
+ // External worker outputs are always nullable.
+ assert(resultAttr.nullable, "The external UDF result attribute must be
nullable")
+
override def output: Seq[Attribute] = child.output :+ resultAttr
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")
+ // 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 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
Review Comment:
**Non-blocking (P2):** Please add a non-default-mode case that keeps the raw
Arrow type visible on both request and response paths. The current schema
checks round-trip through Spark StructType, so they erase the Utf8/LargeUtf8
distinction and would stay green if this flag were ignored or applied in only
one direction.
##########
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-v1"
+ 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,
Review Comment:
**Non-blocking (P2):** Please validate the values and encodings of the
complete task-context map, not just the presence of four keys. As written,
swapped stage/partition IDs, missing barrier/attempt/CPU fields, or malformed
`resources`/`localProperties` JSON would all leave the execution suite green.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/ExecuteExternalUDFExec.scala:
##########
@@ -47,16 +70,253 @@ case class ExecuteExternalUDFExec(
override def workerSpec: UDFWorkerSpecification = udf.workerSpec
+ // External worker outputs are always nullable.
+ assert(resultAttr.nullable, "The external UDF result attribute must be
nullable")
+
override def output: Seq[Attribute] = child.output :+ resultAttr
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")
+ // 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 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.toInt
+ val preparedInit = prepareInit(inputSchema, outputSchema, timeZoneId,
largeVarTypes)
+
+ child.execute().mapPartitionsInternal { rows =>
+ val context = TaskContext.get()
+ 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 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 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) = initialize()
+ while (!rows.hasNext && !exhausted) {
+ if (responses.hasNext) {
+ val batch = ArrowConverters.loadBatch(
+ responses.next().getData.newInput(),
+ activeAllocator)
+ try {
+ new VectorLoader(activeRoot).load(batch)
+ } finally {
+ batch.close()
+ }
+ rows = rowsFromRoot(activeRoot)
+ rowsRemainingInBatch = activeRoot.getRowCount
+ } else {
+ exhausted = true
+ close()
+ }
+ }
+ rows.hasNext
+ } catch {
+ case NonFatal(error) =>
Review Comment:
**Non-blocking (P2):** The malformed-response and both cardinality-mismatch
cases verify only the surfaced exception. Please make failed-task cleanup
observable and assert that the worker session and Arrow response resources
close in each case; the existing success and early-stop close assertions would
not catch a leak that occurs only after an exception.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowConverters.scala:
##########
@@ -231,6 +277,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
+ * checked before appending each row, and the first row is always accepted.
Review Comment:
**Nit (P3):** This sentence is missing `is`: it should read “because its
size is checked before appending each row.”
--
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]
