This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-5421-a51ff5cd886328c91bfb8155bca12cb0192ec83c in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit 6065705c16340c0be293212a71decfd9df4daae4 Author: Chao Sun <[email protected]> AuthorDate: Sun Sep 20 23:19:33 2026 +0000 fix: revert unsafe partial aggregates after final fallback (#5421) * fix: revert unsafe partial aggregates after final fallback * fix: distinguish aggregate buffer compatibility by direction * fix: block unsafe native AVG partials before Spark final * test: cover aggregate fallback for unsupported array hash keys * test: align aggregate fallback checks across Spark versions Capture the Spark baseline inside withSQLConf so Spark 3 does not return Unit in place of the expected rows. Expect AVG to remain in Spark after Celeborn shuffle fallback, consistent with its unsafe empty partial buffer, and retain COUNT as the compatible native-partial control. * fix: clarify aggregate buffer compatibility and fallback * ci: group aggregate policy tests with aggregate suite * test: cover DISTINCT aggregate fallback boundaries Cover nondecimal AVG with SUM(DISTINCT) across global and grouped plans, shuffle modes, and selective lower-exchange fallback. Document why native-only fallback reaches the Final repair and qualify the auto/JVM boundary without changing the production predicate. * refactor: simplify aggregate fallback repair and regression tests * fix: address aggregate fallback review follow-ups Restore conditional diagnostics for known unrepaired buffer paths and focused coverage for direction defaults, immediate fallback reasons, and sticky shuffle refusal. Document the decimal SUM precision-overflow mismatch, preserve its fallback with a cancellation regression, and track remaining audits in #5975. --- .../org/apache/comet/rules/CometExecRule.scala | 121 ++++++- .../RevertNativeForTransitionHeavyStages.scala | 6 +- .../serde/CometAggregateExpressionSerde.scala | 19 +- .../org/apache/comet/serde/QueryPlanSerde.scala | 30 +- .../scala/org/apache/comet/serde/aggregates.scala | 57 +++- .../org/apache/spark/sql/comet/operators.scala | 2 +- .../apache/comet/exec/CometAggregateSuite.scala | 367 ++++++++++++++++++++- .../apache/comet/rules/CometExecRuleSuite.scala | 209 ++++++++++-- .../CometCelebornShufflePlanningSuite.scala | 21 +- 9 files changed, 749 insertions(+), 83 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 265c84a0df..54f8891624 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -220,6 +220,16 @@ case class CometExecRule(session: SparkSession) private def isCometNative(op: SparkPlan): Boolean = op.isInstanceOf[CometNativeExec] + /** + * Restore a Spark Partial while retaining its current children. The tag prevents reconversion + * when AQE replans the exchange without its Final, and records why the Partial stays in Spark. + */ + private def restoreSparkPartial(agg: CometHashAggregateExec, reason: String): SparkPlan = { + val partial = agg.originalPlan.withNewChildren(agg.children) + partial.setTagValue(CometExecRule.COMET_UNSAFE_PARTIAL, reason) + withFallbackReason(partial, reason) + } + /** * A Celeborn exchange can fall back after its child has been converted, for example because of * the partition threshold or an unsupported hash key. Keep incompatible partial aggregate @@ -237,10 +247,8 @@ case class CometExecRule(session: SparkSession) case _: QueryStageExec | _: ShuffleExchangeLike | _: BroadcastExchangeLike => plan case agg: CometHashAggregateExec if agg.modes == Seq(Partial) && - !QueryPlanSerde.allAggsSupportMixedExecution(agg.aggregateExpressions) => - val sparkAggregate = agg.originalPlan.withNewChildren(agg.children) - sparkAggregate.setTagValue(CometExecRule.COMET_UNSAFE_PARTIAL, reason) - withFallbackReason(sparkAggregate, reason) + !QueryPlanSerde.allAggsSupportNativePartialToSparkFinal(agg.aggregateExpressions) => + restoreSparkPartial(agg, reason) // Final output is ordinary SQL data; any partial below it belongs to another aggregate. case agg: CometHashAggregateExec if agg.modes.contains(Final) => agg case agg: BaseAggregateExec if agg.aggregateExpressions.exists(_.mode == Final) => agg @@ -758,7 +766,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) @@ -1130,7 +1138,7 @@ case class CometExecRule(session: SparkSession) val consumerMode: AggregateMode = if (modes.contains(PartialMerge)) PartialMerge else Final if (consumesBuffers && - !QueryPlanSerde.allAggsSupportMixedExecution(agg.aggregateExpressions) && + !QueryPlanSerde.allAggsSupportNativePartialToSparkFinal(agg.aggregateExpressions) && !canAggregateBeConverted(agg, consumerMode)) { findPartialAggInPlan(agg.child).foreach { partial => // Only tag if the Partial would otherwise have been converted. If the Partial itself @@ -1173,6 +1181,103 @@ case class CometExecRule(session: SparkSession) } } + /** + * Inspect a failed repair's buffer path without rewriting it or materializing any stage. Report + * only a native Partial/PartialMerge whose emitted state is not known to be Spark-compatible. + * Spark Partials and completed aggregates establish new buffers, so stop there rather than + * finding an unrelated native producer below them. Only known aggregate and exchange wrappers + * forward the same buffer path; an arbitrary operator is not evidence of a mixed boundary. + */ + private def hasUnrepairedNativeBuffer(plan: SparkPlan): Boolean = plan match { + case agg: CometHashAggregateExec if agg.aggregateExpressions.isEmpty => + hasUnrepairedNativeBuffer(agg.child) + case agg: CometHashAggregateExec => + agg.modes.forall(m => m == Partial || m == PartialMerge) && + !QueryPlanSerde.allAggsSupportNativePartialToSparkFinal(agg.aggregateExpressions) + case agg: BaseAggregateExec + if agg.aggregateExpressions.nonEmpty && + agg.aggregateExpressions.forall(_.mode == Partial) => + false + case agg: BaseAggregateExec => + agg.aggregateExpressions.forall(e => e.mode == Partial || e.mode == PartialMerge) && + hasUnrepairedNativeBuffer(agg.child) + case placeholder: CometSinkPlaceHolder => hasUnrepairedNativeBuffer(placeholder.child) + case read: AQEShuffleReadExec => hasUnrepairedNativeBuffer(read.child) + case stage: ShuffleQueryStageExec => hasUnrepairedNativeBuffer(stage.plan) + case reused: ReusedExchangeExec => hasUnrepairedNativeBuffer(reused.child) + case shuffle: CometShuffleExchangeExec => hasUnrepairedNativeBuffer(shuffle.child) + case shuffle: ShuffleExchangeExec => hasUnrepairedNativeBuffer(shuffle.child) + case _ => false + } + + /** + * The early tagging pass cannot know whether a Final's child will become native. Check the + * actual conversion result before serialization or AQE stage creation, restoring the feeding + * aggregate/exchange chain while keeping native work below its Partial. Return the repaired + * plan, or preserve an unrepairable path and record one warning on its Spark Final if an unsafe + * native producer remains. Existing stages and their buffers are never rewritten by this pass. + */ + private[rules] def revertUnsafePartialAggregates(plan: SparkPlan): SparkPlan = { + def revertChain(node: SparkPlan): Option[SparkPlan] = node match { + case agg: CometHashAggregateExec if agg.modes == Seq(Partial) => + Some( + restoreSparkPartial( + agg, + "Partial aggregate disabled: corresponding final aggregate " + + "cannot be converted to Comet and intermediate buffer formats are incompatible")) + + 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))) + + // Stop at materialized stages and operators outside the feeding aggregate/exchange chain. + 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 { + if (hasUnrepairedNativeBuffer(agg.child)) { + val reason = "Comet could not restore a native intermediate buffer producer " + + "below Spark final aggregate; the remaining buffer may be incompatible" + // AQE can revisit the same consumer. Record the explanation and warn once, + // regardless of whether general fallback logging is enabled. + if (!agg + .getTagValue(CometExplainInfo.FALLBACK_REASONS) + .exists(_.contains(reason))) { + if (!CometConf.COMET_EXPLAIN_FALLBACK_LOG_ENABLED.get()) logWarning(reason) + withFallbackReason(agg, reason) + } + } + 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 @@ -1202,8 +1307,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/rules/RevertNativeForTransitionHeavyStages.scala b/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala index 9a941c28ff..a56e477b03 100644 --- a/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala +++ b/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala @@ -124,7 +124,11 @@ case class RevertNativeForTransitionHeavyStages(session: SparkSession) def visit(plan: SparkPlan): Boolean = plan match { case _ if isStageBoundary(plan) => false case aggregate: CometHashAggregateExec - if !QueryPlanSerde.allAggsSupportMixedExecution(aggregate.aggregateExpressions) => + if !QueryPlanSerde + .allAggsSupportNativePartialToSparkFinal(aggregate.aggregateExpressions) || + QueryPlanSerde + .aggsNotSupportingSparkPartialToNativeFinal(aggregate.aggregateExpressions) + .nonEmpty => val producesBuffer = aggregate.modes.exists(mode => mode == Partial || mode == PartialMerge) val consumesAcrossBoundary = 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..73630a4e8b 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,21 @@ 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. Opt in independently + * from the reverse direction: consuming Spark state does not establish that Comet emits state + * Spark can merge, especially from a never-updated or all-null partial accumulator. Remaining + * forward-compatibility audits are tracked in issue #5975. + */ + def supportsNativePartialToSparkFinal(fn: T): Boolean = false /** * 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 eced2d7fdf..43d65fd418 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala @@ -443,31 +443,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 b3bc5e86b7..ad53cceb31 100644 --- a/spark/src/main/scala/org/apache/comet/serde/aggregates.scala +++ b/spark/src/main/scala/org/apache/comet/serde/aggregates.scala @@ -36,7 +36,10 @@ import org.apache.comet.shims.{CometCollectShim, CometEvalModeUtil, CometTypeShi object CometMin extends CometAggregateExpressionSerde[Min] { - override def supportsMixedPartialFinal(fn: Min): Boolean = true + override def supportsSparkPartialToNativeFinal(fn: Min): Boolean = true + + // Native MIN emits one typed null for empty/all-null input; Spark's least merge ignores it. + override def supportsNativePartialToSparkFinal(fn: Min): Boolean = true override def getSupportLevel(expr: Min): SupportLevel = AggSerde.minMaxSupportLevel(expr.dataType) @@ -72,7 +75,10 @@ object CometMin extends CometAggregateExpressionSerde[Min] { object CometMax extends CometAggregateExpressionSerde[Max] { - override def supportsMixedPartialFinal(fn: Max): Boolean = true + override def supportsSparkPartialToNativeFinal(fn: Max): Boolean = true + + // Native MAX emits one typed null for empty/all-null input; Spark's greatest merge ignores it. + override def supportsNativePartialToSparkFinal(fn: Max): Boolean = true override def getSupportLevel(expr: Max): SupportLevel = AggSerde.minMaxSupportLevel(expr.dataType) @@ -190,6 +196,10 @@ object CometMinBy extends CometMaxMinBy[MinBy] { } 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, @@ -213,7 +223,10 @@ object CometCount extends CometAggregateExpressionSerde[Count] { object CometAverage extends CometAggregateExpressionSerde[Average] { - override def supportsMixedPartialFinal(fn: Average): Boolean = + // Keep the default native-to-Spark restriction until #5420: an untouched native AVG emits + // (null, 0), but Spark's merge needs (0.0, 0). + + 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] @@ -273,7 +286,17 @@ object CometAverage extends CometAggregateExpressionSerde[Average] { object CometSum extends CometAggregateExpressionSerde[Sum] { - override def supportsMixedPartialFinal(fn: Sum): Boolean = + // Non-decimal, non-TRY SUM emits one nullable sum, including null for empty/all-null input; + // Spark's coalesce-based merge accepts it. Decimal SUM has Spark's (sum, isEmpty) layout, + // but native updates make precision overflow sticky (or throw in ANSI mode). Spark's generated + // scalar SUM can recover before emitting its partial: decimal(38,38) inputs 0.6, 0.6, -0.6 + // sum to 0.6. Keep decimal partials in Spark until those update semantics match. Integer TRY + // SUM also remains excluded because its native state contains an extra has_all_nulls column. + override def supportsNativePartialToSparkFinal(fn: Sum): Boolean = + !fn.child.dataType.isInstanceOf[DecimalType] && + CometEvalModeUtil.fromSparkEvalMode(CometEvalModeUtil.sumEvalMode(fn)) != CometEvalMode.TRY + + 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. @@ -390,7 +413,10 @@ object CometLast extends CometAggregateExpressionSerde[Last] { } object CometBitAndAgg extends CometAggregateExpressionSerde[BitAndAgg] { - override def supportsMixedPartialFinal(fn: BitAndAgg): Boolean = true + override def supportsSparkPartialToNativeFinal(fn: BitAndAgg): Boolean = true + + // The single native buffer is null for empty/all-null input; Spark's merge skips nulls. + override def supportsNativePartialToSparkFinal(fn: BitAndAgg): Boolean = true override def getSupportLevel(expr: BitAndAgg): SupportLevel = if (AggSerde.bitwiseAggTypeSupported(expr.dataType)) { @@ -428,7 +454,10 @@ object CometBitAndAgg extends CometAggregateExpressionSerde[BitAndAgg] { } object CometBitOrAgg extends CometAggregateExpressionSerde[BitOrAgg] { - override def supportsMixedPartialFinal(fn: BitOrAgg): Boolean = true + override def supportsSparkPartialToNativeFinal(fn: BitOrAgg): Boolean = true + + // The single native buffer is null for empty/all-null input; Spark's merge skips nulls. + override def supportsNativePartialToSparkFinal(fn: BitOrAgg): Boolean = true override def getSupportLevel(expr: BitOrAgg): SupportLevel = if (AggSerde.bitwiseAggTypeSupported(expr.dataType)) { @@ -466,7 +495,10 @@ object CometBitOrAgg extends CometAggregateExpressionSerde[BitOrAgg] { } object CometBitXOrAgg extends CometAggregateExpressionSerde[BitXorAgg] { - override def supportsMixedPartialFinal(fn: BitXorAgg): Boolean = true + override def supportsSparkPartialToNativeFinal(fn: BitXorAgg): Boolean = true + + // The single native buffer is null for empty/all-null input; Spark's merge skips nulls. + override def supportsNativePartialToSparkFinal(fn: BitXorAgg): Boolean = true override def getSupportLevel(expr: BitXorAgg): SupportLevel = if (AggSerde.bitwiseAggTypeSupported(expr.dataType)) { @@ -1015,7 +1047,11 @@ object CometRegrReplacement object CometBloomFilterAggregate extends CometAggregateExpressionSerde[BloomFilterAggregate] { - override def supportsMixedPartialFinal(fn: BloomFilterAggregate): Boolean = true + override def supportsSparkPartialToNativeFinal(fn: BloomFilterAggregate): Boolean = true + + // Native state is Spark's serialized filter, non-null even for empty/all-null input; only + // the final result may be null, so Spark's deserialize always receives a valid filter. + override def supportsNativePartialToSparkFinal(fn: BloomFilterAggregate): Boolean = true override def getSupportLevel(expr: BloomFilterAggregate): SupportLevel = expr.child.dataType match { @@ -1227,7 +1263,10 @@ 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 + + // Native empty/all-null state contains non-null zero Long words, matching Spark's registers. + override def supportsNativePartialToSparkFinal(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 fe1a2e637a..fa3b5da240 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 @@ -1891,7 +1891,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 7cf5e69d6f..55c2ebdb57 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -30,19 +30,21 @@ import org.apache.spark.sql.{CometTestBase, DataFrame, Row} import org.apache.spark.sql.catalyst.expressions.Cast import org.apache.spark.sql.catalyst.expressions.aggregate.{Final, Partial, PartialMerge} 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.catalyst.plans.physical.{HashPartitioning, RangePartitioning} +import org.apache.spark.sql.comet.{CometFilterExec, CometHashAggregateExec, CometNativeExec, CometProjectExec} import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec import org.apache.spark.sql.execution.SQLExecution -import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AdaptiveSparkPlanHelper} -import org.apache.spark.sql.execution.exchange.ReusedExchangeExec +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AdaptiveSparkPlanHelper, ShuffleQueryStageExec} +import org.apache.spark.sql.execution.aggregate.BaseAggregateExec +import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeExec} import org.apache.spark.sql.functions.{avg, col, count_distinct, expr, sum} import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.types.{DataTypes, StructField, StructType} +import org.apache.spark.sql.types.{ArrayType, 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.serde.RegrSparkVersions import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator, ParquetGenerator, SchemaGenOptions} @@ -356,6 +358,290 @@ 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 and AVG fall back together 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(1 AS BIGINT) 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", + CometConf.COMET_ENABLE_FINAL_HASH_AGGREGATE.key -> "true") { + withParquetTable(path, "count_avg_fallback") { + assert(sql("SELECT * FROM count_avg_fallback").rdd.getNumPartitions == 4) + // A safe COUNT buffer must not admit an unsafe AVG buffer in the same Partial. + // Three partitions have no surviving rows, so AVG has no update_batch call and + // its native state is (null, 0), which poisons Spark Final's sum. Keeping Final + // enabled exercises repair after the shuffle falls back during conversion. + val df = sql("SELECT COUNT(*), AVG(v) FROM count_avg_fallback WHERE id = 1") + val initialPlan = stripAQEPlan(df.queryExecution.executedPlan) + checkAnswer(df, Seq(Row(1L, 1.0))) + 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) + } + withSQLConf( + CometConf.COMET_SHUFFLE_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "native") { + // The native Final can consume its own empty AVG buffers; only the engine split + // is unsafe. Keep the fully native aggregate path enabled. + val native = + sql("SELECT COUNT(*), AVG(v) FROM count_avg_fallback WHERE id = 1") + val initialNativePlan = stripAQEPlan(native.queryExecution.executedPlan) + checkAnswer(native, Seq(Row(1L, 1.0))) + for (plan <- Seq(initialNativePlan, native.queryExecution.executedPlan)) { + assert(collect(plan) { case agg: CometHashAggregateExec => agg }.size == 2) + } + } + } + } + } + } + + 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) { + 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 -> "true") { + for (query <- Seq( + "SELECT _1, COUNT(_2), COUNT(*) FROM count_fallback GROUP BY _1", + "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) + } + } + } + } + } + + for (fn <- Seq("percentile", "collect_list", "sum")) { + test( + s"$fn preserves aggregate buffers with an unsupported array hash key (AQE=$adaptive)") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + SQLConf.SHUFFLE_PARTITIONS.key -> "4", + CometConf.COMET_SHUFFLE_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "native", + CometConf.COMET_SHUFFLE_NATIVE_HASH_PARTITIONING_ENABLED.key -> "true") { + withTempView("array_key_aggregate") { + // The array key itself makes native shuffle ineligible; no feature is disabled. + // https://github.com/apache/datafusion-comet/issues/5419#issuecomment-5464233245 + spark + .range(0, 18, 1, 4) + .selectExpr("id % 3 AS k", "id % 5 AS v") + .createOrReplaceTempView("array_key_aggregate") + val aggregate = if (fn == "percentile") "percentile(v, 0.5)" else s"$fn(v)" + val query = s"SELECT array(k) AS ak, $aggregate " + + "FROM array_key_aggregate GROUP BY array(k)" + + def normalizedRows(df: DataFrame): Seq[Row] = { + df.collect() + .toSeq + .map { row => + // Keep the reported collect_list SQL unchanged, normalizing its order only + // after execution so another expression cannot cause an earlier fallback. + if (fn == "collect_list") { + Row(row.getSeq[Long](0), row.getSeq[Long](1).sorted) + } else { + row + } + } + .sortBy(_.getSeq[Long](0).head) + } + + // Spark 3's withSQLConf returns Unit, so capture the baseline inside its body. + var expected: Seq[Row] = Seq.empty + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + expected = normalizedRows(sql(query)) + } + val df = sql(query) + val initialPlan = stripAQEPlan(df.queryExecution.executedPlan) + // Execute this same DataFrame before inspecting its materialized AQE plan. + assert(normalizedRows(df) == expected) + for (plan <- Seq(initialPlan, df.queryExecution.executedPlan)) { + val exchanges = collect(plan) { case exchange: ShuffleExchangeExec => exchange } + assert(exchanges.size == 1, s"$plan") + assert(exchanges.head.outputPartitioning match { + case HashPartitioning(Seq(key), 4) => key.dataType.isInstanceOf[ArrayType] + case _ => false + }) + assert(collect(plan) { case exchange: CometShuffleExchangeExec => + exchange + }.isEmpty) + val partials = collect(plan) { + case agg: BaseAggregateExec + if agg.aggregateExpressions.map(_.mode).distinct == Seq(Partial) => + agg + } + val finals = collect(plan) { + case agg: BaseAggregateExec + if agg.aggregateExpressions.map(_.mode).distinct == Seq(Final) => + agg + } + assert(finals.size == 1, s"$plan") + val nativeAggregates = collect(plan) { case agg: CometHashAggregateExec => agg } + if (fn == "sum") { + // SUM's Long buffer is safe for Spark's final, so retain its native partial. + assert(nativeAggregates.size == 1, s"$plan") + assert(nativeAggregates.head.modes == Seq(Partial)) + assert(partials.isEmpty, s"$plan") + } else { + assert(nativeAggregates.isEmpty, s"$plan") + assert(partials.size == 1, s"$plan") + assert(partials.head.getTagValue(CometExecRule.COMET_UNSAFE_PARTIAL).isDefined) + assert(collect(partials.head.child) { case project: CometProjectExec => + project + }.nonEmpty) + } + } + if (adaptive) { + val stages = collect(df.queryExecution.executedPlan) { + case stage: ShuffleQueryStageExec => stage + } + assert(stages.nonEmpty && stages.forall(_.isMaterialized)) + } + } + } + } + } + } + test("stddev_pop should return NaN for some cases") { withSQLConf(CometConf.COMET_SHUFFLE_ENABLED.key -> "true") { Seq(true, false).foreach { nullOnDivideByZero => @@ -431,15 +717,56 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } - test("mixed engine sum/avg: Comet partial + Spark final matches Spark") { + test("decimal SUM partial stays in Spark when a later input cancels precision overflow") { + // Keep all three values in one ordered input partition. Generated scalar Spark SUM can + // retain the temporary 1.2 and return 0.6 after cancellation; native decimal SUM instead + // makes that precision overflow sticky, or throws immediately in ANSI mode. A matching + // (sum, isEmpty) buffer schema therefore does not establish forward interoperability. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true", + "spark.sql.files.minPartitionNum" -> "1", + CometConf.COMET_ENABLE_FINAL_HASH_AGGREGATE.key -> "false", + CometConf.COMET_SHUFFLE_ENABLED.key -> "false") { + withTempPath { path => + spark + .range(0, 3, 1, 1) + .selectExpr("CAST(CASE WHEN id < 2 THEN '0.6' ELSE '-0.6' END AS DECIMAL(38,38)) AS v") + .write + .parquet(path.getCanonicalPath) + withParquetTable(path.getCanonicalPath, "decimal_sum_cancellation") { + for (ansi <- Seq(false, true)) { + withSQLConf(SQLConf.ANSI_ENABLED.key -> ansi.toString) { + val df = sql("SELECT SUM(v) FROM decimal_sum_cancellation") + val plan = 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 native: CometNativeExec => native }.nonEmpty) + checkSparkAnswer(df) + checkAnswer(df, Seq(Row(new java.math.BigDecimal("0.6")))) + } + } + } + } + } + } + + test("mixed engine sum/avg falls back when Spark Final would consume native AVG") { val data = (0 until 100).map(i => (i, i.toLong, i.toDouble, i % 7)) withParquetTable(data, "tbl") { withSQLConf( CometConf.COMET_ENABLE_FINAL_HASH_AGGREGATE.key -> "false", CometConf.COMET_SHUFFLE_ENABLED.key -> "true", CometConf.COMET_SHUFFLE_MODE.key -> "jvm") { - checkSparkAnswer( - "SELECT _4, SUM(_1), SUM(_2), SUM(_3), AVG(_1), AVG(_2), AVG(_3) FROM tbl GROUP BY _4") + checkSparkAnswerAndNumOfAggregates( + "SELECT _4, SUM(_1), SUM(_2), SUM(_3), AVG(_1), AVG(_2), AVG(_3) FROM tbl GROUP BY _4", + 0) } } } @@ -709,7 +1036,10 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { checkSparkAnswerAndNumOfAggregates("SELECT _2, COUNT(_1) FROM tbl GROUP BY _2", n) 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) + val avgStages = if (nativeShuffleEnabled) 2 else 0 + checkSparkAnswerAndNumOfAggregates( + "SELECT _2, AVG(_1) FROM tbl GROUP BY _2", + avgStages) } } } @@ -914,26 +1244,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) } } } @@ -1551,14 +1884,14 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } - test("test partial avg") { + test("AVG stays in Spark across a Spark shuffle") { Seq(true, false).foreach { dictionaryEnabled => withParquetTable( (0 until 5).map(i => (i.toDouble, i.toDouble % 2)), "tbl", dictionaryEnabled) { withSQLConf(CometConf.COMET_SHUFFLE_ENABLED.key -> "false") { - checkSparkAnswerAndNumOfAggregates("SELECT _2 , AVG(_1) FROM tbl GROUP BY _2", 1) + checkSparkAnswerAndNumOfAggregates("SELECT _2 , AVG(_1) FROM tbl GROUP BY _2", 0) } } } @@ -1595,7 +1928,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") { - val expectedNumOfCometAggregates = if (nativeShuffleEnabled) 2 else 1 + // Spark rewrites _7 to Long AVG, whose empty native buffer is also unsafe for a + // Spark Final until #5420. Keep all AVG partials in Spark across this boundary. + val expectedNumOfCometAggregates = if (nativeShuffleEnabled) 2 else 0 checkSparkAnswerAndNumOfAggregates( "SELECT _g2, AVG(_7) FROM tbl GROUP BY _g2", 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 79b668444d..9d65dec834 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -21,22 +21,23 @@ package org.apache.comet.rules import scala.util.Random +import org.apache.logging.log4j.Level 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, Partial} +import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, ExpressionInfo, Literal} +import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, BloomFilterAggregate, Final, Min, Partial, PartialMerge} 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.{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, ExtendedExplainInfo} -import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus, isSpark42Plus} -import org.apache.comet.serde.{Compatible, Unsupported} +import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus, isSpark42Plus, withFallbackReason} +import org.apache.comet.serde.{CometAggregateExpressionSerde, Compatible, ExprOuterClass, Unsupported} import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator} /** @@ -353,8 +354,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") @@ -370,11 +370,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) } } } @@ -395,8 +394,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) @@ -483,7 +482,7 @@ class CometExecRuleSuite extends CometTestBase { } } - test("CometExecRule should allow AVG mixed Comet partial and Spark final") { + test("CometExecRule should not allow AVG Comet partial and Spark final before buffer repair") { withTempView("test_data") { createTestDataFrame.createOrReplaceTempView("test_data") val sparkPlan = @@ -493,8 +492,9 @@ class CometExecRuleSuite extends CometTestBase { CometConf.COMET_ENABLE_FINAL_HASH_AGGREGATE.key -> "false", CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { val transformedPlan = applyCometExecRule(sparkPlan) - assert(countOperators(transformedPlan, classOf[HashAggregateExec]) == 1) // final - assert(countOperators(transformedPlan, classOf[CometHashAggregateExec]) == 1) // partial + // Matching field types do not make native AVG's empty (null, 0) state safe for Spark. + assert(countOperators(transformedPlan, classOf[HashAggregateExec]) == 2) + assert(countOperators(transformedPlan, classOf[CometHashAggregateExec]) == 0) } } } @@ -541,6 +541,175 @@ 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(id)" + (if (distinct) ", SUM(DISTINCT id)" 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") { + // Tag only the lowest exchange. A DISTINCT plan's upper exchange must inherit + // the native-only refusal from its now-Spark merge inputs, not from another tag. + val lowerShuffle = stripAQEPlan(sparkPlan).collect { + case shuffle: ShuffleExchangeExec => shuffle + }.last + withFallbackReason(lowerShuffle, fallback) + } + val transformed = applyCometExecRule(sparkPlan) + val nativeExpected = fallback == "none" + + // Shuffle is enabled, but a native-only shuffle can still fall back. The distinct + // rewrite also has intermediate PartialMerge and mixed Partial/PartialMerge stages. + 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("aggregate buffer direction opt-ins are independent") { + // A policy-only handler opts into consuming Spark state. Its inherited producer policy + // must stay false; serializing any expression is outside the scope of this fixture. + val reverseOnly = new CometAggregateExpressionSerde[Min] { + override def supportsSparkPartialToNativeFinal(fn: Min): Boolean = true + + override def convert( + aggExpr: AggregateExpression, + expr: Min, + inputs: Seq[Attribute], + binding: Boolean, + conf: SQLConf): Option[ExprOuterClass.AggExpr] = None + } + val fn = Min(Literal(1L)) + assert(reverseOnly.supportsSparkPartialToNativeFinal(fn)) + assert(!reverseOnly.supportsNativePartialToSparkFinal(fn)) + } + + test("restored partial records its reason when its current child is not native") { + // Wrap a converted input to prevent a re-entrant serde call from supplying the reason. + // This planner-only fixture never executes its synthetic buffer boundary. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false", + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "native") { + withTempView("test_data") { + createTestDataFrame.createOrReplaceTempView("test_data") + val plan = applyCometExecRule( + createSparkPlan(spark, "SELECT AVG(id) FROM test_data GROUP BY (id % 3)")) + val partial = plan.collectFirst { + case agg: CometHashAggregateExec if agg.modes == Seq(Partial) => agg + }.get + val sparkFinal = plan.collectFirst { + case agg: CometHashAggregateExec if agg.modes == Seq(Final) => + agg.originalPlan.asInstanceOf[HashAggregateExec] + }.get + val nonNativeChild = InputAdapter(partial.child) + val restored = CometExecRule(spark).revertUnsafePartialAggregates( + sparkFinal.copy(child = partial.copy(child = nonNativeChild))) + val sparkPartial = restored.children.head + assert(sparkPartial.isInstanceOf[HashAggregateExec]) + assert(sparkPartial.children.head.isInstanceOf[InputAdapter]) + assert(sparkPartial.children.head.output == nonNativeChild.output) + val reason = sparkPartial.getTagValue(CometExecRule.COMET_UNSAFE_PARTIAL).get + assert(sparkPartial.getTagValue(CometExplainInfo.FALLBACK_REASONS).get.contains(reason)) + assert(new ExtendedExplainInfo().getFallbackReasons(sparkPartial).contains(reason)) + } + } + } + + test("unrepaired native aggregate buffers warn once without rewriting query stages") { + // Construct the stage placeholder emitted by CometExchangeSink, including a native merge + // above it. No SQL reproduction or materialization is assumed: this pins the diagnostic + // when repair stops at a stage, and the absence of warnings for unrelated inner producers. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false", + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "native") { + withTempView("test_data") { + createTestDataFrame.createOrReplaceTempView("test_data") + val plan = applyCometExecRule( + createSparkPlan(spark, "SELECT AVG(id) FROM test_data GROUP BY (id % 3)")) + val partial = plan.collectFirst { + case agg: CometHashAggregateExec if agg.modes == Seq(Partial) => agg + }.get + val nativeFinal = plan.collectFirst { + case agg: CometHashAggregateExec if agg.modes == Seq(Final) => agg + }.get + val sparkFinal = nativeFinal.originalPlan.asInstanceOf[HashAggregateExec] + val sparkPartial = partial.originalPlan.asInstanceOf[HashAggregateExec] + val exchange = ShuffleExchangeExec( + org.apache.spark.sql.catalyst.plans.physical.SinglePartition, + partial) + val stage = ShuffleQueryStageExec(0, exchange, exchange.canonicalized) + val placeholder = CometSinkPlaceHolder( + org.apache.comet.serde.OperatorOuterClass.Operator.getDefaultInstance, + stage, + stage) + val nativeMerge = partial.copy( + aggregateExpressions = partial.aggregateExpressions.map(_.copy(mode = PartialMerge)), + child = placeholder) + val warning = "Comet could not restore a native intermediate buffer producer" + val rule = CometExecRule(spark) + for { + (child, shouldWarn) <- Seq( + nativeMerge -> true, + placeholder -> true, + sparkPartial.copy(child = nativeFinal) -> false, + nativeFinal -> false) + logFallback <- Seq("false", "true") + } { + withSQLConf(CometConf.COMET_EXPLAIN_FALLBACK_LOG_ENABLED.key -> logFallback) { + val consumer = sparkFinal.copy(child = child) + val appender = new LogAppender("unrepaired aggregate buffers") + withLogAppender(appender, Seq("org.apache.comet"), Some(Level.WARN)) { + assert(rule.revertUnsafePartialAggregates(consumer) eq consumer) + assert(rule.revertUnsafePartialAggregates(consumer) eq consumer) + } + assert(consumer.child eq child) + val warnings = + appender.loggingEvents.count(_.getMessage.getFormattedMessage.contains(warning)) + assert(warnings == (if (shouldWarn) 1 else 0), s"$child: $warnings") + assert( + new ExtendedExplainInfo() + .getFallbackReasons(consumer) + .exists(_.contains(warning)) == + shouldWarn) + } + } + assert(stage.plan eq exchange) + assert(exchange.child eq partial) + } + } + } + test("CometExecRule should not allow decimal SUM mixed execution") { withTempView("test_data") { createTestDataFrame.createOrReplaceTempView("test_data") @@ -556,9 +725,9 @@ class CometExecRuleSuite extends CometTestBase { CometConf.COMET_ENABLE_FINAL_HASH_AGGREGATE.key -> "false", CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { val transformedPlan = applyCometExecRule(sparkPlan) - // Decimal SUM overflow detection (ANSI throw / Legacy null) does not survive a - // Spark-partial / Comet-final split, so mixed execution is unsafe and the partial - // must also fall back to Spark. + // Native decimal SUM makes precision overflow sticky (or throws eagerly in ANSI), + // while Spark's generated scalar Partial can recover after a later cancelling input. + // Keep the Partial in Spark even though the emitted buffer field types match. assert(countOperators(transformedPlan, classOf[HashAggregateExec]) == 2) assert(countOperators(transformedPlan, classOf[CometHashAggregateExec]) == 0) } diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShufflePlanningSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShufflePlanningSuite.scala index a27f7c7b14..b2c6c0f8a3 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShufflePlanningSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShufflePlanningSuite.scala @@ -548,7 +548,7 @@ class CometCelebornShufflePlanningSuite extends CometTestBase { for { fallback <- Seq("partition threshold", "unsupported array hash key") - function <- Seq("collect_list", "collect_set", "avg") + function <- Seq("collect_list", "collect_set", "avg", "count") } { test(s"native $fallback preserves $function aggregate buffers with AQE=$adaptive") { val complexKey = fallback == "unsupported array hash key" @@ -563,8 +563,11 @@ class CometCelebornShufflePlanningSuite extends CometTestBase { SQLConf.SHUFFLE_PARTITIONS.key -> "4", CometConf.COMET_SHUFFLE_MODE.key -> "native") { val grouping = if (complexKey) "array(id % 3)" else "id % 3" - val aggregate = - if (function == "avg") "avg(value)" else s"sort_array($function(value))" + val aggregate = if (function.startsWith("collect_")) { + s"sort_array($function(value))" + } else { + s"$function(value)" + } val query = spark .range(0, 18, 1, 4) .selectExpr(s"$grouping AS grouping_key", "id AS value") @@ -579,13 +582,15 @@ class CometCelebornShufflePlanningSuite extends CometTestBase { val nativeAggregates = collect(executedPlan) { case aggregate: CometHashAggregateExec => aggregate } - if (function == "avg") { - // AVG's intermediate state is Spark-compatible; native partials remain safe. - assert(nativeAggregates.nonEmpty, s"$executedPlan") + if (function == "count") { + // COUNT's non-null Long buffer is safe for Spark Final to consume. + assert(nativeAggregates.size == 1, s"$executedPlan") + assert(nativeAggregates.head.modes == Seq(Partial), s"$executedPlan") } else { // A Spark final cannot deserialize Comet's ArrayType collect_list/collect_set - // state as its BinaryType buffer. Both halves must agree when the exchange falls - // back, not just when an aggregate operator itself is unsupported. + // state as BinaryType, or safely merge AVG's never-updated (null, 0) buffer. + // Both halves must agree when the exchange falls back, not just when an + // aggregate operator itself is unsupported. assert(nativeAggregates.isEmpty, s"$executedPlan") } } --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
