cloud-fan commented on code in PR #57670: URL: https://github.com/apache/spark/pull/57670#discussion_r3920863766
########## 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: **Nit (P3):** `Window.output` already prepends every attribute from `child.output`, so adding this ordinary child attribute to `windowProjectList` exposes the same ExprId a second time. This is reachable for `externalUDF(windowExpression, input)`. Please leave direct child attributes out of `Window.windowExpressions` and add a focused test that checks unique Window output ExprIds while preserving both UDF arguments. ########## 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: **Nit (P3):** The class Scaladoc currently says this node evaluates the UDF in a worker, but `doExecute` always throws `methodNotImplementedError`, and the PR intentionally leaves scalar execution out of scope. Please describe this as a planning placeholder with unimplemented execution until the worker path exists. ########## 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: **Nit (P3):** The suite exercises only the `otherConditions.isEmpty` path. Please add a mixed condition such as `leftKey = rightKey AND externalUDF(leftValue, rightValue)` and assert that the equality remains on the `Join` while only the external-UDF predicate moves to the post-join `Filter`; otherwise a regression here can silently turn a keyed join into a Cartesian plan. ########## 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: **Non-blocking (P2):** `PullOutNondeterministic` processes this Aggregate before `PlanExternalUDFs`, but `NondeterministicExpressionCollection` only collects `Nondeterministic` and nondeterministic `UserDefinedExpression` leaves. `ExternalUserDefinedFunction` implements neither, so a grouping UDF with `udfDeterministic = false` is left in the Aggregate and triggers the internal "should not appear in grouping expression" error before this assertion is reached. Please teach the collector to hoist nondeterministic external UDFs and add a grouping-key regression that verifies the UDF is evaluated below the Aggregate. **Recommended change:** Teach NondeterministicExpressionCollection to hoist nondeterministic ExternalUserDefinedFunction leaves and add a focused aggregate-grouping regression test. **Why this works:** Add an ExternalUserDefinedFunction case guarded by nondeterminism alongside the existing user-defined-expression case, allowing PullOutNondeterministic to materialize an alias in the child Project and replace the Aggregate grouping key before PlanExternalUDFs runs. **Scope:** Catalyst nondeterministic-expression collection plus the focused scalar external-UDF planning test suite. **Compatibility:** This changes only internal analysis of a currently failing unified external-UDF query; deterministic grouping behavior and the disabled configuration path remain unchanged. **Risks:** The new collection case must preserve canonicalized deduplication and must not hoist deterministic external UDFs. **Constraints:** Keep the existing PullOutNondeterministic ordering and its requirement that Aggregate grouping expressions are deterministic after rewriting. **Success:** A nondeterministic external UDF used as both grouping key and selected grouping expression analyzes without an internal error, is evaluated in the child plan, and leaves a deterministic attribute in the Aggregate. -- 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]
