andygrove commented on code in PR #4053: URL: https://github.com/apache/datafusion-comet/pull/4053#discussion_r3139425977
########## spark/src/main/scala/org/apache/comet/rules/CometReuseSubquery.scala: ########## @@ -0,0 +1,67 @@ +/* + * 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.comet.rules + +import scala.collection.mutable + +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.catalyst.trees.TreePattern.PLAN_EXPRESSION +import org.apache.spark.sql.execution.{BaseSubqueryExec, ExecSubqueryExpression, ReusedSubqueryExec, SparkPlan} + +/** + * Re-applies subquery deduplication after Comet node conversions. + * + * Spark's ReuseAdaptiveSubquery runs as a queryStageOptimizerRule before postStageCreationRules, + * which is where CometScanRule/CometExecRule replace Spark operators with Comet equivalents. The + * Comet rules copy expressions from the original Spark nodes, which can disrupt subquery reuse + * that was already applied by Spark's rule. This rule runs after Comet conversions to restore + * proper deduplication. + * + * Uses the same algorithm as Spark's ReuseExchangeAndSubquery (subquery portion): top-down + * traversal via transformAllExpressionsWithPruning, caching by canonical form. + * + * For non-AQE, Spark's ReuseExchangeAndSubquery runs after ApplyColumnarRulesAndInsertTransitions + * in QueryExecution.preparations and handles reuse correctly without this rule. + * + * @see + * ReuseExchangeAndSubquery (Spark's non-AQE subquery reuse) + * @see + * ReuseAdaptiveSubquery (Spark's AQE subquery reuse) + */ +case object CometReuseSubquery extends Rule[SparkPlan] { Review Comment: We should start thinking about an audit process for areas where we are replicating Spark logic, so that when we upgrade to future versions of Spark we can check that we staying consistent. Maybe a starting point is adding annotations to sections of code where we replicate upstream logic? No action needed on this PR ... just thinking aloud -- 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]
