diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index c69602fc80..adc86f7b2a 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -32,7 +32,7 @@ import org.apache.spark.sql.comet._ import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometNativeShuffle, CometShuffleExchangeExec} import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.execution._ -import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, BroadcastQueryStageExec, ShuffleQueryStageExec} +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, BroadcastQueryStageExec, LogicalQueryStage, ShuffleQueryStageExec} import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec, ObjectHashAggregateExec} import org.apache.spark.sql.execution.command.{DataWritingCommandExec, ExecutedCommandExec} import org.apache.spark.sql.execution.datasources.WriteFilesExec @@ -646,6 +646,15 @@ case class CometExecRule(session: SparkSession) // Set up logical links newPlan = newPlan.transform { + case op: CometExec + if op + .getTagValue(SparkPlan.LOGICAL_PLAN_TAG) + .exists(_.isInstanceOf[LogicalQueryStage]) => + // AQE replanning reuses this physical root and links it to the current logical stage. + // originalPlan can still point to a subtree hidden inside that logical leaf, which + // AQE cannot replace in the current logical plan. Only preserve a direct stage link, + // not a link inherited from an ancestor. + op case op: CometExec => if (op.originalPlan.logicalLink.isEmpty) { op.unsetTagValue(SparkPlan.LOGICAL_PLAN_TAG) diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala index 4eb6d00178..c117e13327 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -32,12 +32,12 @@ import org.apache.spark.sql._ import org.apache.spark.sql.catalyst.{FunctionIdentifier, TableIdentifier} import org.apache.spark.sql.catalyst.catalog.{BucketSpec, CatalogStatistics, CatalogTable} import org.apache.spark.sql.catalyst.expressions.{DynamicPruningExpression, Expression, ExpressionInfo, Hex, Literal} -import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateMode, BloomFilterAggregate} +import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateMode, BloomFilterAggregate, Final} import org.apache.spark.sql.comet._ import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometShuffleExchangeExec} import org.apache.spark.sql.connector.catalog.InMemoryTableCatalog import org.apache.spark.sql.execution._ -import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, BroadcastQueryStageExec} +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, BroadcastQueryStageExec, LogicalQueryStage} import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, BroadcastExchangeLike, ReusedExchangeExec, ShuffleExchangeExec} import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, CartesianProductExec, SortMergeJoinExec} @@ -2110,6 +2110,54 @@ class CometExecSuite extends CometTestBase { } } + test("AQE broadcasts native aggregates after replanning") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10485760", + SQLConf.SHUFFLE_PARTITIONS.key -> "4", + CometConf.COMET_SHUFFLE_MODE.key -> "native", + CometConf.COMET_SPARK_TO_ARROW_SUPPORTED_OPERATOR_LIST.key -> "Range") { + val df = sql(""" + |WITH s AS ( + | SELECT id % 64 AS k, SUM(id) AS v FROM range(0, 4096, 1, 4) GROUP BY id % 64 + |), r1 AS ( + | SELECT id % 64 AS k, SUM(id + 1) AS v FROM range(0, 3072, 1, 4) GROUP BY id % 64 + |), r2 AS ( + | SELECT id % 64 AS k, SUM(id + 7) AS v FROM range(0, 2048, 1, 4) GROUP BY id % 64 + |), g AS ( + | SELECT SUM(id) AS v FROM range(0, 1024, 1, 4) + |) + |SELECT SUM(s.v + COALESCE(r1.v, 0) + COALESCE(r2.v, 0) + g.v) + |FROM s LEFT JOIN r1 ON s.k = r1.k LEFT JOIN r2 ON s.k = r2.k CROSS JOIN g + |""".stripMargin) + val adaptive = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec] + assert(collect(adaptive.executedPlan) { case b: CometBroadcastHashJoinExec => b }.isEmpty) + + checkAnswer(df, Seq(Row(48738816L))) + + val finalPlan = adaptive.executedPlan + assert(collect(finalPlan) { case b: CometBroadcastHashJoinExec => b }.size == 2) + val broadcasts = collect(finalPlan) { case b: CometBroadcastExchangeExec => b } + val aggregates = broadcasts.flatMap { broadcast => + collect(broadcast.child) { + case a: CometHashAggregateExec + if a.modes.contains(Final) && a.groupingExpressions.nonEmpty => + a + } + } + assert(aggregates.size == 2) + aggregates.foreach { aggregate => + assert(aggregate.longMetric("output_rows").value == 64) + assert(aggregate.longMetric("elapsed_compute").value > 0) + assert( + aggregate + .getTagValue(SparkPlan.LOGICAL_PLAN_TAG) + .exists(_.isInstanceOf[LogicalQueryStage])) + } + } + } + test("CometShuffleExchangeExec logical link should be correct") { withTempView("v") { spark.sparkContext diff --git a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala index 5444a89fa3..8dd370700c 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -25,12 +25,14 @@ import org.apache.spark.sql._ import org.apache.spark.sql.catalyst.FunctionIdentifier import org.apache.spark.sql.catalyst.expressions.{Expression, ExpressionInfo} import org.apache.spark.sql.catalyst.expressions.aggregate.BloomFilterAggregate +import org.apache.spark.sql.catalyst.plans.logical.LocalRelation import org.apache.spark.sql.comet._ import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec import org.apache.spark.sql.execution._ -import org.apache.spark.sql.execution.adaptive.QueryStageExec +import org.apache.spark.sql.execution.adaptive.{LogicalQueryStage, QueryStageExec, ShuffleQueryStageExec} import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAggregateExec} import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ShuffleExchangeExec} +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{DataTypes, StructField, StructType} import org.apache.comet.{CometConf, CometExplainInfo} @@ -77,6 +79,85 @@ class CometExecRuleSuite extends CometTestBase { }.sum } + /** A native final aggregate over a shuffle stage, as reused by AQE replanning. */ + private def createAdaptiveAggregate(): CometHashAggregateExec = { + val plan = createSparkPlan( + spark, + "SELECT id % 3 AS k, SUM(id) AS total FROM range(0, 100, 1, 2) GROUP BY id % 3") + val aggregate = applyCometExecRule(plan).asInstanceOf[CometHashAggregateExec] + val shuffle = aggregate.child.asInstanceOf[CometShuffleExchangeExec] + aggregate + .withNewChildren(Seq(ShuffleQueryStageExec(0, shuffle, shuffle.canonicalized))) + .asInstanceOf[CometHashAggregateExec] + } + + test("CometExecRule preserves the current direct AQE logical link") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false", + CometConf.COMET_SPARK_TO_ARROW_SUPPORTED_OPERATOR_LIST.key -> "Range") { + val originalTags = + Seq(Some(SparkPlan.LOGICAL_PLAN_TAG), Some(SparkPlan.LOGICAL_PLAN_INHERITED_TAG), None) + originalTags.foreach { originalTag => + withClue(s"original logical tag: $originalTag") { + val aggregate = createAdaptiveAggregate() + val original = aggregate.originalPlan + val originalLogicalPlan = original.logicalLink.get + original.unsetTagValue(SparkPlan.LOGICAL_PLAN_TAG) + original.unsetTagValue(SparkPlan.LOGICAL_PLAN_INHERITED_TAG) + originalTag.foreach(original.setTagValue(_, originalLogicalPlan)) + + var current: SparkPlan = aggregate + (1 to 2).foreach { _ => + val logicalStage = LogicalQueryStage(originalLogicalPlan, current) + val replanned = spark.sessionState.planner.plan(logicalStage).next() + assert(replanned eq current) + assert(replanned.getTagValue(SparkPlan.LOGICAL_PLAN_TAG).exists(_ eq logicalStage)) + + current = applyCometExecRule(replanned) + assert(current.getTagValue(SparkPlan.LOGICAL_PLAN_TAG).exists(_ eq logicalStage)) + } + } + } + } + } + + test("CometExecRule repairs ordinary and inherited logical links from the original plan") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false", + CometConf.COMET_SPARK_TO_ARROW_SUPPORTED_OPERATOR_LIST.key -> "Range") { + val originalTags = + Seq(Some(SparkPlan.LOGICAL_PLAN_TAG), Some(SparkPlan.LOGICAL_PLAN_INHERITED_TAG), None) + for (originalTag <- originalTags; hasDirectLink <- Seq(false, true)) { + withClue(s"original logical tag: $originalTag, ordinary direct link: $hasDirectLink") { + val aggregate = createAdaptiveAggregate() + val original = aggregate.originalPlan + val originalLogicalPlan = original.logicalLink.get + original.unsetTagValue(SparkPlan.LOGICAL_PLAN_TAG) + original.unsetTagValue(SparkPlan.LOGICAL_PLAN_INHERITED_TAG) + originalTag.foreach(original.setTagValue(_, originalLogicalPlan)) + + aggregate.unsetTagValue(SparkPlan.LOGICAL_PLAN_TAG) + aggregate.setTagValue( + SparkPlan.LOGICAL_PLAN_INHERITED_TAG, + LogicalQueryStage(originalLogicalPlan, aggregate)) + if (hasDirectLink) { + aggregate.setTagValue(SparkPlan.LOGICAL_PLAN_TAG, LocalRelation(aggregate.output)) + } + + val transformed = applyCometExecRule(aggregate) + if (originalTag.isDefined) { + assert(transformed.logicalLink.exists(_ eq originalLogicalPlan)) + } else { + assert(transformed.getTagValue(SparkPlan.LOGICAL_PLAN_TAG).isEmpty) + assert(transformed.getTagValue(SparkPlan.LOGICAL_PLAN_INHERITED_TAG).isEmpty) + } + } + } + } + } + test("expression-level fallback reasons are rolled up onto the operator that falls back") { // Extended explain only walks plan nodes, so a reason recorded on a sub-expression is // invisible unless CometExecRule lifts it onto the enclosing operator. Disabling a single