haiyangsun-db commented on code in PR #57670: URL: https://github.com/apache/spark/pull/57670#discussion_r3926983275
########## sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/PlanExternalUDFs.scala: ########## @@ -0,0 +1,266 @@ +/* + * 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.internal.Logging +import org.apache.spark.internal.LogKeys.JOIN_CONDITION +import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression +import org.apache.spark.sql.catalyst.plans.InnerLike +import org.apache.spark.sql.catalyst.plans.logical._ +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.catalyst.trees.TreePattern.{AGGREGATE, EXTERNAL_UDF, JOIN} +import org.apache.spark.sql.errors.QueryCompilationErrors +import org.apache.spark.sql.internal.SQLConf + +/** + * Converts each scalar external UDF expression into a separate logical evaluation node. + * Join-condition handling mirrors `ExtractPythonUDFFromJoinCondition`. + * + * TODO(SPARK-55278): Add an external UDF equivalent of `ExtractPythonUDFFromLambda`. + * TODO(SPARK-55278): Revisit sharing placement logic with the Python UDF extractors after + * external UDF planning semantics stabilize. + */ +private[sql] object PlanExternalUDFs + extends Rule[LogicalPlan] with Logging with PredicateHelper { + + 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 QueryCompilationErrors.externalUDFsDisabledError( + SQLConf.UNIFIED_UDF_EXECUTION_ENABLED.key) + } + plan + case _ => + var preparedPlan = extractExternalUDFFromJoinCondition(plan) + preparedPlan = extractExternalUDFFromAggregate(preparedPlan) + preparedPlan = extractGroupingExternalUDFFromAggregate(preparedPlan) + preparedPlan.transformUpWithPruning(_.containsPattern(EXTERNAL_UDF)) { + // These nodes already own their external UDF expressions. + case udfPlan: ExternalUDF => udfPlan + case other => extract(other) + } + } + + private def hasUnevaluableExternalUDF(expression: Expression, join: Join): Boolean = { + expression.exists { + case udf: ExternalUserDefinedFunction => + !canEvaluate(udf, join.left) && !canEvaluate(udf, join.right) + case _ => false + } + } + + private def extractExternalUDFFromJoinCondition(plan: LogicalPlan): LogicalPlan = { + plan.transformUpWithPruning(_.containsAllPatterns(EXTERNAL_UDF, JOIN)) { + case join @ Join(_, _, joinType, Some(condition), _) + if hasUnevaluableExternalUDF(condition, join) => + if (!joinType.isInstanceOf[InnerLike]) { + // Match `PYTHON_UDF_IN_ON_CLAUSE`: moving a cross-side UDF to a post-join filter + // changes the semantics of non-inner joins. + throw QueryCompilationErrors.useExternalUDFInJoinConditionUnsupportedError(joinType) + } + + val (udfConditions, otherConditions) = splitConjunctivePredicates(condition) + .partition(hasUnevaluableExternalUDF(_, join)) + val newCondition = if (otherConditions.isEmpty) { + logWarning(log"The join condition:${MDC(JOIN_CONDITION, condition)} " + + log"of the join plan contains external UDFs only, " + + log"so it will be moved out and the join plan will become a cross join.") + None + } else { + Some(otherConditions.reduceLeft(And)) + } + Filter(udfConditions.reduceLeft(And), join.copy(condition = newCondition)) + } + } + + private def belongsToAggregate( + expression: Expression, + groupingExpressions: ExpressionSet): Boolean = { + expression.isInstanceOf[AggregateExpression] || + groupingExpressions.contains(expression) + } + + private def hasExternalUDFOverAggregate( + expression: Expression, + groupingExpressions: ExpressionSet): Boolean = { + expression.exists { + case udf: ExternalUserDefinedFunction => + udf.references.isEmpty || udf.exists(belongsToAggregate(_, groupingExpressions)) + case _ => false + } + } + + private def extractExternalUDFFromAggregate(plan: LogicalPlan): LogicalPlan = { + plan.transformUpWithPruning(_.containsAllPatterns(EXTERNAL_UDF, AGGREGATE)) { + case aggregate: Aggregate => + val groupingExpressions = ExpressionSet(aggregate.groupingExpressions) + if (!aggregate.aggregateExpressions.exists( + hasExternalUDFOverAggregate(_, groupingExpressions))) { + aggregate + } else { + val projectExpressions = ArrayBuffer.empty[NamedExpression] + val aggregateExpressions = ArrayBuffer.empty[NamedExpression] + aggregate.aggregateExpressions.foreach { expression => + if (hasExternalUDFOverAggregate(expression, groupingExpressions)) { + val newExpression = expression.transformDown { + case child: Expression if belongsToAggregate(child, groupingExpressions) => + val alias = child match { + case named: NamedExpression => named + case other => Alias(other, "agg")() + } + aggregateExpressions += alias + alias.toAttribute + } + projectExpressions += newExpression.asInstanceOf[NamedExpression] + } else { + aggregateExpressions += expression + projectExpressions += expression.toAttribute + } + } + Project( + projectExpressions.toSeq, + aggregate.copy(aggregateExpressions = aggregateExpressions.toSeq)) + } + } + } + + private def hasExternalUDF(expression: Expression): Boolean = { + expression.exists(_.isInstanceOf[ExternalUserDefinedFunction]) + } + + private def extractGroupingExternalUDFFromAggregate(plan: LogicalPlan): LogicalPlan = { + plan.transformUpWithPruning(_.containsAllPatterns(EXTERNAL_UDF, AGGREGATE)) { + case aggregate: Aggregate if aggregate.groupingExpressions.exists(hasExternalUDF) => + val projectExpressions = ArrayBuffer.empty[NamedExpression] + val groupingExpressions = ArrayBuffer.empty[Expression] + val attributeMap = ArrayBuffer.empty[ + (ExternalUserDefinedFunction, NamedExpression)] + + def mappedAttribute(udf: ExternalUserDefinedFunction): Option[NamedExpression] = { + attributeMap.collectFirst { + case (candidate, attribute) if sameUDF(candidate, udf) => attribute + } + } + + aggregate.groupingExpressions.foreach { expression => + if (hasExternalUDF(expression)) { + val newExpression = expression.transformDown { + case udf: ExternalUserDefinedFunction => + assert(udf.udfDeterministic, Review Comment: Fixed. `NondeterministicExpressionCollection` now hoists nondeterministic `ExternalUserDefinedFunction` expressions, and the regression test verifies that the UDF is evaluated below `Aggregate` while the grouping expressions are deterministic. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/ExtractExternalUDFFromWindow.scala: ########## @@ -0,0 +1,75 @@ +/* + * 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 + +import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, Expression, ExprId, + ExternalUserDefinedFunction, NamedExpression, WindowExpression} +import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, Project, Window} +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.catalyst.trees.TreePattern.{EXTERNAL_UDF, WINDOW} + +/** + * Extracts external UDFs that are parents of window expressions from a [[Window]] operator. + * The window expressions are evaluated by the [[Window]], and [[PlanExternalUDFs]] subsequently + * converts the external UDFs in the new [[Project]] into evaluation nodes above it. + */ +private[sql] object ExtractExternalUDFFromWindow extends Rule[LogicalPlan] { + + private def containsExternalUDFOverWindowExpression(expression: Expression): Boolean = { + expression.exists { + case udf: ExternalUserDefinedFunction => + udf.exists(_.isInstanceOf[WindowExpression]) + case _ => false + } + } + + override def apply(plan: LogicalPlan): LogicalPlan = { + plan.transformWithPruning( + _.containsAllPatterns(EXTERNAL_UDF, WINDOW)) { + case window: Window + if window.windowExpressions.exists(containsExternalUDFOverWindowExpression) => + val windowProjectExprIds = mutable.Set.empty[ExprId] + val windowProjectList = mutable.ArrayBuffer.empty[NamedExpression] + val externalUdfProjectList = window.windowExpressions.map { expression => + if (containsExternalUDFOverWindowExpression(expression)) { + expression.transformDown { + case windowExpression: WindowExpression => + val alias = Alias(windowExpression, s"w_${windowProjectList.size}")() + windowProjectList += alias + alias.toAttribute + case attribute: Attribute if !windowProjectExprIds.contains(attribute.exprId) => + windowProjectList += attribute Review Comment: Fixed. Direct child attributes are no longer added to `Window.windowExpressions`. The focused regression keeps both UDF arguments and verifies that the Window output ExprIds are unique. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/PlanExternalUDFs.scala: ########## @@ -0,0 +1,266 @@ +/* + * 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.internal.Logging +import org.apache.spark.internal.LogKeys.JOIN_CONDITION +import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression +import org.apache.spark.sql.catalyst.plans.InnerLike +import org.apache.spark.sql.catalyst.plans.logical._ +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.catalyst.trees.TreePattern.{AGGREGATE, EXTERNAL_UDF, JOIN} +import org.apache.spark.sql.errors.QueryCompilationErrors +import org.apache.spark.sql.internal.SQLConf + +/** + * Converts each scalar external UDF expression into a separate logical evaluation node. + * Join-condition handling mirrors `ExtractPythonUDFFromJoinCondition`. + * + * TODO(SPARK-55278): Add an external UDF equivalent of `ExtractPythonUDFFromLambda`. + * TODO(SPARK-55278): Revisit sharing placement logic with the Python UDF extractors after + * external UDF planning semantics stabilize. + */ +private[sql] object PlanExternalUDFs + extends Rule[LogicalPlan] with Logging with PredicateHelper { + + 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 QueryCompilationErrors.externalUDFsDisabledError( + SQLConf.UNIFIED_UDF_EXECUTION_ENABLED.key) + } + plan + case _ => + var preparedPlan = extractExternalUDFFromJoinCondition(plan) + preparedPlan = extractExternalUDFFromAggregate(preparedPlan) + preparedPlan = extractGroupingExternalUDFFromAggregate(preparedPlan) + preparedPlan.transformUpWithPruning(_.containsPattern(EXTERNAL_UDF)) { + // These nodes already own their external UDF expressions. + case udfPlan: ExternalUDF => udfPlan + case other => extract(other) + } + } + + private def hasUnevaluableExternalUDF(expression: Expression, join: Join): Boolean = { + expression.exists { + case udf: ExternalUserDefinedFunction => + !canEvaluate(udf, join.left) && !canEvaluate(udf, join.right) + case _ => false + } + } + + private def extractExternalUDFFromJoinCondition(plan: LogicalPlan): LogicalPlan = { + plan.transformUpWithPruning(_.containsAllPatterns(EXTERNAL_UDF, JOIN)) { + case join @ Join(_, _, joinType, Some(condition), _) + if hasUnevaluableExternalUDF(condition, join) => + if (!joinType.isInstanceOf[InnerLike]) { + // Match `PYTHON_UDF_IN_ON_CLAUSE`: moving a cross-side UDF to a post-join filter + // changes the semantics of non-inner joins. + throw QueryCompilationErrors.useExternalUDFInJoinConditionUnsupportedError(joinType) + } + + val (udfConditions, otherConditions) = splitConjunctivePredicates(condition) + .partition(hasUnevaluableExternalUDF(_, join)) + val newCondition = if (otherConditions.isEmpty) { + logWarning(log"The join condition:${MDC(JOIN_CONDITION, condition)} " + + log"of the join plan contains external UDFs only, " + + log"so it will be moved out and the join plan will become a cross join.") + None + } else { + Some(otherConditions.reduceLeft(And)) Review Comment: Fixed. Added a mixed join-condition regression that verifies the equality remains on the `Join` while only the external-UDF predicate moves to the upper `Filter`. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/ExecuteExternalUDFExec.scala: ########## @@ -0,0 +1,61 @@ +/* + * 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.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.errors.QueryExecutionErrors +import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.udf.worker.UDFWorkerSpecification + +/** + * :: Experimental :: + * Physical plan node that evaluates one scalar UDF in an external worker process. + * + * @param udf UDF expression evaluated by the worker session. + * @param resultAttr Output attribute for the UDF expression. + * @param child Child plan providing input rows. + */ +@Experimental +case class ExecuteExternalUDFExec( + udf: ExternalUserDefinedFunction, Review Comment: Fixed. The Scaladoc now describes this as a scalar-UDF planning node and explicitly says worker execution remains unimplemented. -- 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]
