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..7c14d1ae4d 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -616,7 +616,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) @@ -974,7 +974,7 @@ case class CometExecRule(session: SparkSession) // PartialMerge stages of a distinct-aggregate rewrite. See issues #1389 and #4813. val modes = agg.aggregateExpressions.map(_.mode).distinct if (modes == Seq(Final) && - !QueryPlanSerde.allAggsSupportMixedExecution(agg.aggregateExpressions) && + !QueryPlanSerde.allAggsSupportNativePartialToSparkFinal(agg.aggregateExpressions) && !canAggregateBeConverted(agg, Final)) { findPartialAggInPlan(agg.child).foreach { partial => // Only tag if the Partial would otherwise have been converted. If the Partial itself @@ -1016,6 +1016,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.allAggsSupportNativePartialToSparkFinal(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 @@ -1045,8 +1102,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/main/scala/org/apache/comet/serde/CometAggregateExpressionSerde.scala b/spark/src/main/scala/org/apache/comet/serde/CometAggregateExpressionSerde.scala index a52d600821..091b57bf19 100644 --- a/spark/src/main/scala/org/apache/comet/serde/CometAggregateExpressionSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/CometAggregateExpressionSerde.scala @@ -82,16 +82,23 @@ trait CometAggregateExpressionSerde[T <: AggregateFunction] { def getSupportLevel(expr: T): SupportLevel = Compatible(None) /** - * Whether this aggregate's intermediate buffer format is compatible between Spark and Comet for - * the given function instance, making it safe to run the Partial in one engine and the Final in - * the other. Aggregates with simple single-value buffers (MIN, MAX, bitwise) are always safe; - * SUM and non-decimal AVG match Spark's buffer and are safe except where noted per instance - * (e.g. TRY-mode SUM uses a Comet-internal flag column). COUNT is intentionally excluded - * despite a matching buffer: mixed COUNT partial/final regressed AQE's + * Whether a Comet aggregate can consume this function's Spark intermediate buffer. This covers + * Spark Partial to Comet Final, including intermediate PartialMerge stages. COUNT is excluded + * despite a matching buffer: a Comet Final above a Spark Partial regressed AQE's * PropagateEmptyRelationAfterAQE pattern (which matches BaseAggregateExec only) and the Spark * 4.0 count-bug decorrelation for correlated IN subqueries. */ - def supportsMixedPartialFinal(fn: T): Boolean = false + def supportsSparkPartialToNativeFinal(fn: T): Boolean = false + + /** + * Whether Spark can consume this function's Comet intermediate buffer. Keep this separate from + * the reverse direction: planner restrictions on a Comet Final do not necessarily prohibit a + * Comet Partial feeding a Spark Final. Existing bidirectional implementations share the + * default; handlers may admit an additional forward direction only after validating it + * independently. + */ + def supportsNativePartialToSparkFinal(fn: T): Boolean = + supportsSparkPartialToNativeFinal(fn) /** * Convert a Spark expression into a protocol buffer representation that can be passed into diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala index 6802dfaa64..236cbd8377 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala @@ -421,31 +421,35 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { classOf[VarianceSamp] -> CometVarianceSamp) /** - * Returns true if all aggregate expressions in the list have intermediate buffer formats that - * are compatible between Spark and Comet, making it safe to run Partial in one engine and Final - * in the other. + * Returns true if Spark can consume all the intermediate buffers produced by Comet. Used when a + * Spark Final would otherwise consume a native Partial, including after shuffle fallback. */ - def allAggsSupportMixedExecution(aggExprs: Seq[AggregateExpression]): Boolean = { - aggExprs.forall(aggExpr => supportsMixedExecution(aggExpr.aggregateFunction)) + def allAggsSupportNativePartialToSparkFinal(aggExprs: Seq[AggregateExpression]): Boolean = { + aggExprs.forall { aggExpr => + val fn = aggExpr.aggregateFunction + aggrSerdeMap.get(fn.getClass).exists { handler => + handler + .asInstanceOf[CometAggregateExpressionSerde[AggregateFunction]] + .supportsNativePartialToSparkFinal(fn) + } + } } /** - * Returns the aggregate functions in the list whose intermediate buffer formats are not known - * to be compatible between Spark and Comet. These are the functions that prevent a Spark Final - * aggregate (without a Comet Partial) from running, since the buffer produced by one engine - * cannot be safely consumed by the other. + * Returns functions whose Spark intermediate buffers cannot safely be consumed by a Comet Final + * or PartialMerge. This is independent of native Partial to Spark Final compatibility. */ - def aggsNotSupportingMixedExecution( + def aggsNotSupportingSparkPartialToNativeFinal( aggExprs: Seq[AggregateExpression]): Seq[AggregateFunction] = { - aggExprs.map(_.aggregateFunction).filterNot(supportsMixedExecution) + aggExprs.map(_.aggregateFunction).filterNot(supportsSparkPartialToNativeFinal) } - private def supportsMixedExecution(fn: AggregateFunction): Boolean = { + private def supportsSparkPartialToNativeFinal(fn: AggregateFunction): Boolean = { aggrSerdeMap.get(fn.getClass) match { case Some(handler) => handler .asInstanceOf[CometAggregateExpressionSerde[AggregateFunction]] - .supportsMixedPartialFinal(fn) + .supportsSparkPartialToNativeFinal(fn) case None => false } } diff --git a/spark/src/main/scala/org/apache/comet/serde/aggregates.scala b/spark/src/main/scala/org/apache/comet/serde/aggregates.scala index db435e5b5d..7eaaa2b04a 100644 --- a/spark/src/main/scala/org/apache/comet/serde/aggregates.scala +++ b/spark/src/main/scala/org/apache/comet/serde/aggregates.scala @@ -35,7 +35,7 @@ import org.apache.comet.shims.{CometCollectShim, CometEvalModeUtil} object CometMin extends CometAggregateExpressionSerde[Min] { - override def supportsMixedPartialFinal(fn: Min): Boolean = true + override def supportsSparkPartialToNativeFinal(fn: Min): Boolean = true override def getSupportLevel(expr: Min): SupportLevel = AggSerde.minMaxSupportLevel(expr.dataType) @@ -71,7 +71,7 @@ object CometMin extends CometAggregateExpressionSerde[Min] { object CometMax extends CometAggregateExpressionSerde[Max] { - override def supportsMixedPartialFinal(fn: Max): Boolean = true + override def supportsSparkPartialToNativeFinal(fn: Max): Boolean = true override def getSupportLevel(expr: Max): SupportLevel = AggSerde.minMaxSupportLevel(expr.dataType) @@ -106,6 +106,10 @@ object CometMax extends CometAggregateExpressionSerde[Max] { } object CometCount extends CometAggregateExpressionSerde[Count] { + // Both buffers are a single non-null Long. The AQE/count-bug restrictions documented on the + // reverse direction concern a Comet Final; retaining Spark's Final preserves those rewrites. + override def supportsNativePartialToSparkFinal(fn: Count): Boolean = true + override def convert( aggExpr: AggregateExpression, expr: Count, @@ -129,7 +133,7 @@ object CometCount extends CometAggregateExpressionSerde[Count] { object CometAverage extends CometAggregateExpressionSerde[Average] { - override def supportsMixedPartialFinal(fn: Average): Boolean = + override def supportsSparkPartialToNativeFinal(fn: Average): Boolean = // Non-decimal AVG has a (sum: double, count: long) buffer matching Spark. Decimal AVG is // deferred (overflow nulls count differently) and stays unsafe for mixed execution. !fn.child.dataType.isInstanceOf[DecimalType] @@ -189,7 +193,7 @@ object CometAverage extends CometAggregateExpressionSerde[Average] { object CometSum extends CometAggregateExpressionSerde[Sum] { - override def supportsMixedPartialFinal(fn: Sum): Boolean = + override def supportsSparkPartialToNativeFinal(fn: Sum): Boolean = // Decimal SUM is excluded: overflow detection (ANSI throw / Legacy null) does not survive a // Spark-partial / Comet-final split, so the required ArithmeticException is never raised. // TRY-mode integer SUM carries a Comet-internal has_all_nulls column that Spark cannot read. @@ -306,7 +310,7 @@ object CometLast extends CometAggregateExpressionSerde[Last] { } object CometBitAndAgg extends CometAggregateExpressionSerde[BitAndAgg] { - override def supportsMixedPartialFinal(fn: BitAndAgg): Boolean = true + override def supportsSparkPartialToNativeFinal(fn: BitAndAgg): Boolean = true override def getSupportLevel(expr: BitAndAgg): SupportLevel = if (AggSerde.bitwiseAggTypeSupported(expr.dataType)) { @@ -344,7 +348,7 @@ object CometBitAndAgg extends CometAggregateExpressionSerde[BitAndAgg] { } object CometBitOrAgg extends CometAggregateExpressionSerde[BitOrAgg] { - override def supportsMixedPartialFinal(fn: BitOrAgg): Boolean = true + override def supportsSparkPartialToNativeFinal(fn: BitOrAgg): Boolean = true override def getSupportLevel(expr: BitOrAgg): SupportLevel = if (AggSerde.bitwiseAggTypeSupported(expr.dataType)) { @@ -382,7 +386,7 @@ object CometBitOrAgg extends CometAggregateExpressionSerde[BitOrAgg] { } object CometBitXOrAgg extends CometAggregateExpressionSerde[BitXorAgg] { - override def supportsMixedPartialFinal(fn: BitXorAgg): Boolean = true + override def supportsSparkPartialToNativeFinal(fn: BitXorAgg): Boolean = true override def getSupportLevel(expr: BitXorAgg): SupportLevel = if (AggSerde.bitwiseAggTypeSupported(expr.dataType)) { @@ -765,7 +769,7 @@ object CometCorr extends CometAggregateExpressionSerde[Corr] { object CometBloomFilterAggregate extends CometAggregateExpressionSerde[BloomFilterAggregate] { - override def supportsMixedPartialFinal(fn: BloomFilterAggregate): Boolean = true + override def supportsSparkPartialToNativeFinal(fn: BloomFilterAggregate): Boolean = true override def getSupportLevel(expr: BloomFilterAggregate): SupportLevel = expr.child.dataType match { @@ -934,7 +938,7 @@ object CometApproxCountDistinct extends CometAggregateExpressionSerde[HyperLogLo // The register buffer uses Spark's identical packed-`Long` layout (`numWords` `Long` columns), // matching Spark's `aggBufferSchema`, so a Comet partial and Spark final (or the reverse) can // be mixed in one plan. - override def supportsMixedPartialFinal(fn: HyperLogLogPlusPlus): Boolean = true + override def supportsSparkPartialToNativeFinal(fn: HyperLogLogPlusPlus): Boolean = true // Types that Comet's native `xxhash64` hashes identically to Spark's `XxHash64Function`. // `StringType` here is the default UTF8_BINARY collation; a collated `StringType(collationId)` diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala index 216f9f1e4b..f39f09b605 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala @@ -1656,7 +1656,7 @@ trait CometBaseAggregate { if (missingCometProducer) { val incompatibleAggs = - QueryPlanSerde.aggsNotSupportingMixedExecution(aggregate.aggregateExpressions) + QueryPlanSerde.aggsNotSupportingSparkPartialToNativeFinal(aggregate.aggregateExpressions) if (incompatibleAggs.nonEmpty) { val names = incompatibleAggs.map(_.prettyName).distinct.sorted.mkString(", ") withFallbackReason( 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 143248f551..fb763b0b3c 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -31,10 +31,11 @@ import org.apache.spark.sql.catalyst.expressions.Cast import org.apache.spark.sql.catalyst.expressions.aggregate.{Final, Partial} import org.apache.spark.sql.catalyst.optimizer.EliminateSorts import org.apache.spark.sql.catalyst.plans.physical.RangePartitioning -import org.apache.spark.sql.comet.CometHashAggregateExec +import org.apache.spark.sql.comet.{CometFilterExec, CometHashAggregateExec} import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec import org.apache.spark.sql.execution.SQLExecution -import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +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} @@ -42,6 +43,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} /** @@ -299,6 +301,145 @@ 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) + } + } + } + } + } + } + + for (adaptive <- Seq(false, true)) { + test(s"COUNT preserves safe native partials across a Spark shuffle (AQE=$adaptive)") { + val data = Seq((0, None), (0, None), (1, Some(3)), (1, None), (1, Some(4))) + withParquetTable(data, "count_fallback", false) { + for (finalEnabled <- Seq(false, true)) { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + SQLConf.PARQUET_FILTER_PUSHDOWN_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_ENABLED.key -> "false", + CometConf.COMET_ENABLE_FINAL_HASH_AGGREGATE.key -> finalEnabled.toString) { + for (query <- Seq( + "SELECT _1, COUNT(_2), COUNT(*) FROM count_fallback GROUP BY _1", + "SELECT COUNT(_2), COUNT(*) FROM count_fallback WHERE _1 = 0", + "SELECT COUNT(_2), COUNT(*) FROM count_fallback WHERE _1 < 0")) { + val df = sql(query) + val initialPlan = stripAQEPlan(df.queryExecution.executedPlan) + assert(collect(initialPlan) { + case agg: CometHashAggregateExec if agg.modes == Seq(Partial) => agg + }.size == 1) + assert(collect(initialPlan) { + case agg: BaseAggregateExec + if agg.aggregateExpressions.map(_.mode).distinct == Seq(Final) => + agg + }.size == 1) + checkSparkAnswer(df) + } + } + } + } + } + + for (fn <- Seq("collect_list", "collect_set")) { + test(s"$fn falls back when enabled native shuffle is ineligible (AQE=$adaptive)") { + val data = (0 until 30).map(i => (i % 3, if (i % 7 == 0) None else Some(i % 5))) + withParquetTable(data, "collect_fallback", false) { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + SQLConf.PARQUET_FILTER_PUSHDOWN_ENABLED.key -> "false", + SQLConf.USE_OBJECT_HASH_AGG.key -> "true", + CometConf.COMET_SHUFFLE_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "native", + CometConf.COMET_SHUFFLE_NATIVE_HASH_PARTITIONING_ENABLED.key -> "false") { + // Integer keys isolate this from the wide-decimal shuffle restriction in #5420. + // The native Partial emits an Array buffer, but Spark's Final expects Binary. + val query = s"SELECT _1, sort_array($fn(_2)), COUNT(*) " + + "FROM collect_fallback WHERE _1 >= 0 GROUP BY _1" + val df = sql(query) + val initialPlan = stripAQEPlan(df.queryExecution.executedPlan) + checkSparkAnswer(df) + 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.map(_.mode).distinct == Seq(Partial) => + agg + } + assert(partials.size == 1) + assert(partials.head.getTagValue(CometExecRule.COMET_UNSAFE_PARTIAL).isDefined) + assert(collect(plan) { case filter: CometFilterExec => filter }.nonEmpty) + } + // A fully native producer/consumer pair can still use its native buffer format. + withSQLConf(CometConf.COMET_SHUFFLE_NATIVE_HASH_PARTITIONING_ENABLED.key -> "true") { + val native = sql(query) + checkSparkAnswer(native) + assert(getNumCometHashAggregate(native) == 2) + } + } + } + } + } + } + test("stddev_pop should return NaN for some cases") { withSQLConf(CometConf.COMET_SHUFFLE_ENABLED.key -> "true") { Seq(true, false).foreach { nullOnDivideByZero => @@ -857,26 +998,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) } } } @@ -1461,7 +1605,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", @@ -1469,11 +1615,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", @@ -1481,11 +1628,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 5444a89fa3..38df4add64 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} /** @@ -233,8 +234,7 @@ class CometExecRuleSuite extends CometTestBase { } } - // Regression test for https://github.com/apache/datafusion-comet/issues/1389 - test("CometExecRule should not allow Comet partial and Spark final hash aggregate") { + test("CometExecRule should allow COUNT Comet partial and Spark final hash aggregate") { withTempView("test_data") { createTestDataFrame.createOrReplaceTempView("test_data") @@ -250,11 +250,10 @@ class CometExecRuleSuite extends CometTestBase { CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { val transformedPlan = applyCometExecRule(sparkPlan) - // COUNT is intentionally excluded from mixed execution (AQE / count-bug reasons), so if - // the final aggregate cannot be converted to Comet, neither should the partial. - assert( - countOperators(transformedPlan, classOf[HashAggregateExec]) == originalHashAggCount) - assert(countOperators(transformedPlan, classOf[CometHashAggregateExec]) == 0) + // COUNT's buffer is compatible in this direction. Keeping the Final in Spark also keeps + // the AQE/count-bug rewrites that prevent the reverse direction from being admitted. + assert(countOperators(transformedPlan, classOf[HashAggregateExec]) == 1) + assert(countOperators(transformedPlan, classOf[CometHashAggregateExec]) == 1) } } } @@ -275,8 +274,8 @@ class CometExecRuleSuite extends CometTestBase { CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { val transformedPlan = applyCometExecRule(sparkPlan) - // COUNT blocks mixed execution, so if the partial cannot be converted, neither should - // the final. + // COUNT still blocks Spark Partial to Comet Final, independently of the safe reverse + // direction, so if the partial cannot be converted, neither should the final. assert( countOperators(transformedPlan, classOf[HashAggregateExec]) == originalHashAggCount) assert(countOperators(transformedPlan, classOf[CometHashAggregateExec]) == 0) @@ -421,6 +420,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")