haiyangsun-db commented on code in PR #57670: URL: https://github.com/apache/spark/pull/57670#discussion_r3926984937
########## sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/ExtractExternalUDFFromWindow.scala: ########## @@ -0,0 +1,71 @@ +/* + * 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, 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 + }.asInstanceOf[NamedExpression] + } else { + if (!windowProjectExprIds.contains(expression.exprId)) { + windowProjectList += expression + windowProjectExprIds += expression.exprId + } + expression.toAttribute + } + } + Project( + externalUdfProjectList, Review Comment: Fixed. The rewritten `Project` now prepends `window.child.output` in its original order, and the regression verifies that both the input column and UDF result are preserved. ########## 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) => Review Comment: Addressed with a different lifecycle choice. The planner is no longer retained by `SessionState`, and the unified-execution flag is now a static SQL config. Plan construction and optimizer gating therefore observe the same immutable value. Attempts to change the flag in either direction after session construction fail explicitly with `CANNOT_MODIFY_STATIC_CONFIG`, with focused coverage for both directions. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/SparkOptimizer.scala: ########## @@ -147,7 +154,7 @@ class SparkOptimizer( * batch executing the [[ExperimentalMethods]] optimizer rules. This hook can be used to add * custom optimizer batches to the Spark optimizer. * - * Note that 'Extract Python UDFs' batch is an exception and ran after the batches defined here. + * Note that 'Extract UDFs' batch is an exception and ran after the batches defined here. Review Comment: Fixed. Updated the sentence to use the present tense: “is an exception and runs after”. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala: ########## @@ -2452,6 +2456,7 @@ object PushPredicateThroughNonJoin extends Rule[LogicalPlan] with PredicateHelpe case _: RebalancePartitions => true case _: ScriptTransformation => true case _: Sort => true + case _: ExecuteExternalUDF => true Review Comment: Fixed. The predicate-pushdown regression now verifies both sides: the child-only predicate moves below `ExecuteExternalUDF`, and the UDF-result predicate remains in the upper `Filter`. -- 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]
