haiyangsun-db commented on code in PR #58988:
URL: https://github.com/apache/spark/pull/58988#discussion_r4102723598
##########
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(
Review Comment:
Added an execution test for a zero-argument scalar UDF over five input rows.
It verifies the output cardinality, returned constant values, request-batch
count, and session cleanup.
##########
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:
Added the execution test "scalar external UDF uses large variable Arrow
types for requests and responses". It enables
ARROW_EXECUTION_USE_LARGE_VAR_TYPES, validates the raw Init schemas with large
variable types, echoes the request batch through the response path, and
verifies the result.
##########
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 = "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,
Review Comment:
Added ensureTaskContextSupported before withUDFWorkerSession, so
BarrierTaskContext is rejected with a structured error before any worker
session is created. A regression test covers the rejection.
##########
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) =>
+ close()
+ throw error
+ }
+ }
+
+ override def next(): InternalRow = {
+ if (!hasNext) Iterator.empty.next()
+ val row = rows.next()
+ rowsRemainingInBatch -= 1
+ row
+ }
+
+ private def initialize(): (BufferAllocator, VectorSchemaRoot) = {
+ if (!initialized) {
+ initialized = true
+ val newAllocator = ArrowUtils.rootAllocator.newChildAllocator(
+ "externalUdfArrowResponse",
+ 0,
+ Long.MaxValue)
+ allocator = Some(newAllocator)
+ root = Some(VectorSchemaRoot.create(
+ ArrowUtils.toArrowSchema(
+ expectedSchema,
+ timeZoneId,
+ true,
+ largeVarTypes),
+ newAllocator))
+ }
+ (allocator, root) match {
+ case (Some(activeAllocator), Some(activeRoot)) =>
+ (activeAllocator, activeRoot)
+ case _ =>
+ throw SparkException.internalError(
+ "The Arrow response row iterator was not initialized correctly.")
+ }
+ }
+
+ private def rowsFromRoot(root: VectorSchemaRoot): Iterator[InternalRow] = {
+ val columns: Array[ColumnVector] =
root.getFieldVectors.asScala.iterator.map { vector =>
Review Comment:
Updated ArrowResponseRowIterator to create the VectorSchemaRoot,
ArrowColumnVector wrappers, and ColumnarBatch once during initialization. Each
response now only reloads the persistent root and updates the row count.
##########
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:
Fixed the Scaladoc wording to say that the size is checked before appending
each row.
##########
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:
Added checkFailedExecutionCleanup, which verifies both worker-session
closure and that no Arrow response allocator remains open. The
malformed-response and both cardinality-mismatch tests now call it.
##########
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:
Expanded validateTaskContext to check the complete key set and exact values
for barrier, stage, partition, attempt, task, and CPU fields, and to
deserialize and compare the resources and localProperties JSON payloads.
--
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]
