[SPARK-59744][UDF] Execute scalar external UDFs through worker sessions

common/utils/src/main/resources/error/error-conditions.json

@@ -7088,6 +7088,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.",
@@ -9069,6 +9075,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."

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 {

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,

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))