[SPARK-55278][SQL] Add scalar external UDF planning support

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

@@ -8592,6 +8592,11 @@
           "Drop the namespace <namespace>."
         ]
       },
+      "EXTERNAL_UDF" : {
+        "message" : [
+          "External UDF expressions. Set <config> to true to enable them."
+        ]
+      },
       "GEOSPATIAL_DISABLED" : {
         "message" : [
           "Geospatial feature is disabled."
@@ -8742,6 +8747,11 @@
           "Python 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."
         ]
       },
+      "PYTHON_UDF_TO_EXTERNAL_UDF" : {
+        "message" : [
+          "Legacy scalar PythonUDF expressions in the unified execution path. Set <config> to true to convert them to external UDF expressions."
+        ]
+      },
       "QUERY_ONLY_CORRUPT_RECORD_COLUMN" : {
         "message" : [
           "Queries from raw JSON/CSV/XML files are disallowed when the",

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ExternalUserDefinedFunction.scala

@@ -20,6 +20,7 @@ package org.apache.spark.sql.catalyst.expressions
 import org.apache.spark.annotation.Experimental
 import org.apache.spark.sql.catalyst.trees.TreePattern.{EXTERNAL_UDF, TreePattern}
 import org.apache.spark.sql.types.DataType
+import org.apache.spark.udf.worker.UDFWorkerSpecification
 
 /**
  * :: Experimental ::
@@ -37,6 +38,7 @@ import org.apache.spark.sql.types.DataType
  * to execute.
  *
  * @param name             Optional name of the UDF.
+ * @param workerSpec       Specification of the worker that executes this UDF.
  * @param payload          Opaque serialized function definition.
  * @param dataType         Return type of the UDF.
  * @param children         Input argument expressions.
@@ -48,6 +50,7 @@ import org.apache.spark.sql.types.DataType
 @Experimental
 case class ExternalUserDefinedFunction(
     name: Option[String],
+    workerSpec: UDFWorkerSpecification,
     payload: Array[Byte],
     dataType: DataType,
     children: Seq[Expression],

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala

@@ -1004,6 +1004,10 @@ object LimitPushDown extends Rule[LogicalPlan] {
       LocalLimit(le, udf.copy(child = maybePushLocalLimit(le, udf.child)))
     case LocalLimit(le, p @ Project(_, udf: ArrowEvalPython)) =>
       LocalLimit(le, p.copy(child = udf.copy(child = maybePushLocalLimit(le, udf.child))))
+    case LocalLimit(le, udf: EvalExternalUDF) =>
+      LocalLimit(le, udf.copy(child = maybePushLocalLimit(le, udf.child)))
+    case LocalLimit(le, p @ Project(_, udf: EvalExternalUDF)) =>
+      LocalLimit(le, p.copy(child = udf.copy(child = maybePushLocalLimit(le, udf.child))))
   }
 }
 
@@ -2311,6 +2315,7 @@ object PushPredicateThroughNonJoin extends Rule[LogicalPlan] with PredicateHelpe
     case _: Sort => true
     case _: BatchEvalPython => true
     case _: ArrowEvalPython => true
+    case _: EvalExternalUDF => true
     case _: Expand => true
     case _: BinBy => true
     case _ => false

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/logicalExternalUDFOperators.scala

@@ -19,7 +19,7 @@ package org.apache.spark.sql.catalyst.plans.logical
 
 import org.apache.spark.annotation.Experimental
 import org.apache.spark.resource.ResourceProfile
-import org.apache.spark.sql.catalyst.expressions.{Attribute,
+import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeSet,
   ExternalUserDefinedFunction}
 import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes
 import org.apache.spark.sql.types.StructType
@@ -38,6 +38,54 @@ trait ExternalUDF extends UnaryNode {
   def workerSpec: UDFWorkerSpecification
 }
 
+/**
+ * :: Experimental ::
+ * Logical plan node for evaluating scalar UDF expressions in one external
+ * worker session.
+ *
+ * Every UDF call in this node uses [[workerSpec]]. Workers without UDF
+ * chaining support receive exactly one non-nested UDF. Workers with chaining
+ * support may receive nested chains or multiple independent UDF roots.
+ *
+ * @param workerSpec Specification describing the UDF worker.
+ * @param udfs UDF expression roots evaluated by the worker session.
+ * @param resultAttrs One output attribute for each UDF expression root.
+ * @param child Input relation for the UDFs.
+ */
+@Experimental
+case class EvalExternalUDF(
+    workerSpec: UDFWorkerSpecification,
+    udfs: Seq[ExternalUserDefinedFunction],
+    resultAttrs: Seq[Attribute],
+    child: LogicalPlan)
+  extends ExternalUDF {
+
+  assert(udfs.length == resultAttrs.length,
+    "Input UDFs and result attributes must have the same size")
+  assert(udfs.zip(resultAttrs).forall { case (udf, resultAttr) =>
+    udf.dataType == resultAttr.dataType && udf.nullable == resultAttr.nullable
+  }, "Input UDFs and result attributes must have matching types and nullability")
+
+  private val allUdfs = udfs.flatMap(_.collect {
+    case udf: ExternalUserDefinedFunction => udf
+  })
+  assert(allUdfs.forall(_.workerSpec == workerSpec),
+    "All UDFs in one evaluation node must use the same worker specification")
+  assert(workerSpec.getCapabilities.getSupportsUdfChaining || allUdfs.size == 1,
+    "A worker without UDF chaining support can evaluate only one UDF per node")
+
+  override def output: Seq[Attribute] = child.output ++ resultAttrs
+
+  override def producedAttributes: AttributeSet = AttributeSet(resultAttrs)
+
+  override def maxRows: Option[Long] = child.maxRows
+
+  override def maxRowsPerPartition: Option[Long] = child.maxRowsPerPartition
+
+  override protected def withNewChildInternal(newChild: LogicalPlan): EvalExternalUDF =
+    copy(child = newChild)
+}
+
 /**
  * :: Experimental ::
  * Logical plan node for mapPartitions-style UDF execution in an
@@ -59,6 +107,9 @@ case class MapPartitionsExternalUDF(
     child: LogicalPlan)
   extends ExternalUDF {
 
+  assert(function.workerSpec == workerSpec,
+    "The map partitions UDF must use the node's worker specification")
+
   val nodeOutputAttributes = toAttributes(
     function.dataType.asInstanceOf[StructType]
   )

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala

@@ -5143,15 +5143,24 @@ object SQLConf {
 
   val UNIFIED_UDF_EXECUTION_ENABLED =
     buildConf("spark.sql.execution.udf.unified.execution.enabled")
-      .doc("When true, UDFs that support the language-agnostic " +
-        "UDF worker protocol are executed via the unified, " +
-        "external UDF worker framework instead of the " +
-        "language-specific runners. Experimental.")
+      .doc("When true, enable planning and execution through the language-agnostic " +
+        "external UDF worker framework. When false, external UDF expressions are rejected. " +
+        "Experimental.")
       .version("4.2.0")
       .withBindingPolicy(ConfigBindingPolicy.SESSION)
       .booleanConf
       .createWithDefault(false)
 
+  val UNIFIED_UDF_EXECUTION_CONVERT_PYTHON_UDF_ENABLED =
+    buildConf("spark.sql.execution.udf.unified.convertPythonUDF.enabled")
+      .doc("When unified UDF execution is enabled, convert legacy scalar PythonUDF " +
+        "expressions to external UDF expressions. When false, legacy scalar PythonUDF " +
+        "expressions are rejected by the unified execution path. Experimental.")
+      .version("4.3.0")
+      .withBindingPolicy(ConfigBindingPolicy.SESSION)
+      .booleanConf
+      .createWithDefault(false)
+
   val PYTHON_UDF_ARROW_ENABLED =
     buildConf("spark.sql.execution.pythonUDF.arrow.enabled")
       .doc("Enable Arrow optimization in regular Python UDFs. This optimization " +

sql/core/src/main/scala/org/apache/spark/sql/execution/SparkOptimizer.scala

@@ -17,6 +17,7 @@
 
 package org.apache.spark.sql.execution
 
+import org.apache.spark.SparkConf
 import org.apache.spark.sql.ExperimentalMethods
 import org.apache.spark.sql.catalyst.catalog.SessionCatalog
 import org.apache.spark.sql.catalyst.optimizer._
@@ -26,14 +27,18 @@ import org.apache.spark.sql.connector.catalog.CatalogManager
 import org.apache.spark.sql.execution.datasources.{MarkSingleTaskExecution, PruneFileSourcePartitions, PullOutVariantExtractions, PushVariantIntoScan, SchemaPruning, V1Writes}
 import org.apache.spark.sql.execution.datasources.v2.{GroupBasedRowLevelOperationScanPlanning, OptimizeMetadataOnlyDeleteFromTable, V2ScanPartitioningAndOrdering, V2ScanRelationPushDown, V2Writes}
 import org.apache.spark.sql.execution.dynamicpruning.{CleanupDynamicPruningFilters, PartitionPruning, RowLevelOperationRuntimeGroupFiltering}
+import org.apache.spark.sql.execution.externalUDF.ExtractExternalUDFs
 import org.apache.spark.sql.execution.python.{ExtractGroupingPythonUDFFromAggregate, ExtractPythonUDFFromAggregate, ExtractPythonUDFs, ExtractPythonUDTFs}
 
 class SparkOptimizer(
     catalogManager: CatalogManager,
     catalog: SessionCatalog,
-    experimentalMethods: ExperimentalMethods)
+    experimentalMethods: ExperimentalMethods,
+    sparkConf: SparkConf)
   extends Optimizer(catalogManager) {
 
+  private val extractExternalUDFs = new ExtractExternalUDFs(sparkConf)
+
   override def earlyScanPushDownRules: Seq[Rule[LogicalPlan]] =
     // TODO: move SchemaPruning into catalyst
     Seq(
@@ -88,6 +93,7 @@ class SparkOptimizer(
       ExtractPythonUDFFromAggregate,
       // This must be executed after `ExtractPythonUDFFromAggregate` and before `ExtractPythonUDFs`.
       ExtractGroupingPythonUDFFromAggregate,
+      extractExternalUDFs,
       ExtractPythonUDFs,
       ExtractPythonUDTFs,
       // The eval-python node may be between Project/Filter and the scan node, which breaks
@@ -114,6 +120,7 @@ class SparkOptimizer(
       ExtractPythonUDFFromJoinCondition.ruleName,
       ExtractPythonUDFFromAggregate.ruleName,
       ExtractGroupingPythonUDFFromAggregate.ruleName,
+      extractExternalUDFs.ruleName,
       ExtractPythonUDFs.ruleName,
       GroupBasedRowLevelOperationScanPlanning.ruleName,
       V2ScanRelationPushDown.ruleName,

sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala

@@ -1084,6 +1084,9 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
         execution.externalUDF.MapPartitionsExternalUDFExec(
           workerSpec, functionExpr,
           isBarrier, profile, planLater(child)) :: Nil
+      case logical.EvalExternalUDF(workerSpec, udfs, resultAttrs, child) =>
+        execution.externalUDF.EvalExternalUDFExec(
+          workerSpec, udfs, resultAttrs, planLater(child)) :: Nil
       case logical.AttachDistributedSequence(attr, child, cache) =>
         execution.python.AttachDistributedSequenceExec(attr, planLater(child), cache) :: Nil
       case logical.PythonWorkerLogs(jsonAttr) =>

sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/EvalExternalUDFExec.scala

@@ -0,0 +1,59 @@
+/*
+ * 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 org.apache.spark.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.execution.SparkPlan
+import org.apache.spark.udf.worker.UDFWorkerSpecification
+
+/**
+ * :: Experimental ::
+ * Physical plan node that evaluates scalar UDF expressions in one external
+ * worker session.
+ */
+@Experimental
+case class EvalExternalUDFExec(
+    workerSpec: UDFWorkerSpecification,
+    udfs: Seq[ExternalUserDefinedFunction],
+    resultAttrs: Seq[Attribute],
+    child: SparkPlan)
+  extends ExternalUDFExec {
+
+  override def output: Seq[Attribute] = child.output ++ resultAttrs
+
+  override def producedAttributes: AttributeSet = AttributeSet(resultAttrs)
+
+  override protected def doExecute(): RDD[InternalRow] = {
+    child.execute().mapPartitionsInternal { _ =>
+      withUDFWorkerSession(TaskContext.get(), securityScope = None) { _ =>
+        // TODO [SPARK-55278]: Stream scalar UDF rows to and from the worker.
+        // scalastyle:off throwerror
+        throw new NotImplementedError("doExecute() is not yet implemented.")
+        // scalastyle:on throwerror
+      }
+    }
+  }
+
+  override protected def withNewChildInternal(newChild: SparkPlan): EvalExternalUDFExec =
+    copy(child = newChild)
+}

sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/ExternalUDFPlanner.scala

@@ -105,6 +105,7 @@ class UnifiedExternalUDFPlanner(
         pythonUdf.func, conf)
     val udf = ExternalUserDefinedFunction(
       name = Some(pythonUdf.name),
+      workerSpec = workerSpec,
       payload = pythonUdf.func.command.toArray,
       dataType = pythonUdf.dataType,
       children = Seq.empty,