cloud-fan commented on code in PR #58988:
URL: https://github.com/apache/spark/pull/58988#discussion_r4085215998
##########
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:
**Non-blocking (P2):** Could we add an execution test for a zero-argument
scalar UDF over nonempty input? This path creates zero-column Arrow batches,
whose row count must survive without any vectors; the existing planner test
proves the shape is supported, but every worker-backed execution test currently
projects at least one argument, so a cardinality regression here would go
undetected.
##########
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:
**Non-blocking (P2):** This rebuilds the `ArrowColumnVector` accessor tree
and `ColumnarBatch` for every response batch even though `VectorLoader` reloads
the same persistent root in place. For small batches or nested return types,
that adds schema-proportional allocation on the hot path. Could we construct
these wrappers once when the root is initialized and only update the batch row
count after each load?
--
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]
