diff --git a/common/utils/src/main/resources/error/error-conditions.json b/common/utils/src/main/resources/error/error-conditions.json index d27b0c02b1d5c..b164c0a2a5e3d 100644 --- a/common/utils/src/main/resources/error/error-conditions.json +++ b/common/utils/src/main/resources/error/error-conditions.json @@ -8592,6 +8592,11 @@ "Drop the namespace ." ] }, + "EXTERNAL_UDF" : { + "message" : [ + "External UDF expressions. Set to true to enable them." + ] + }, "GEOSPATIAL_DISABLED" : { "message" : [ "Geospatial feature is disabled." @@ -8742,6 +8747,11 @@ "Python UDF in the ON clause of a 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 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", diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ExternalUserDefinedFunction.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ExternalUserDefinedFunction.scala index ede4adf32e119..b2b9b0071c75b 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ExternalUserDefinedFunction.scala +++ b/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], diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala index 15833bb74ba38..68f21139f1149 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala +++ b/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 diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/logicalExternalUDFOperators.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/logicalExternalUDFOperators.scala index f29c90772d4f5..d466c5705f161 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/logicalExternalUDFOperators.scala +++ b/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] ) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index ee9d066833609..613bc96193f7f 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/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 " + diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkOptimizer.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkOptimizer.scala index 0994220afe0c1..df59f60868d0d 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkOptimizer.scala +++ b/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, diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala index 996ecac61b728..73307b34b7663 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala +++ b/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) => diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/EvalExternalUDFExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/EvalExternalUDFExec.scala new file mode 100644 index 0000000000000..ca476812f1d20 --- /dev/null +++ b/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) +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/ExternalUDFPlanner.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/ExternalUDFPlanner.scala index c90a8af89783f..54cd23a0a1f89 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/ExternalUDFPlanner.scala +++ b/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, diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/ExtractExternalUDFs.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/ExtractExternalUDFs.scala new file mode 100644 index 0000000000000..d014808584eb6 --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/ExtractExternalUDFs.scala @@ -0,0 +1,199 @@ +/* + * 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.collection.mutable.ArrayBuffer + +import org.apache.spark.{SparkConf, SparkException, SparkUnsupportedOperationException} +import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.catalyst.expressions.PythonUDF.isScalarPythonUDF +import org.apache.spark.sql.catalyst.plans.logical._ +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.catalyst.trees.TreePattern.{EXTERNAL_UDF, PYTHON_UDF} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.udf.worker.UDFWorkerSpecification + +/** + * Extracts scalar external UDF expressions into logical evaluation nodes. + * + * Each [[EvalExternalUDF]] represents one worker session. A worker that does + * not advertise UDF chaining receives one UDF call per node. A worker that + * supports chaining may evaluate unary nested chains and fuse independent + * UDF roots that use the same worker specification. + * + * When unified UDF execution and legacy Python UDF conversion are enabled, + * scalar [[PythonUDF]] expressions are converted to + * [[ExternalUserDefinedFunction]] expressions before extraction. + */ +private[sql] class ExtractExternalUDFs(sparkConf: SparkConf) extends Rule[LogicalPlan] { + + override def apply(plan: LogicalPlan): LogicalPlan = plan match { + // A correlated subquery is rewritten as a join and revisits this rule later. + case subquery: Subquery if subquery.correlated => plan + case _ if !conf.getConf(SQLConf.UNIFIED_UDF_EXECUTION_ENABLED) => + if (plan.containsPattern(EXTERNAL_UDF)) { + throw new SparkUnsupportedOperationException( + errorClass = "UNSUPPORTED_FEATURE.EXTERNAL_UDF", + messageParameters = Map( + "config" -> SQLConf.UNIFIED_UDF_EXECUTION_ENABLED.key)) + } + plan + case _ => + val externalPlan = convertPythonUDFs(plan) + externalPlan.transformUpWithPruning(_.containsPattern(EXTERNAL_UDF)) { + // These nodes already own their external UDF expressions. + case udfPlan: ExternalUDF => udfPlan + case other => extract(other) + } + } + + private def convertPythonUDFs(plan: LogicalPlan): LogicalPlan = { + val convertPythonUDF = + conf.getConf(SQLConf.UNIFIED_UDF_EXECUTION_CONVERT_PYTHON_UDF_ENABLED) + plan.transformUpWithPruning(_.containsPattern(PYTHON_UDF)) { operator => + operator.transformExpressionsUpWithPruning(_.containsPattern(PYTHON_UDF)) { + case udf: PythonUDF if isScalarPythonUDF(udf) => + if (!convertPythonUDF) { + throw new SparkUnsupportedOperationException( + errorClass = "UNSUPPORTED_FEATURE.PYTHON_UDF_TO_EXTERNAL_UDF", + messageParameters = Map( + "config" -> + SQLConf.UNIFIED_UDF_EXECUTION_CONVERT_PYTHON_UDF_ENABLED.key)) + } + ExternalUserDefinedFunction( + name = Some(udf.name), + workerSpec = PythonUDFWorkerSpecification.fromPythonFunction(udf.func, sparkConf), + payload = udf.func.command.toArray, + dataType = udf.dataType, + children = udf.children, + inputTypes = None, + udfDeterministic = udf.udfDeterministic, + udfNullable = udf.nullable, + resultId = udf.resultId) + } + } + } + + private def supportsUDFChaining(workerSpec: UDFWorkerSpecification): Boolean = { + workerSpec.hasCapabilities && + workerSpec.getCapabilities.hasSupportsUdfChaining && + workerSpec.getCapabilities.getSupportsUdfChaining + } + + private def containsExternalUDF(expression: Expression): Boolean = { + expression.exists(_.isInstanceOf[ExternalUserDefinedFunction]) + } + + /** Returns whether `udf` can be evaluated in one worker session. */ + @scala.annotation.tailrec + private def isEvaluable(udf: ExternalUserDefinedFunction): Boolean = { + udf.children match { + case Seq(child: ExternalUserDefinedFunction) + if supportsUDFChaining(udf.workerSpec) && child.workerSpec == udf.workerSpec => + isEvaluable(child) + case children => + !children.exists(containsExternalUDF) + } + } + + private def sameUDF( + left: ExternalUserDefinedFunction, + right: ExternalUserDefinedFunction): Boolean = { + if (left.deterministic && right.deterministic) { + left.semanticEquals(right) + } else { + left.resultId == right.resultId + } + } + + private def collectEvaluableUDFs(expression: Expression): Seq[ExternalUserDefinedFunction] = { + expression match { + case udf: ExternalUserDefinedFunction if isEvaluable(udf) => Seq(udf) + case other => other.children.flatMap(collectEvaluableUDFs) + } + } + + /** Selects the UDF roots that can share the next worker session. */ + private def collectSessionUDFs(plan: LogicalPlan): Seq[ExternalUserDefinedFunction] = { + val candidates = plan.expressions + .flatMap(collectEvaluableUDFs) + .filter(_.references.subsetOf(plan.inputSet)) + + val distinct = ArrayBuffer.empty[ExternalUserDefinedFunction] + candidates.foreach { candidate => + if (!distinct.exists(sameUDF(_, candidate))) { + distinct += candidate + } + } + + distinct.headOption.toSeq.flatMap { first => + if (supportsUDFChaining(first.workerSpec)) { + distinct.filter(_.workerSpec == first.workerSpec) + } else { + Seq(first) + } + } + } + + /** Extracts one worker session and recursively extracts the remaining UDFs. */ + private def extract(plan: LogicalPlan): LogicalPlan = { + val udfs = collectSessionUDFs(plan) + if (udfs.isEmpty) { + plan + } else { + val udfToAttribute = ArrayBuffer.empty[(ExternalUserDefinedFunction, Attribute)] + + def resultFor(udf: ExternalUserDefinedFunction): Option[Attribute] = { + udfToAttribute.collectFirst { + case (candidate, result) if sameUDF(candidate, udf) => result + } + } + + val newChildren = plan.children.map { child => + val validUdfs = udfs.filter(_.references.subsetOf(child.outputSet)) + if (validUdfs.isEmpty) { + child + } else { + val resultAttrs = validUdfs.zipWithIndex.map { case (udf, index) => + AttributeReference(s"externalUDF$index", udf.dataType, udf.nullable)() + } + val evaluation = EvalExternalUDF( + validUdfs.head.workerSpec, validUdfs, resultAttrs, child) + udfToAttribute ++= validUdfs.zip(resultAttrs) + evaluation + } + } + + udfs.filter(resultFor(_).isEmpty).foreach { udf => + throw SparkException.internalError( + s"Invalid external UDF $udf, requires attributes from more than one child.") + } + + val rewritten = plan.withNewChildren(newChildren).transformExpressions { + case udf: ExternalUserDefinedFunction => resultFor(udf).getOrElse(udf) + } + + val newPlan = extract(rewritten) + if (newPlan.output != plan.output) { + Project(plan.output, newPlan) + } else { + newPlan + } + } + } +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/MapPartitionsExternalUDFExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/MapPartitionsExternalUDFExec.scala index dbb5bfef007e7..a38142d2e2261 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/MapPartitionsExternalUDFExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/MapPartitionsExternalUDFExec.scala @@ -51,6 +51,9 @@ case class MapPartitionsExternalUDFExec( child: SparkPlan) extends ExternalUDFExec { + assert(function.workerSpec == workerSpec, + "The map partitions UDF must use the node's worker specification") + // Map partitions always operate on StructTypes override def output: Seq[Attribute] = toAttributes( function.dataType.asInstanceOf[StructType] diff --git a/sql/core/src/main/scala/org/apache/spark/sql/internal/BaseSessionStateBuilder.scala b/sql/core/src/main/scala/org/apache/spark/sql/internal/BaseSessionStateBuilder.scala index 52c6821d00011..da16396972591 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/internal/BaseSessionStateBuilder.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/internal/BaseSessionStateBuilder.scala @@ -323,7 +323,7 @@ abstract class BaseSessionStateBuilder( * Note: this depends on `catalog` and `experimentalMethods` fields. */ protected def optimizer: Optimizer = { - new SparkOptimizer(catalogManager, catalog, experimentalMethods) { + new SparkOptimizer(catalogManager, catalog, experimentalMethods, session.sparkContext.conf) { override def earlyScanPushDownRules: Seq[Rule[LogicalPlan]] = super.earlyScanPushDownRules ++ customEarlyScanPushDownRules diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/externalUDF/ExternalUDFPlanningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/externalUDF/ExternalUDFPlanningSuite.scala index 7f10e349d0dac..c9a22b18be809 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/externalUDF/ExternalUDFPlanningSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/externalUDF/ExternalUDFPlanningSuite.scala @@ -20,13 +20,13 @@ import scala.collection.mutable.ArrayBuffer import scala.jdk.CollectionConverters._ import scala.reflect.ClassTag -import org.apache.spark.SparkConf +import org.apache.spark.{SparkConf, SparkUnsupportedOperationException} import org.apache.spark.api.python.{PythonEvalType, SimplePythonFunction} import org.apache.spark.sql.DataFrame import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, - MapInPandas, MapPartitionsExternalUDF} + EvalExternalUDF, MapInPandas, MapPartitionsExternalUDF} import org.apache.spark.sql.execution.SparkPlan -import org.apache.spark.sql.execution.python.{MapInPandasExec, +import org.apache.spark.sql.execution.python.{BatchEvalPythonExec, MapInPandasExec, UserDefinedPythonFunction} import org.apache.spark.sql.functions.col import org.apache.spark.sql.internal.SQLConf @@ -62,11 +62,24 @@ trait ExternalUDFPlanningTestBase extends SharedSparkSession { pythonEvalType = PythonEvalType.SQL_MAP_PANDAS_ITER_UDF, udfDeterministic = true) + protected val scalarPythonUDF: UserDefinedPythonFunction = + UserDefinedPythonFunction( + name = "dummyScalarPythonUDF", + func = dummyPythonFunction, + dataType = IntegerType, + pythonEvalType = PythonEvalType.SQL_BATCHED_UDF, + udfDeterministic = true) + protected def applyMapInPandas(): DataFrame = { val inputDF = Seq(("hello", 1)).toDF("a", "b") inputDF.mapInPandas(mapInPandasUDF(col("a"), col("b"))) } + protected def applyScalarPythonUDF(): DataFrame = { + val inputDF = Seq(("hello", 1)).toDF("a", "b") + inputDF.select(scalarPythonUDF(col("b"))) + } + protected def assertLogicalNode[T <: LogicalPlan: ClassTag]( df: DataFrame): Unit = { val tag = implicitly[ClassTag[T]] @@ -78,6 +91,17 @@ trait ExternalUDFPlanningTestBase extends SharedSparkSession { " in logical plan") } + protected def assertOptimizedLogicalNode[T <: LogicalPlan: ClassTag]( + df: DataFrame): Unit = { + val tag = implicitly[ClassTag[T]] + val node = df.queryExecution.optimizedPlan.collectFirst { + case n if tag.runtimeClass.isInstance(n) => n + } + assert(node.isDefined, + s"Expected ${tag.runtimeClass.getSimpleName}" + + " in optimized logical plan") + } + protected def assertPhysicalNode[T <: SparkPlan: ClassTag]( df: DataFrame): Unit = { val tag = implicitly[ClassTag[T]] @@ -107,6 +131,11 @@ class ClassicUDFPlanningSuite val result = applyMapInPandas() assertPhysicalNode[MapInPandasExec](result) } + + test("scalar Python UDF uses BatchEvalPythonExec physical node") { + val result = applyScalarPythonUDF() + assertPhysicalNode[BatchEvalPythonExec](result) + } } /** @@ -117,8 +146,9 @@ class UnifiedUDFPlanningSuite extends ExternalUDFPlanningTestBase { override def sparkConf: SparkConf = - super.sparkConf.set( - SQLConf.UNIFIED_UDF_EXECUTION_ENABLED.key, "true") + super.sparkConf + .set(SQLConf.UNIFIED_UDF_EXECUTION_ENABLED.key, "true") + .set(SQLConf.UNIFIED_UDF_EXECUTION_CONVERT_PYTHON_UDF_ENABLED.key, "true") test("mapInPandas uses MapPartitionsExternalUDF logical node") { val result = applyMapInPandas() @@ -130,4 +160,37 @@ class UnifiedUDFPlanningSuite val result = applyMapInPandas() assertPhysicalNode[MapPartitionsExternalUDFExec](result) } + + test("scalar Python UDF uses EvalExternalUDF logical node") { + val result = applyScalarPythonUDF() + assertOptimizedLogicalNode[EvalExternalUDF](result) + } + + test("scalar Python UDF uses EvalExternalUDFExec physical node") { + val result = applyScalarPythonUDF() + assertPhysicalNode[EvalExternalUDFExec](result) + } +} + +/** + * Tests that the unified execution path rejects legacy scalar Python UDF + * expressions when their conversion config is disabled. + */ +class UnsupportedLegacyPythonUDFPlanningSuite + extends ExternalUDFPlanningTestBase { + + override def sparkConf: SparkConf = + super.sparkConf.set( + SQLConf.UNIFIED_UDF_EXECUTION_ENABLED.key, "true") + + test("legacy scalar Python UDF is rejected") { + val exception = intercept[SparkUnsupportedOperationException] { + applyScalarPythonUDF().queryExecution.optimizedPlan + } + checkError( + exception = exception, + condition = "UNSUPPORTED_FEATURE.PYTHON_UDF_TO_EXTERNAL_UDF", + parameters = Map( + "config" -> SQLConf.UNIFIED_UDF_EXECUTION_CONVERT_PYTHON_UDF_ENABLED.key)) + } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/externalUDF/ExtractExternalUDFsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/externalUDF/ExtractExternalUDFsSuite.scala new file mode 100644 index 0000000000000..54901b7a90a96 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/externalUDF/ExtractExternalUDFsSuite.scala @@ -0,0 +1,160 @@ +/* + * 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 java.nio.charset.StandardCharsets + +import org.apache.spark.{SparkConf, SparkUnsupportedOperationException} +import org.apache.spark.sql.catalyst.expressions.{Add, Alias, AttributeReference, + ExternalUserDefinedFunction, Expression} +import org.apache.spark.sql.catalyst.plans.PlanTest +import org.apache.spark.sql.catalyst.plans.logical.{EvalExternalUDF, LocalRelation, + LogicalPlan, Project} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.IntegerType +import org.apache.spark.udf.worker.{DirectWorker, ProcessCallable, UDFWorkerProperties, + UDFWorkerSpecification, WorkerCapabilities, WorkerEnvironment} + +class ExtractExternalUDFsSuite extends PlanTest { + + private val input = AttributeReference("input", IntegerType, nullable = false)() + private val relation = LocalRelation(Seq(input)) + private val extractRule = new ExtractExternalUDFs(new SparkConf(false)) + + private def workerSpec(name: String, supportsChaining: Boolean): UDFWorkerSpecification = { + val capabilities = WorkerCapabilities.newBuilder() + .setSupportsUdfChaining(supportsChaining) + val direct = DirectWorker.newBuilder() + .setRunner(ProcessCallable.newBuilder().addCommand(name)) + .setProperties(UDFWorkerProperties.newBuilder()) + + UDFWorkerSpecification.newBuilder() + .setEnvironment(WorkerEnvironment.newBuilder()) + .setCapabilities(capabilities) + .setDirect(direct) + .build() + } + + private def udf( + name: String, + spec: UDFWorkerSpecification, + children: Seq[Expression]): ExternalUserDefinedFunction = { + ExternalUserDefinedFunction( + name = Some(name), + workerSpec = spec, + payload = name.getBytes(StandardCharsets.UTF_8), + dataType = IntegerType, + children = children, + inputTypes = None, + udfDeterministic = true, + udfNullable = false) + } + + private def extract(expression: Expression): LogicalPlan = { + withSQLConf(SQLConf.UNIFIED_UDF_EXECUTION_ENABLED.key -> "true") { + extractRule(Project(Seq(Alias(expression, "result")()), relation)) + } + } + + private def evalNodes(plan: LogicalPlan): Seq[EvalExternalUDF] = { + plan.collect { case eval: EvalExternalUDF => eval } + } + + private def callCount(expression: Expression): Int = { + expression.collect { case _: ExternalUserDefinedFunction => 1 }.size + } + + test("a worker without chaining support gets one UDF call per node") { + val spec = workerSpec("single", supportsChaining = false) + val expression = udf("outer", spec, + Seq(udf("middle", spec, Seq(udf("inner", spec, Seq(input)))))) + + val nodes = evalNodes(extract(expression)) + assert(nodes.size == 3) + assert(nodes.forall(_.udfs.size == 1)) + assert(nodes.forall(_.udfs.forall(callCount(_) == 1))) + } + + test("a chaining worker evaluates a unary UDF chain in one node") { + val spec = workerSpec("chained", supportsChaining = true) + val expression = udf("outer", spec, + Seq(udf("middle", spec, Seq(udf("inner", spec, Seq(input)))))) + + val nodes = evalNodes(extract(expression)) + assert(nodes.size == 1) + assert(nodes.head.udfs.size == 1) + assert(callCount(nodes.head.udfs.head) == 3) + } + + test("a chaining worker fuses independent UDF roots in one node") { + val spec = workerSpec("fused", supportsChaining = true) + val expression = Add(udf("left", spec, Seq(input)), udf("right", spec, Seq(input))) + + val nodes = evalNodes(extract(expression)) + assert(nodes.size == 1) + assert(nodes.head.udfs.size == 2) + assert(nodes.head.udfs.forall(callCount(_) == 1)) + } + + test("dependent branches use separate session nodes") { + val spec = workerSpec("branched", supportsChaining = true) + val left = udf("left", spec, Seq(input)) + val right = udf("right", spec, Seq(input)) + val expression = udf("outer", spec, Seq(Add(left, right))) + + val nodes = evalNodes(extract(expression)) + assert(nodes.size == 2) + assert(nodes.map(_.udfs.size).sorted == Seq(1, 2)) + assert(nodes.flatMap(_.udfs).map(callCount).sorted == Seq(1, 1, 1)) + } + + test("UDFs with different worker specifications use separate nodes") { + val outerSpec = workerSpec("outer-worker", supportsChaining = true) + val innerSpec = workerSpec("inner-worker", supportsChaining = true) + val expression = udf("outer", outerSpec, Seq(udf("inner", innerSpec, Seq(input)))) + + val nodes = evalNodes(extract(expression)) + assert(nodes.size == 2) + assert(nodes.map(_.workerSpec).toSet == Set(outerSpec, innerSpec)) + assert(nodes.forall(_.udfs.forall(callCount(_) == 1))) + } + + test("independent UDFs are separate sessions when chaining is unsupported") { + val spec = workerSpec("single", supportsChaining = false) + val expression = Add(udf("left", spec, Seq(input)), udf("right", spec, Seq(input))) + + val nodes = evalNodes(extract(expression)) + assert(nodes.size == 2) + assert(nodes.forall(_.udfs.size == 1)) + } + + test("external UDF expressions are rejected when unified execution is disabled") { + val spec = workerSpec("disabled", supportsChaining = true) + val plan = Project(Seq(Alias(udf("external", spec, Seq(input)), "result")()), relation) + + val exception = withSQLConf(SQLConf.UNIFIED_UDF_EXECUTION_ENABLED.key -> "false") { + intercept[SparkUnsupportedOperationException] { + extractRule(plan) + } + } + checkError( + exception = exception, + condition = "UNSUPPORTED_FEATURE.EXTERNAL_UDF", + parameters = Map("config" -> SQLConf.UNIFIED_UDF_EXECUTION_ENABLED.key)) + } +} diff --git a/udf/worker/proto/src/main/protobuf/worker_spec.proto b/udf/worker/proto/src/main/protobuf/worker_spec.proto index 83dac4f962e5f..8c34828a06928 100644 --- a/udf/worker/proto/src/main/protobuf/worker_spec.proto +++ b/udf/worker/proto/src/main/protobuf/worker_spec.proto @@ -161,7 +161,15 @@ message WorkerCapabilities { // (Optional) optional bool supports_reuse = 4; - // To be extended with UDF chaining, ... + // Whether one UDF session can evaluate multiple UDF calls. This includes + // nested chains, where one UDF consumes another UDF's output, and fused + // independent calls over the same input. + // + // When this is not supported, the engine creates a separate session for + // every UDF call. + // + // (Optional) + optional bool supports_udf_chaining = 5; } // The worker that will be created to process UDFs