From 644e38fa9984b4f9c99d86c7ab5834cf4cded138 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Sat, 22 Aug 2026 12:57:56 -0700 Subject: [PATCH] fix: revert unsafe partial aggregates after final fallback --- .../apache/comet/rules/CometExecRule.scala | 63 ++++++++++- .../comet/exec/CometAggregateSuite.scala | 101 ++++++++++++++++-- .../comet/rules/CometExecRuleSuite.scala | 53 +++++++++ 3 files changed, 203 insertions(+), 14 deletions(-) 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 3c2668cfe92..bce77e9d534 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -614,7 +614,7 @@ case class CometExecRule(session: SparkSession) // during the bottom-up conversion. Tags persist through AQE stage creation. tagUnsafePartialAggregates(planWithJoinRewritten) - var newPlan = transform(planWithJoinRewritten) + var newPlan = revertUnsafePartialAggregates(transform(planWithJoinRewritten)) // if the plan cannot be run fully natively then explain why (when appropriate // config is enabled) @@ -1014,6 +1014,63 @@ case class CometExecRule(session: SparkSession) } } + /** + * The early tagging pass cannot know whether a Final's child will become native. Check the + * actual conversion result as well, before native blocks are serialized or AQE launches stages. + * Restore only the feeding aggregate/exchange chain; keep native work below its Partial. + */ + private def revertUnsafePartialAggregates(plan: SparkPlan): SparkPlan = { + def revertChain(node: SparkPlan): Option[SparkPlan] = node match { + case agg: CometHashAggregateExec if agg.modes == Seq(Partial) => + val partial = agg.originalPlan.withNewChildren(Seq(agg.child)) + partial.setTagValue( + CometExecRule.COMET_UNSAFE_PARTIAL, + "Partial aggregate disabled: corresponding final aggregate " + + "cannot be converted to Comet and intermediate buffer formats are incompatible") + Some(partial) + + case agg: CometHashAggregateExec + if agg.modes.forall(m => m == Partial || m == PartialMerge) => + revertChain(agg.child).map(child => agg.originalPlan.withNewChildren(Seq(child))) + + case agg: BaseAggregateExec + if agg.aggregateExpressions.nonEmpty && + agg.aggregateExpressions.forall(_.mode == Partial) => + // This producer already emits Spark buffers. Do not reach through it to an unrelated + // aggregate below it. + None + + case agg: BaseAggregateExec + if agg.aggregateExpressions.forall(e => e.mode == Partial || e.mode == PartialMerge) => + revertChain(agg.child).map(child => agg.withNewChildren(Seq(child))) + + case CometSinkPlaceHolder(_, _, shuffle: CometShuffleExchangeExec) => + revertChain(shuffle) + case shuffle: CometShuffleExchangeExec => + revertChain(shuffle.child).map(child => shuffle.originalPlan.withNewChildren(Seq(child))) + case shuffle: ShuffleExchangeExec => + revertChain(shuffle.child).map(child => shuffle.withNewChildren(Seq(child))) + + case _: ShuffleQueryStageExec | _: ReusedExchangeExec => + // A stage owns (and may already have materialized) its buffers. Never rewrite it here. + // The whole-plan QueryStagePrep pass must tag the Partial before stages are created; + // that tag keeps it in Spark when the rule is reapplied to the exchange in isolation. + None + case _ => None + } + + plan.transformUp { + case agg: BaseAggregateExec + if agg.aggregateExpressions.map(_.mode).distinct == Seq(Final) && + !QueryPlanSerde.allAggsSupportMixedExecution(agg.aggregateExpressions) => + revertChain(agg.child) + // Rebuild native consumers and shuffles from their original Spark operators. Merely + // replacing their children would leave a native protobuf reading the old buffers. + .map(child => transform(agg.withNewChildren(Seq(child)))) + .getOrElse(agg) + } + } + /** * Look for the bottom Partial-mode aggregate that feeds into the given plan (the child of a * Final). Walks through exchanges and AQE stages, and continues down through intermediate @@ -1043,8 +1100,8 @@ case class CometExecRule(session: SparkSession) /** * Conservative check for whether an aggregate could be converted to Comet. Checks operator * enablement, grouping expressions, aggregate expressions, and result expressions. - * Intentionally skips the sparkFinalMode / child-native checks since those depend on - * transformation state. + * Intentionally skips the child-native checks since those depend on transformation state; + * [[revertUnsafePartialAggregates]] checks the actual conversion result before execution. * * WARNING: this intentionally mirrors the predicate checks in `CometBaseAggregate.doConvert` * (operators.scala). Any change to the convertibility rules there must be reflected here or diff --git a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala index 2aaf3eb8290..635b24464a5 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -25,9 +25,11 @@ import org.apache.hadoop.fs.Path import org.apache.spark.SparkConf import org.apache.spark.sql.{CometTestBase, DataFrame, Row} import org.apache.spark.sql.catalyst.expressions.Cast +import org.apache.spark.sql.catalyst.expressions.aggregate.Partial import org.apache.spark.sql.catalyst.optimizer.EliminateSorts -import org.apache.spark.sql.comet.CometHashAggregateExec -import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.comet.{CometFilterExec, CometHashAggregateExec} +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, ShuffleQueryStageExec} +import org.apache.spark.sql.execution.aggregate.BaseAggregateExec import org.apache.spark.sql.functions.{avg, col, count_distinct, sum} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{DataTypes, StructField, StructType} @@ -35,6 +37,7 @@ import org.apache.spark.sql.types.{DataTypes, StructField, StructType} import org.apache.comet.CometConf import org.apache.comet.CometConf.COMET_EXEC_STRICT_FLOATING_POINT import org.apache.comet.CometSparkSessionExtensions.isSpark41Plus +import org.apache.comet.rules.CometExecRule import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator, ParquetGenerator, SchemaGenOptions} /** @@ -190,6 +193,73 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } + for (adaptive <- Seq(false, true)) { + test(s"decimal AVG falls back across a Spark shuffle (AQE=$adaptive)") { + withTempDir { dir => + val path = s"${dir.getAbsolutePath}/data" + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + spark + .range(0L, 8L, 1L, 4) + .selectExpr("id", "CAST(200 AS DECIMAL(20, 2)) AS v") + .write + .parquet(path) + } + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + SQLConf.FILES_MAX_PARTITION_BYTES.key -> "1048576", + SQLConf.PARQUET_FILTER_PUSHDOWN_ENABLED.key -> "false", + CometConf.COMET_NATIVE_SCAN_ENABLED.key -> "false", + CometConf.COMET_CONVERT_FROM_PARQUET_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_ENABLED.key -> "false") { + withParquetTable(path, "decimal_avg_fallback") { + // The filter leaves three input partitions empty. Decimal AVG is not safe to mix + // between engines: a native empty partial can poison the Spark final's sum buffer. + val df = sql("SELECT AVG(v) FROM decimal_avg_fallback WHERE id = 1") + val initialPlan = stripAQEPlan(df.queryExecution.executedPlan) + checkAnswer(df, Seq(Row(new java.math.BigDecimal("200.000000")))) + for (plan <- Seq(initialPlan, df.queryExecution.executedPlan)) { + assert(collect(plan) { case agg: CometHashAggregateExec => agg }.isEmpty) + val partials = collect(plan) { + case agg: BaseAggregateExec + if agg.aggregateExpressions.forall(_.mode == Partial) => + agg + } + assert(partials.size == 1) + assert(partials.forall(_.getTagValue(CometExecRule.COMET_UNSAFE_PARTIAL).isDefined)) + // Falling back the aggregate must not discard the native filter/scan conversion. + assert(collect(plan) { case filter: CometFilterExec => filter }.nonEmpty) + } + if (adaptive) { + val stages = collect(df.queryExecution.executedPlan) { + case stage: ShuffleQueryStageExec => stage + } + assert(stages.nonEmpty && stages.forall(_.isMaterialized)) + } + + // Compatible buffers may still use a native Partial and a Spark Final. + val safe = sql("SELECT MIN(v), MAX(v) FROM decimal_avg_fallback WHERE id = 1") + checkAnswer( + safe, + Seq(Row(new java.math.BigDecimal("200.00"), new java.math.BigDecimal("200.00")))) + assert(collect(safe.queryExecution.executedPlan) { case agg: CometHashAggregateExec => + agg + }.size == 1) + + // The same unsafe buffer is valid when both aggregate stages execute in Comet. + withSQLConf(CometConf.COMET_SHUFFLE_ENABLED.key -> "true") { + val native = sql("SELECT AVG(v) FROM decimal_avg_fallback WHERE id = 1") + checkAnswer(native, Seq(Row(new java.math.BigDecimal("200.000000")))) + assert(collect(native.queryExecution.executedPlan) { + case agg: CometHashAggregateExec => agg + }.size == 2) + } + } + } + } + } + } + test("stddev_pop should return NaN for some cases") { withSQLConf(CometConf.COMET_SHUFFLE_ENABLED.key -> "true") { Seq(true, false).foreach { nullOnDivideByZero => @@ -513,7 +583,10 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { dictionaryEnabled) { val n = if (nativeShuffleEnabled) 2 else 1 checkSparkAnswerAndNumOfAggregates("SELECT _2, SUM(_1) FROM tbl GROUP BY _2", n) - checkSparkAnswerAndNumOfAggregates("SELECT _2, COUNT(_1) FROM tbl GROUP BY _2", n) + // COUNT is not declared safe for mixed execution, unlike the other aggregates here. + checkSparkAnswerAndNumOfAggregates( + "SELECT _2, COUNT(_1) FROM tbl GROUP BY _2", + if (nativeShuffleEnabled) 2 else 0) checkSparkAnswerAndNumOfAggregates("SELECT _2, MIN(_1) FROM tbl GROUP BY _2", n) checkSparkAnswerAndNumOfAggregates("SELECT _2, MAX(_1) FROM tbl GROUP BY _2", n) checkSparkAnswerAndNumOfAggregates("SELECT _2, AVG(_1) FROM tbl GROUP BY _2", n) @@ -721,26 +794,29 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { val path = new Path(dir.toURI.toString, "test") makeParquetFile(path, 1000, 20, dictionaryEnabled) withParquetTable(path.toUri.toString, "tbl") { + // Spark rewrites _7's small decimal SUM to Long; _8 and _9 remain decimal and + // cannot use a native Partial when the Final runs in Spark. val expectedNumOfCometAggregates = if (nativeShuffleEnabled) 2 else 1 + val expectedNumOfDecimalAggregates = if (nativeShuffleEnabled) 2 else 0 checkSparkAnswerAndNumOfAggregates( "SELECT _g2, SUM(_7) FROM tbl GROUP BY _g2", expectedNumOfCometAggregates) checkSparkAnswerAndNumOfAggregates( "SELECT _g3, SUM(_8) FROM tbl GROUP BY _g3", - expectedNumOfCometAggregates) + expectedNumOfDecimalAggregates) checkSparkAnswerAndNumOfAggregates( "SELECT _g4, SUM(_9) FROM tbl GROUP BY _g4", - expectedNumOfCometAggregates) + expectedNumOfDecimalAggregates) checkSparkAnswerAndNumOfAggregates( "SELECT SUM(_7) FROM tbl", expectedNumOfCometAggregates) checkSparkAnswerAndNumOfAggregates( "SELECT SUM(_8) FROM tbl", - expectedNumOfCometAggregates) + expectedNumOfDecimalAggregates) checkSparkAnswerAndNumOfAggregates( "SELECT SUM(_9) FROM tbl", - expectedNumOfCometAggregates) + expectedNumOfDecimalAggregates) } } } @@ -1325,7 +1401,9 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { val path = new Path(dir.toURI.toString, "test") makeParquetFile(path, 1000, 20, dictionaryEnabled) withParquetTable(path.toUri.toString, "tbl") { + // Only _7 is rewritten to a mixed-safe Long AVG by Spark's decimal optimizer. val expectedNumOfCometAggregates = if (nativeShuffleEnabled) 2 else 1 + val expectedNumOfDecimalAggregates = if (nativeShuffleEnabled) 2 else 0 checkSparkAnswerAndNumOfAggregates( "SELECT _g2, AVG(_7) FROM tbl GROUP BY _g2", @@ -1333,11 +1411,12 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { checkSparkAnswerWithTolerance("SELECT _g3, AVG(_8) FROM tbl GROUP BY _g3") assert(getNumCometHashAggregate( - sql("SELECT _g3, AVG(_8) FROM tbl GROUP BY _g3")) == expectedNumOfCometAggregates) + sql( + "SELECT _g3, AVG(_8) FROM tbl GROUP BY _g3")) == expectedNumOfDecimalAggregates) checkSparkAnswerAndNumOfAggregates( "SELECT _g4, AVG(_9) FROM tbl GROUP BY _g4", - expectedNumOfCometAggregates) + expectedNumOfDecimalAggregates) checkSparkAnswerAndNumOfAggregates( "SELECT AVG(_7) FROM tbl", @@ -1345,11 +1424,11 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { checkSparkAnswerWithTolerance("SELECT AVG(_8) FROM tbl") assert(getNumCometHashAggregate( - sql("SELECT AVG(_8) FROM tbl")) == expectedNumOfCometAggregates) + sql("SELECT AVG(_8) FROM tbl")) == expectedNumOfDecimalAggregates) checkSparkAnswerAndNumOfAggregates( "SELECT AVG(_9) FROM tbl", - expectedNumOfCometAggregates) + expectedNumOfDecimalAggregates) } } } 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 5444a89fa36..908dfcf8dec 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -35,6 +35,7 @@ import org.apache.spark.sql.types.{DataTypes, StructField, StructType} import org.apache.comet.{CometConf, CometExplainInfo} import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus, isSpark42Plus} +import org.apache.comet.CometSparkSessionExtensions.withFallbackReason import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator} /** @@ -421,6 +422,58 @@ class CometExecRuleSuite extends CometTestBase { } } + for (distinct <- Seq(false, true)) { + test( + s"unsafe aggregate buffers fall back when native shuffle is ineligible (distinct=$distinct)") { + withTempView("test_data") { + createTestDataFrame.createOrReplaceTempView("test_data") + val aggregates = "AVG(CAST(id AS DECIMAL(20, 2)))" + + (if (distinct) ", COUNT(DISTINCT name)" else "") + + for (fallback <- Seq("disabled hash partitioning", "prior shuffle fallback", "none")) { + withSQLConf( + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "native", + CometConf.COMET_SHUFFLE_NATIVE_HASH_PARTITIONING_ENABLED.key -> + (fallback != "disabled hash partitioning").toString) { + val sparkPlan = + createSparkPlan(spark, s"SELECT $aggregates FROM test_data GROUP BY (id % 3)") + val aggregateCount = countOperators(sparkPlan, classOf[HashAggregateExec]) + assert(aggregateCount == (if (distinct) 4 else 2)) + if (fallback == "prior shuffle fallback") { + foreach(sparkPlan) { + case shuffle: ShuffleExchangeExec => + withFallbackReason(shuffle, "prior shuffle fallback") + case _ => + } + } + val transformed = applyCometExecRule(sparkPlan) + + // Shuffle is enabled, but a native-only shuffle can still fall back. The distinct + // rewrite also has intermediate PartialMerge and mixed Partial/PartialMerge stages. + val nativeExpected = fallback == "none" + for (plan <- Seq(transformed, applyCometExecRule(transformed))) { + assert( + countOperators(plan, classOf[CometHashAggregateExec]) == + (if (nativeExpected) aggregateCount else 0)) + assert( + countOperators(plan, classOf[HashAggregateExec]) == + (if (nativeExpected) 0 else aggregateCount)) + } + // AQE reapplies the rule to an exchange without its Final aggregate. The tagged + // Partial must remain in Spark in that stage-only pass too. + transformed.collect { case shuffle: ShuffleExchangeExec => shuffle }.foreach { + shuffle => + val stage = applyCometExecRule(shuffle) + assert(countOperators(stage, classOf[CometHashAggregateExec]) == 0) + } + } + } + } + } + } + test("CometExecRule should not allow decimal SUM mixed execution") { withTempView("test_data") { createTestDataFrame.createOrReplaceTempView("test_data")