diff --git a/docs/source/user-guide/latest/compatibility/operators.md b/docs/source/user-guide/latest/compatibility/operators.md index e039df23b76..f89635bfbdc 100644 --- a/docs/source/user-guide/latest/compatibility/operators.md +++ b/docs/source/user-guide/latest/compatibility/operators.md @@ -35,6 +35,16 @@ Sampling with replacement (`df.sample(withReplacement = true, ...)`) falls back it draws from a Poisson distribution that Comet does not implement natively ([#5109](https://github.com/apache/datafusion-comet/issues/5109)). +## Aggregation + +Ungrouped `AVG` and `TRY_AVG` on `DECIMAL(p, s)` fall back to Spark when the intermediate sum +has reached precision 38 (`p >= 28`). Spark's generated aggregation code can retain a wider +intermediate sum and check precision after dividing by the count. Comet instead records an +overflow while accumulating that sum, which can discard a valid average. + +Both partial and final aggregates stay in Spark, including when a shuffle separates them. +Grouped averages and narrower decimal averages remain eligible for native execution. + ## Window Functions Comet runs `WindowExec` natively and it is enabled by default (`spark.comet.exec.window.enabled`). A broad set of @@ -61,6 +71,8 @@ incorrect result. When any single window expression in a `WindowExec` falls back support as the batch aggregates, so these fall back in both contexts. - `sum` or `avg` on `DECIMAL` with a sliding (non ever-expanding) frame, because the sliding path would wrap on overflow instead of returning Spark's `NULL`. +- `avg` or `try_avg` on `DECIMAL(p, s)` with an expanding frame when `p >= 28`, because Spark can preserve + a wider intermediate sum until division, as described under aggregation above. - `RANGE` frame with an explicit offset when the `ORDER BY` column is `DATE` or `DECIMAL` ([#4834](https://github.com/apache/datafusion-comet/issues/4834)). - `first_value` / `last_value` on a `RANGE` frame with a literal offset diff --git a/docs/source/user-guide/latest/tuning.md b/docs/source/user-guide/latest/tuning.md index 5499052464c..ef57a1a7052 100644 --- a/docs/source/user-guide/latest/tuning.md +++ b/docs/source/user-guide/latest/tuning.md @@ -220,6 +220,12 @@ Comet provides a fully native shuffle implementation, which generally provides t supports `HashPartitioning`, `RangePartitioning` and `SinglePartitioning` but currently only supports primitive type partitioning keys. Columns that are not partitioning keys may contain complex types like maps, structs, and arrays. +Hash partitioning on decimal keys with precision greater than 18 falls back because native hashing does not match +Spark's partition assignments. This can affect decimal aggregate overflow behavior, including `AVG(DISTINCT ...)`. +With `spark.comet.shuffle.mode=auto`, Comet uses Columnar Shuffle when eligible; with `native`, it uses Spark shuffle. +The restriction applies only to hash partitioning keys: wider decimals remain supported as payload columns and +range partitioning keys. + #### Columnar (JVM) Shuffle Comet Columnar shuffle is JVM-based and supports `HashPartitioning`, `RoundRobinPartitioning`, `RangePartitioning`, and diff --git a/native/spark-expr/src/agg_funcs/avg.rs b/native/spark-expr/src/agg_funcs/avg.rs index 24a9a309911..fa89bee1c85 100644 --- a/native/spark-expr/src/agg_funcs/avg.rs +++ b/native/spark-expr/src/agg_funcs/avg.rs @@ -146,8 +146,11 @@ pub struct AvgAccumulator { impl Accumulator for AvgAccumulator { fn state(&mut self) -> Result> { + // Spark's final AVG adds partial sums without coalescing nulls. Empty partials + // must therefore carry a zero sum, even when update_batch was never called. + let sum = if self.count == 0 { Some(0.0) } else { self.sum }; Ok(vec![ - ScalarValue::Float64(self.sum), + ScalarValue::Float64(sum), ScalarValue::from(self.count), ]) } @@ -345,3 +348,50 @@ where + self.sums.capacity() * std::mem::size_of::() } } + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::Float64Array; + + #[test] + fn empty_partial_state_matches_spark() -> Result<()> { + for values in [None, Some(vec![]), Some(vec![None, None])] { + let mut acc = AvgAccumulator::default(); + if let Some(values) = values { + acc.update_batch(&[Arc::new(Float64Array::from(values))])?; + } + assert_eq!( + acc.state()?, + vec![ScalarValue::Float64(Some(0.0)), ScalarValue::Int64(Some(0))] + ); + assert_eq!(acc.evaluate()?, ScalarValue::Float64(None)); + } + Ok(()) + } + + #[test] + fn merge_empty_partial_states() -> Result<()> { + for nonempty in [false, true] { + let mut final_acc = AvgAccumulator::default(); + let values = nonempty.then(|| vec![Some(2.0), Some(4.0)]); + for values in [None, values, Some(vec![None])] { + let mut partial = AvgAccumulator::default(); + if let Some(values) = values { + partial.update_batch(&[Arc::new(Float64Array::from(values))])?; + } + let state = partial + .state()? + .into_iter() + .map(|value| value.to_array_of_size(1)) + .collect::>>()?; + final_acc.merge_batch(&state)?; + } + assert_eq!( + final_acc.evaluate()?, + ScalarValue::Float64(nonempty.then_some(3.0)) + ); + } + Ok(()) + } +} diff --git a/native/spark-expr/src/agg_funcs/avg_decimal.rs b/native/spark-expr/src/agg_funcs/avg_decimal.rs index 186e0272445..b6b139d044d 100644 --- a/native/spark-expr/src/agg_funcs/avg_decimal.rs +++ b/native/spark-expr/src/agg_funcs/avg_decimal.rs @@ -300,14 +300,22 @@ fn make_decimal128(value: Option, precision: u8, scale: i8) -> ScalarValue impl Accumulator for AvgDecimalAccumulator { fn state(&mut self) -> Result> { + // Spark distinguishes an empty partial (zero sum) from an overflow (null sum). + let sum = if !self.is_not_null { + None + } else if self.count == 0 { + Some(0) + } else { + self.sum + }; Ok(vec![ - ScalarValue::Decimal128(self.sum, self.sum_precision, self.sum_scale), + ScalarValue::Decimal128(sum, self.sum_precision, self.sum_scale), ScalarValue::from(self.count), ]) } fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { - if !self.is_empty && !self.is_not_null { + if !self.is_not_null { // This means there's a overflow in decimal, so we will just skip the rest // of the computation return Ok(()); @@ -342,6 +350,15 @@ impl Accumulator for AvgDecimalAccumulator { // counts are summed self.count += sum(partial_counts).unwrap_or_default(); + // Empty partials have non-null zero buffers. A null sum is an overflow marker; + // grouped partials also null the count. Do not let sum() skip these markers or + // let a later empty partial revive an accumulator that has already overflowed. + if !self.is_not_null || partial_sums.null_count() > 0 || partial_counts.null_count() > 0 { + self.is_not_null = false; + self.sum = None; + return Ok(()); + } + // sums are summed if let Some(x) = sum(partial_sums) { let v = self.sum.get_or_insert(0); @@ -359,24 +376,23 @@ impl Accumulator for AvgDecimalAccumulator { } fn evaluate(&mut self) -> Result { - // Check for overflow during sum accumulation in ANSI mode. - // This matches Spark's DecimalDivideWithOverflowCheck behavior. - // `count` guards against reporting an overflow when there was nothing to sum: an - // empty or all-null input also leaves `sum` as None, and `is_empty` cannot - // distinguish those cases because the counts merged in `merge_batch` are never null. - if self.sum.is_none() - && !self.is_empty - && self.count > 0 - && self.eval_mode == EvalMode::Ansi - { + // A grouped overflow can have a null count, leaving the merged count at zero. + // Check the overflow marker before treating a zero count as empty input. + let has_overflow = + !self.is_not_null || (self.sum.is_none() && !self.is_empty && self.count > 0); + if has_overflow && self.eval_mode == EvalMode::Ansi { let error = decimal_sum_overflow_error("avg"); return Err(self.wrap_error_with_context(error)); } - // Also check if is_not_null is false (indicates overflow) - if !self.is_not_null && self.count > 0 && self.eval_mode == EvalMode::Ansi { - let error = decimal_sum_overflow_error("avg"); - return Err(self.wrap_error_with_context(error)); + // An all-empty native final now merges zero sums. It must still return NULL, + // rather than dividing that sum by zero, in every evaluation mode. + if self.count == 0 || has_overflow { + return Ok(make_decimal128( + None, + self.target_precision, + self.target_scale, + )); } let scaler = 10_i128.pow(self.target_scale.saturating_sub(self.sum_scale) as u32); @@ -608,9 +624,9 @@ impl GroupsAccumulator for AvgDecimalGroupsAccumulator { let target_max = MAX_DECIMAL128_FOR_EACH_PRECISION[self.target_precision as usize]; for (is_not_null, (sum, count)) in nulls.into_iter().zip(iter) { - // Check for overflow during sum accumulation in ANSI mode. - // This matches Spark's DecimalDivideWithOverflowCheck behavior. - if !is_not_null && count > 0 && self.eval_mode == EvalMode::Ansi { + // A null state marks overflow even if a shuffle zeroed its null count + // payload. Empty/all-null groups keep valid zero buffers instead. + if !is_not_null && self.eval_mode == EvalMode::Ansi { let error = decimal_sum_overflow_error("avg"); return Err(self.wrap_error_with_context(error)); } @@ -687,3 +703,282 @@ fn avg(sum: i128, count: i128, target_min: i128, target_max: i128, scaler: i128) None } } + +#[cfg(test)] +mod tests { + use super::*; + + fn accumulator(mode: EvalMode) -> AvgDecimalAccumulator { + AvgDecimalAccumulator::new(2, 17, 11, 6, mode, None, crate::create_query_context_map()) + } + + fn assert_overflow(acc: &mut AvgDecimalAccumulator) -> Result<()> { + if acc.eval_mode == EvalMode::Ansi { + assert!(acc.evaluate().is_err()); + } else { + assert_eq!( + acc.evaluate()?, + ScalarValue::Decimal128(None, acc.target_precision, acc.target_scale) + ); + } + Ok(()) + } + + fn overflowed_accumulator(mode: EvalMode) -> Result { + let mut acc = accumulator(mode); + // Both partial sums fit, but their merge exceeds the sum's decimal precision. + let sums = + Decimal128Array::from(vec![10_i128.pow(17) - 1, 1]).with_precision_and_scale(17, 2)?; + acc.merge_batch(&[Arc::new(sums), Arc::new(Int64Array::from(vec![1, 1]))])?; + Ok(acc) + } + + #[test] + fn empty_partial_state_matches_spark() -> Result<()> { + for mode in [EvalMode::Legacy, EvalMode::Ansi, EvalMode::Try] { + for values in [None, Some(vec![]), Some(vec![None, None])] { + let mut acc = accumulator(mode); + if let Some(values) = values { + let values = Decimal128Array::from(values).with_precision_and_scale(7, 2)?; + acc.update_batch(&[Arc::new(values)])?; + } + assert_eq!( + acc.state()?, + vec![ + ScalarValue::Decimal128(Some(0), 17, 2), + ScalarValue::Int64(Some(0)) + ] + ); + assert_eq!(acc.evaluate()?, ScalarValue::Decimal128(None, 11, 6)); + } + } + Ok(()) + } + + #[test] + fn merge_empty_partial_states() -> Result<()> { + for mode in [EvalMode::Legacy, EvalMode::Ansi, EvalMode::Try] { + for nonempty in [false, true] { + let mut final_acc = accumulator(mode); + let values = nonempty.then(|| vec![Some(100), Some(300)]); + for values in [None, values, Some(vec![None])] { + let mut partial = accumulator(mode); + if let Some(values) = values { + let values = + Decimal128Array::from(values).with_precision_and_scale(7, 2)?; + partial.update_batch(&[Arc::new(values)])?; + } + let state = partial + .state()? + .into_iter() + .map(|value| value.to_array_of_size(1)) + .collect::>>()?; + final_acc.merge_batch(&state)?; + } + assert_eq!( + final_acc.evaluate()?, + ScalarValue::Decimal128(nonempty.then_some(2_000_000), 11, 6) + ); + } + } + Ok(()) + } + + #[test] + fn overflow_partial_state_is_not_empty() -> Result<()> { + for mode in [EvalMode::Legacy, EvalMode::Ansi, EvalMode::Try] { + let mut acc = overflowed_accumulator(mode)?; + assert_eq!( + acc.state()?, + vec![ + ScalarValue::Decimal128(None, 17, 2), + ScalarValue::Int64(Some(2)) + ] + ); + assert_overflow(&mut acc)?; + } + Ok(()) + } + + #[test] + fn update_overflow_partial_state_is_preserved() -> Result<()> { + for mode in [EvalMode::Legacy, EvalMode::Ansi, EvalMode::Try] { + let mut acc = AvgDecimalAccumulator::new( + 2, + 38, + 38, + 6, + mode, + None, + crate::create_query_context_map(), + ); + let values = Decimal128Array::from(vec![10_i128.pow(38) - 1; 2]) + .with_precision_and_scale(38, 2)?; + acc.update_batch(&[Arc::new(values)])?; + // update_single retains the last valid sum on overflow. It must not leak + // into the partial buffer or be evaluated as a valid average. + assert_eq!(acc.state()?[0], ScalarValue::Decimal128(None, 38, 2)); + assert_overflow(&mut acc)?; + } + Ok(()) + } + + #[test] + fn empty_partial_does_not_mask_overflow() -> Result<()> { + for mode in [EvalMode::Legacy, EvalMode::Ansi, EvalMode::Try] { + for reverse in [false, true] { + for batch_size in [1, 2] { + let mut states = [ + accumulator(mode).state()?, + overflowed_accumulator(mode)?.state()?, + ]; + if reverse { + states.reverse(); + } + let mut final_acc = accumulator(mode); + for batch in states.chunks(batch_size) { + let arrays = (0..2) + .map(|i| { + ScalarValue::iter_to_array( + batch.iter().map(|state| state[i].clone()), + ) + }) + .collect::>>()?; + final_acc.merge_batch(&arrays)?; + } + assert_eq!( + final_acc.state()?, + vec![ + ScalarValue::Decimal128(None, 17, 2), + ScalarValue::Int64(Some(2)) + ] + ); + assert_overflow(&mut final_acc)?; + } + } + } + Ok(()) + } + + #[test] + fn grouped_overflow_with_null_count_is_not_empty() -> Result<()> { + for mode in [EvalMode::Legacy, EvalMode::Ansi, EvalMode::Try] { + let mut partial = AvgDecimalGroupsAccumulator::new( + &DataType::Decimal128(38, 6), + &DataType::Decimal128(38, 2), + 38, + 6, + 38, + 2, + mode, + None, + crate::create_query_context_map(), + ); + let values = Decimal128Array::from(vec![10_i128.pow(38) - 1; 2]) + .with_precision_and_scale(38, 2)?; + partial.update_batch(&[Arc::new(values)], &[0, 0], None, 1)?; + let overflow = partial.state(EmitTo::All)?; + // Grouped partials encode overflow by nulling both the sum and count. + assert!(overflow[0].is_null(0)); + assert!(overflow[1].is_null(0)); + + let new_acc = || { + AvgDecimalAccumulator::new( + 2, + 38, + 38, + 6, + mode, + None, + crate::create_query_context_map(), + ) + }; + for reverse in [false, true] { + let empty = new_acc() + .state()? + .into_iter() + .map(|value| value.to_array_of_size(1)) + .collect::>>()?; + let mut states = [overflow.clone(), empty]; + if reverse { + states.reverse(); + } + let mut final_acc = new_acc(); + for state in states { + final_acc.merge_batch(&state)?; + } + assert_overflow(&mut final_acc)?; + } + + // A null count can leave is_empty set after merge. A subsequent update + // must still respect the overflow marker, regardless of that empty flag. + let mut final_acc = new_acc(); + final_acc.merge_batch(&overflow)?; + let values = Decimal128Array::from(vec![100]).with_precision_and_scale(38, 2)?; + final_acc.update_batch(&[Arc::new(values)])?; + assert_overflow(&mut final_acc)?; + } + Ok(()) + } + + #[test] + fn grouped_overflow_with_zeroed_null_count_is_not_empty() -> Result<()> { + for mode in [EvalMode::Legacy, EvalMode::Ansi, EvalMode::Try] { + let new_acc = || { + AvgDecimalGroupsAccumulator::new( + &DataType::Decimal128(38, 38), + &DataType::Decimal128(38, 38), + 38, + 38, + 38, + 38, + mode, + None, + crate::create_query_context_map(), + ) + }; + let mut partial = new_acc(); + // Group 0 has no input, group 1 has only nulls, and group 2 overflows. + let values = Decimal128Array::from(vec![ + None, + None, + Some(6 * 10_i128.pow(37)), + Some(6 * 10_i128.pow(37)), + ]) + .with_precision_and_scale(38, 38)?; + partial.update_batch(&[Arc::new(values)], &[1, 1, 2, 2], None, 3)?; + let state = partial.state(EmitTo::All)?; + + // Rebuilding the logical values models a columnar shuffle's row roundtrip: + // overflow stays null, but the null count's underlying payload becomes 0. + let sums = state[0] + .as_primitive::() + .iter() + .collect::() + .with_precision_and_scale(38, 38)?; + let counts = state[1] + .as_primitive::() + .iter() + .collect::(); + assert!(counts.is_null(2)); + assert_eq!(counts.value(2), 0); + + let mut final_acc = new_acc(); + final_acc.merge_batch(&[Arc::new(sums), Arc::new(counts)], &[0, 1, 2], None, 3)?; + let empty = final_acc.evaluate(EmitTo::First(2))?; + assert_eq!(empty.len(), 2); + assert_eq!(empty.null_count(), 2); + + let result = final_acc.evaluate(EmitTo::All); + if mode == EvalMode::Ansi { + let error = result.unwrap_err().to_string(); + assert!(error.contains("ARITHMETIC_OVERFLOW"), "{error}"); + } else { + let result = result?; + assert_eq!(result.len(), 1); + assert!(result.is_null(0)); + } + } + Ok(()) + } +} 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 c69602fc80b..aeee9ff35e5 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) @@ -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.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 @@ -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/spark/sql/comet/CometWindowExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowExec.scala index 4d4abf3d2e0..e65f1b2961a 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowExec.scala @@ -333,6 +333,28 @@ object CometWindowExec extends CometOperatorSerde[WindowExec] { case _: SpecifiedWindowFrame => false case _ => true } + if (isEverExpanding) { + windowExpr.windowFunction match { + case agg: AggregateExpression => + agg.aggregateFunction match { + case a: Average => + a.sumDataType match { + case decimal: DecimalType if decimal.precision == DecimalType.MAX_PRECISION => + // Spark can preserve an out-of-precision intermediate sum in a window + // buffer until AVG divides it by the count. Comet's decimal accumulator + // instead records overflow immediately, so keep these windows in Spark + // when the intermediate precision cannot be widened any further. + withFallbackReason( + windowExpr, + "AVG on DECIMAL with maximum-precision intermediate state is not supported") + return None + case _ => + } + case _ => + } + case _ => + } + } if (!isEverExpanding) { windowExpr.windowFunction match { case agg: AggregateExpression => diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala index 32fe282c955..c70742c37a3 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala @@ -397,12 +397,10 @@ object CometShuffleExchangeExec _: FloatType | _: DoubleType | _: StringType | _: BinaryType | _: TimestampType | _: TimestampNTZType | _: DateType => true - case _: DecimalType => - // TODO enforce this check - // https://github.com/apache/datafusion-comet/issues/3079 - // Decimals with precision > 18 require Java BigDecimal conversion before hashing - // d.precision <= 18 - true + case d: DecimalType => + // Spark hashes wider decimals through BigInteger bytes, which native hashing does not + // match. Different partition assignments can change decimal AVG overflow behavior. + d.precision <= 18 case dt if isTimeType(dt) => true case _ => 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 216f9f1e4bd..5505cea8cca 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 @@ -31,7 +31,7 @@ import org.apache.spark.internal.Logging import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, AttributeSet, Expression, ExpressionSet, Generator, NamedExpression, SortOrder} -import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, AggregateMode, CollectList, CollectSet, Final, First, Last, Partial, PartialMerge, Percentile} +import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, AggregateMode, Average, CollectList, CollectSet, Final, First, Last, Partial, PartialMerge, Percentile} import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} import org.apache.spark.sql.catalyst.plans._ import org.apache.spark.sql.catalyst.plans.physical._ @@ -1627,6 +1627,30 @@ case class CometUnionExec( trait CometBaseAggregate { + protected def aggregateSupportLevel(op: BaseAggregateExec): SupportLevel = { + val unsupportedAverage = op.groupingExpressions.isEmpty && + op.aggregateExpressions.exists(_.aggregateFunction match { + case avg: Average => + avg.sumDataType match { + case decimal: DecimalType => decimal.precision == DecimalType.MAX_PRECISION + case _ => false + } + case _ => false + }) + + if (unsupportedAverage) { + // Spark's ungrouped codegen buffers can retain a wider decimal sum until AVG divides + // by the count, including while merging partials. Comet records overflow immediately. + // Both conversion and unsafe-partial tagging consult this operator support check, so + // the partial and final fall back together without exchanging incompatible buffers. + Unsupported( + Some( + "Ungrouped AVG on DECIMAL with maximum-precision intermediate state is not supported")) + } else { + Compatible() + } + } + def doConvert( aggregate: BaseAggregateExec, builder: Operator.Builder, @@ -1976,7 +2000,7 @@ object CometHashAggregateExec op.aggregateExpressions.exists(_.mode == Final)) { return Unsupported(Some("Final aggregates disabled via test config")) } - Compatible() + aggregateSupportLevel(op) } override def convert( @@ -2018,7 +2042,7 @@ object CometObjectHashAggregateExec op.aggregateExpressions.exists(_.mode == Final)) { return Unsupported(Some("Final aggregates disabled via test config")) } - Compatible() + aggregateSupportLevel(op) } override def convert( diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q49/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q49/extended.txt index a2dc4cfd61e..16ee30d1e5e 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q49/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q49/extended.txt @@ -1,7 +1,7 @@ CometColumnarToRow +- CometTakeOrderedAndProject +- CometHashAggregate - +- CometExchange + +- CometColumnarExchange +- CometHashAggregate +- CometUnion :- CometProject diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark3_5/q70a/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark3_5/q70a/extended.txt index 8e5298ac981..4fb957fda26 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark3_5/q70a/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark3_5/q70a/extended.txt @@ -5,7 +5,7 @@ CometColumnarToRow +- CometSort +- CometExchange +- CometHashAggregate - +- CometExchange + +- CometColumnarExchange +- CometHashAggregate +- CometUnion :- CometHashAggregate diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark4_1/q14a/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark4_1/q14a/extended.txt index 65b3bc63498..fa1a6aa613b 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark4_1/q14a/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark4_1/q14a/extended.txt @@ -1,7 +1,7 @@ CometColumnarToRow +- CometTakeOrderedAndProject +- CometHashAggregate - +- CometExchange + +- CometColumnarExchange +- CometHashAggregate +- CometUnion :- CometHashAggregate diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q14a/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q14a/extended.txt index 975b8d63cd0..c4dbb812fd3 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q14a/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q14a/extended.txt @@ -1,7 +1,7 @@ CometColumnarToRow +- CometTakeOrderedAndProject +- CometHashAggregate - +- CometExchange + +- CometColumnarExchange +- CometHashAggregate +- CometUnion :- CometHashAggregate diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q36a/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q36a/extended.txt index c512a09c2c4..36e6b2d03e1 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q36a/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q36a/extended.txt @@ -5,7 +5,7 @@ CometColumnarToRow +- CometSort +- CometExchange +- CometHashAggregate - +- CometExchange + +- CometColumnarExchange +- CometHashAggregate +- CometUnion :- CometHashAggregate diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q49/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q49/extended.txt index a2dc4cfd61e..16ee30d1e5e 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q49/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q49/extended.txt @@ -1,7 +1,7 @@ CometColumnarToRow +- CometTakeOrderedAndProject +- CometHashAggregate - +- CometExchange + +- CometColumnarExchange +- CometHashAggregate +- CometUnion :- CometProject diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q5a/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q5a/extended.txt index 50f60ae8c1a..6289695c112 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q5a/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q5a/extended.txt @@ -1,7 +1,7 @@ CometColumnarToRow +- CometTakeOrderedAndProject +- CometHashAggregate - +- CometExchange + +- CometColumnarExchange +- CometHashAggregate +- CometUnion :- CometHashAggregate diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q70a/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q70a/extended.txt index b3435650dae..d6e7c0c320f 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q70a/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q70a/extended.txt @@ -5,7 +5,7 @@ CometColumnarToRow +- CometSort +- CometExchange +- CometHashAggregate - +- CometExchange + +- CometColumnarExchange +- CometHashAggregate +- CometUnion :- CometHashAggregate diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q77a/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q77a/extended.txt index c88625a73f9..fd77c21f332 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q77a/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q77a/extended.txt @@ -1,7 +1,7 @@ CometColumnarToRow +- CometTakeOrderedAndProject +- CometHashAggregate - +- CometExchange + +- CometColumnarExchange +- CometHashAggregate +- CometUnion :- CometHashAggregate diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q80a/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q80a/extended.txt index a2ebd27e56f..bee6cc882cb 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q80a/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q80a/extended.txt @@ -1,7 +1,7 @@ CometColumnarToRow +- CometTakeOrderedAndProject +- CometHashAggregate - +- CometExchange + +- CometColumnarExchange +- CometHashAggregate +- CometUnion :- CometHashAggregate diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q86a/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q86a/extended.txt index a4db2a8ee0a..e4652ba10d4 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q86a/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q86a/extended.txt @@ -5,7 +5,7 @@ CometColumnarToRow +- CometSort +- CometExchange +- CometHashAggregate - +- CometExchange + +- CometColumnarExchange +- CometHashAggregate +- CometUnion :- CometHashAggregate diff --git a/spark/src/test/scala/org/apache/comet/CometFuzzAggregateSuite.scala b/spark/src/test/scala/org/apache/comet/CometFuzzAggregateSuite.scala index 3674887f80f..32b21cb92ad 100644 --- a/spark/src/test/scala/org/apache/comet/CometFuzzAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometFuzzAggregateSuite.scala @@ -19,6 +19,10 @@ package org.apache.comet +import org.apache.spark.sql.execution.aggregate.HashAggregateExec +import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec +import org.apache.spark.sql.types.DecimalType + import org.apache.comet.DataTypeSupport.isComplexType class CometFuzzAggregateSuite extends CometFuzzTestBase { @@ -31,7 +35,16 @@ class CometFuzzAggregateSuite extends CometFuzzTestBase { val (_, cometPlan) = checkSparkAnswer(sql) assert(1 == collectNativeScans(cometPlan).length) - checkSparkAnswerAndOperator(sql) + val hasWideDecimalKey = df.schema(col).dataType match { + case d: DecimalType => d.precision > 18 + case _ => false + } + // Wide decimal hash keys require Spark shuffle when columnar shuffle is disabled. + if (CometConf.COMET_SHUFFLE_MODE.get() == "native" && hasWideDecimalKey) { + checkSparkAnswerAndOperator(sql, classOf[HashAggregateExec], classOf[ShuffleExchangeExec]) + } else { + checkSparkAnswerAndOperator(sql) + } } } @@ -55,7 +68,18 @@ class CometFuzzAggregateSuite extends CometFuzzTestBase { val (_, cometPlan) = checkSparkAnswer(sql) assert(1 == collectNativeScans(cometPlan).length) - checkSparkAnswerAndOperator(sql) + val hasWideDecimalKey = Seq("c1", "c2", "c3", col).exists { key => + df.schema(key).dataType match { + case d: DecimalType => d.precision > 18 + case _ => false + } + } + // Check both GROUP BY and DISTINCT keys, not unrelated decimal payload columns. + if (CometConf.COMET_SHUFFLE_MODE.get() == "native" && hasWideDecimalKey) { + checkSparkAnswerAndOperator(sql, classOf[HashAggregateExec], classOf[ShuffleExchangeExec]) + } else { + checkSparkAnswerAndOperator(sql) + } } } diff --git a/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala b/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala index 846fd633a3e..9e485921f0c 100644 --- a/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala @@ -21,8 +21,15 @@ package org.apache.comet import org.apache.spark.SparkConf import org.apache.spark.sql._ +import org.apache.spark.sql.catalyst.expressions.{Ascending, SortOrder} +import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, RangePartitioning, SinglePartition} +import org.apache.spark.sql.comet.CometSinkPlaceHolder +import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometNativeShuffle, CometShuffleExchangeExec, CometShuffleManager} +import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec import org.apache.spark.sql.internal.SQLConf +import org.apache.comet.serde.OperatorOuterClass + class CometSparkSessionExtensionsSuite extends CometTestBase { import CometSparkSessionExtensions._ @@ -72,6 +79,66 @@ class CometSparkSessionExtensionsSuite extends CometTestBase { assert(isCometLoaded(conf)) } + test("wide decimal hash keys use Spark-compatible shuffle partitioning") { + for { + mode <- Seq("native", "auto", "jvm") + precision <- Seq(18, 19, 38) + } { + withSQLConf( + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> mode, + CometConf.COMET_SHUFFLE_NATIVE_HASH_PARTITIONING_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_NATIVE_RANGE_PARTITIONING_ENABLED.key -> "true", + "spark.shuffle.manager" -> classOf[CometShuffleManager].getName) { + val originalChild = spark + .range(1) + .selectExpr(s"CAST(id AS DECIMAL($precision, 0)) AS d", "id") + .queryExecution + .executedPlan + val child = CometSinkPlaceHolder( + OperatorOuterClass.Operator.getDefaultInstance, + originalChild, + originalChild) + val shuffle = ShuffleExchangeExec(HashPartitioning(Seq(child.output.head), 2), child) + val expected = if (mode == "jvm" || (mode == "auto" && precision > 18)) { + Some(CometColumnarShuffle) + } else if (precision <= 18) { + Some(CometNativeShuffle) + } else { + None + } + + withClue(s"mode=$mode, precision=$precision: ") { + assert(CometShuffleExchangeExec.shuffleSupported(shuffle) == expected) + if (expected.isEmpty) { + assert( + shuffle + .getTagValue(CometExplainInfo.FALLBACK_REASONS) + .getOrElse(Set.empty[String]) + .exists(_.contains("unsupported hash partitioning data type for native shuffle"))) + } else { + // A native failure must not tag an exchange that can use the columnar path. + assert(shuffle.getTagValue(CometExplainInfo.FALLBACK_REASONS).isEmpty) + } + + if (mode == "native" && precision == 38) { + // Wide decimals remain supported as payloads, range keys, and in a single partition. + Seq( + HashPartitioning(Seq(child.output(1)), 2), + RangePartitioning(Seq(SortOrder(child.output.head, Ascending)), 2), + SinglePartition).foreach { partitioning => + val supported = ShuffleExchangeExec(partitioning, child) + assert( + CometShuffleExchangeExec.shuffleSupported(supported).contains(CometNativeShuffle)) + } + } + } + } + } + } + test("Arrow properties") { NativeBase.setLoaded(false) NativeBase.load() 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 143248f551c..b2237802da8 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -30,18 +30,20 @@ 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} 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.execution.shuffle.CometShuffleExchangeExec +import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, RangePartitioning} +import org.apache.spark.sql.comet.{CometFilterExec, CometHashAggregateExec} +import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometNativeShuffle, 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, HashAggregateExec} 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} +import org.apache.spark.sql.types.{DataTypes, DecimalType, 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,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 => @@ -599,6 +668,48 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } + test("avg with empty partial input and Spark final") { + import org.apache.spark.sql.catalyst.expressions.aggregate.{Final, Partial} + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + 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 -> "false") { + withTempPath { path => + // Four files remain separate scan partitions. Only one partition has a matching row, + // so the other native partials must emit Spark's (sum = 0, count = 0) AVG buffer. + spark + .range(0L, 8L, 1L, 4) + .selectExpr( + "id", + "cast(id + 1 as int) as quantity", + "cast(id + 10 as decimal(7, 2)) as price") + .write + .parquet(path.getCanonicalPath) + withParquetTable(path.getCanonicalPath, "avg_empty_partials") { + for (predicate <- Seq("id = 1", "id < 0")) { + // Spark optimizes this narrow decimal AVG to AVG(UnscaledValue(price)), so both + // averages exercise the mixed non-decimal accumulator path. + val df = + sql(s"SELECT avg(quantity), avg(price) FROM avg_empty_partials WHERE $predicate") + checkSparkAnswer(df) + val plan = df.queryExecution.executedPlan + assert( + plan.collect { case a: CometHashAggregateExec => a.modes }.flatten == + Seq(Partial)) + assert(plan.collect { case a: HashAggregateExec => + a.aggregateExpressions.map(_.mode).distinct + }.flatten == Seq(Final)) + } + } + } + } + } + test("count, avg with null") { Seq(false, true).foreach { dictionary => withSQLConf("parquet.enable.dictionary" -> dictionary.toString) { @@ -649,7 +760,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) @@ -857,26 +971,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) } } } @@ -1341,20 +1458,24 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { // There is no sum to overflow when the input is empty or entirely null, so avg must // return null rather than raising ARITHMETIC_OVERFLOW in ANSI mode. // https://github.com/apache/datafusion-comet/issues/5148 - Seq(true, false).foreach { ansiEnabled => + for { + ansiEnabled <- Seq(true, false) + precision <- Seq(27, 38) + } { withSQLConf( SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString, CometConf.COMET_SHUFFLE_ENABLED.key -> "true", CometConf.COMET_SHUFFLE_MODE.key -> "native") { - val table = s"avg_decimal_empty_ansi_$ansiEnabled" + val table = s"avg_decimal_empty_${precision}_ansi_$ansiEnabled" + val nativeAggregates = if (precision == 38) 0 else 2 withTable(table) { - sql(s"create table $table(a decimal(38, 2), b INT) using parquet") + sql(s"create table $table(a decimal($precision, 2), b INT) using parquet") // no rows at all - checkSparkAnswer(s"select avg(a) from $table") + checkSparkAnswerAndNumOfAggregates(s"select avg(a) from $table", nativeAggregates) checkSparkAnswer(s"select b, avg(a) from $table group by b") // rows, but every value to average is null sql(s"insert into $table values(null, 1), (null, 2)") - checkSparkAnswer(s"select avg(a) from $table") + checkSparkAnswerAndNumOfAggregates(s"select avg(a) from $table", nativeAggregates) checkSparkAnswer(s"select b, avg(a) from $table group by b") } } @@ -1383,6 +1504,242 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } + test("high-precision global decimal AVG and TRY_AVG fall back to Spark") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "native") { + withTempDir { dir => + withTempView("high_precision_global_avg") { + Seq((1, 1, "0.6"), (1, 2, "0.6"), (2, 1, "0.6"), (2, 2, "0.5")) + .toDF("g", "ord", "raw_v") + .selectExpr( + "g", + "ord", + "CAST(raw_v AS DECIMAL(27,27)) AS v27", + "CAST(raw_v AS DECIMAL(28,28)) AS v28", + "CAST(raw_v AS DECIMAL(38,38)) AS v38") + .write + .mode("overwrite") + .parquet(dir.toString) + spark.read + .parquet(dir.toString) + .coalesce(1) + .createOrReplaceTempView("high_precision_global_avg") + + val fallbackReason = + "Ungrouped AVG on DECIMAL with maximum-precision intermediate state is not supported" + for { + ansiEnabled <- Seq(false, true) + (group, expected) <- Seq((1, "0.6"), (2, "0.55")) + aggregate <- Seq("AVG", "TRY_AVG") + } { + withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString) { + val query = + s"SELECT $aggregate(v38) FROM high_precision_global_avg WHERE g = $group" + checkAnswer(sql(query), Seq(Row(new java.math.BigDecimal(expected).setScale(38)))) + val (_, cometPlan) = checkSparkAnswerAndFallbackReason(query, fallbackReason) + assert(collect(cometPlan) { case agg: CometHashAggregateExec => agg }.isEmpty) + val sparkAggregates = collect(cometPlan) { case agg: HashAggregateExec => agg } + assert(sparkAggregates.size == 2) + // Both stages must remain adjacent so Spark can divide the wider sum before + // materializing a decimal buffer. Falling back only the final is unsafe. + assert(sparkAggregates.head.child.isInstanceOf[HashAggregateExec]) + } + } + + val (_, boundaryPlan) = checkSparkAnswerAndFallbackReason( + "SELECT AVG(v28) FROM high_precision_global_avg", + fallbackReason) + assert(collect(boundaryPlan) { case agg: CometHashAggregateExec => agg }.isEmpty) + checkSparkAnswerAndNumOfAggregates("SELECT AVG(v27) FROM high_precision_global_avg", 2) + + // Exercise grouped intermediate buffers in the distinct rewrite and the separate + // ObjectHashAggregate support check used when AVG shares a node with collect_list. + for (aggregates <- Seq( + "AVG(v38), COUNT(DISTINCT ord)", + "AVG(v38), sort_array(collect_list(ord))")) { + val (_, cometPlan) = checkSparkAnswerAndFallbackReason( + s"SELECT $aggregates FROM high_precision_global_avg WHERE g = 1", + fallbackReason) + assert(collect(cometPlan) { case agg: CometHashAggregateExec => agg }.isEmpty) + } + + // A shuffled final can also retain a wider sum after merging individually valid + // partials. Keep that shape in Spark as well, not just adjacent partial/final pairs. + spark.read + .parquet(dir.toString) + .repartition(2, col("ord")) + .createOrReplaceTempView("high_precision_global_avg") + val shuffledQuery = + "SELECT AVG(v38), TRY_AVG(v38) FROM high_precision_global_avg WHERE g = 1" + for (ansiEnabled <- Seq(false, true)) { + withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString) { + val expected = new java.math.BigDecimal("0.6").setScale(38) + checkAnswer(sql(shuffledQuery), Seq(Row(expected, expected))) + val (_, cometPlan) = + checkSparkAnswerAndFallbackReason(shuffledQuery, fallbackReason) + assert(collect(cometPlan) { case agg: CometHashAggregateExec => agg }.isEmpty) + } + } + } + } + } + } + + test("high-precision DISTINCT decimal AVG preserves Spark shuffle semantics") { + withSQLConf( + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_NATIVE_HASH_PARTITIONING_ENABLED.key -> "true") { + withTempDir { dir => + withTempView("high_precision_distinct_avg") { + Seq((1, "0.6"), (2, "0.6"), (3, "0.2"), (4, "0.3")) + .toDF("ord", "raw_v") + .selectExpr("ord", "CAST(raw_v AS DECIMAL(38,38)) AS v") + .write + .mode("overwrite") + .parquet(dir.toString) + spark.read + .parquet(dir.toString) + .repartition(2, col("ord")) + .createOrReplaceTempView("high_precision_distinct_avg") + + for { + shuffleMode <- Seq("native", "auto") + adaptiveEnabled <- Seq(false, true) + ansiEnabled <- Seq(false, true) + aggregate <- Seq("AVG", "TRY_AVG") + } { + withSQLConf( + CometConf.COMET_SHUFFLE_MODE.key -> shuffleMode, + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptiveEnabled.toString, + SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString) { + withClue( + s"mode=$shuffleMode, aqe=$adaptiveEnabled, ansi=$ansiEnabled, $aggregate: ") { + val df = sql(s"SELECT $aggregate(DISTINCT v) FROM high_precision_distinct_avg") + // Spark hashes all three distinct values to one partition. Materializing its + // partial sum overflows; different decimal hashing can hide that overflow. + if (ansiEnabled && aggregate == "AVG") { + checkSparkAnswerMaybeThrows(df) match { + case (Some(sparkExc), Some(cometExc)) => + assert(sparkExc.getMessage.contains("ARITHMETIC_OVERFLOW")) + assert(cometExc.getMessage.contains("ARITHMETIC_OVERFLOW")) + case _ => fail("Both Spark and Comet must report decimal AVG overflow") + } + } else { + checkSparkAnswer(df) + checkAnswer(df, Seq(Row(null))) + } + + val nativeDecimalHashShuffles = collect(df.queryExecution.executedPlan) { + case exchange: CometShuffleExchangeExec + if exchange.shuffleType == CometNativeShuffle && + (exchange.outputPartitioning match { + case HashPartitioning(expressions, _) => + expressions.exists(_.dataType match { + case d: DecimalType => d.precision > 18 + case _ => false + }) + case _ => false + }) => + exchange + } + assert(nativeDecimalHashShuffles.isEmpty) + } + } + } + } + } + } + } + + test("grouped decimal AVG preserves overflow across columnar shuffle") { + withSQLConf( + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_NATIVE_HASH_PARTITIONING_ENABLED.key -> "true") { + withTempDir { dir => + withTempView("grouped_decimal_avg") { + Seq((1, 0, "0.6"), (1, 0, "0.6"), (2, 2, "0.1"), (2, 2, "0.2"), (3, 4, null)) + .toDF("k", "part", "raw_v") + .selectExpr( + "k", + "part", + "CAST(k AS DECIMAL(19,0)) AS k19", + "CAST(k AS DECIMAL(38,0)) AS k38", + "CAST(raw_v AS DECIMAL(38,38)) AS v") + .coalesce(1) + .write + .mode("overwrite") + .parquet(dir.toString) + spark.read + .parquet(dir.toString) + .repartition(2, col("part")) + .createOrReplaceTempView("grouped_decimal_avg") + + val validAverage = new java.math.BigDecimal("0.15").setScale(38) + val validGroups = Seq( + Row(new java.math.BigDecimal("2"), validAverage), + Row(new java.math.BigDecimal("3"), null)) + for { + shuffleMode <- Seq("native", "auto") + adaptiveEnabled <- Seq(false, true) + ansiEnabled <- Seq(false, true) + precision <- Seq(19, 38) + aggregate <- Seq("AVG", "TRY_AVG") + } { + withSQLConf( + CometConf.COMET_SHUFFLE_MODE.key -> shuffleMode, + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptiveEnabled.toString, + SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString) { + withClue(s"mode=$shuffleMode, aqe=$adaptiveEnabled, ansi=$ansiEnabled, " + + s"key precision=$precision, $aggregate: ") { + val df = sql(s"SELECT k$precision, $aggregate(v) FROM grouped_decimal_avg " + + s"GROUP BY k$precision") + // Both 0.6 values reach one partial aggregate. Its overflowed state has a + // null count whose payload is zeroed by the columnar shuffle's row conversion. + if (ansiEnabled && aggregate == "AVG") { + checkSparkAnswerMaybeThrows(df) match { + case (Some(sparkExc), Some(cometExc)) => + assert(sparkExc.getMessage.contains("ARITHMETIC_OVERFLOW")) + assert(cometExc.getMessage.contains("ARITHMETIC_OVERFLOW")) + case _ => + fail("Both Spark and Comet must report grouped decimal AVG overflow") + } + // A zero count in a valid all-null group must still produce NULL, not an error. + val controls = sql( + s"SELECT k$precision, AVG(v) FROM grouped_decimal_avg WHERE k >= 2 " + + s"GROUP BY k$precision") + checkSparkAnswer(controls) + checkAnswer(controls, validGroups) + } else { + checkSparkAnswer(df) + checkAnswer(df, Row(new java.math.BigDecimal("1"), null) +: validGroups) + } + + val plan = df.queryExecution.executedPlan + val nativeAggregates = collect(plan) { case agg: CometHashAggregateExec => agg } + assert(nativeAggregates.size == (if (shuffleMode == "auto") 2 else 0)) + if (shuffleMode == "auto") { + assert(collect(plan) { + case exchange: CometShuffleExchangeExec + if exchange.shuffleType == CometColumnarShuffle => + exchange + }.nonEmpty) + } + } + } + } + } + } + } + } + test("final decimal avg") { withSQLConf( CometConf.COMET_SHUFFLE_ENABLED.key -> "true", @@ -1404,13 +1761,13 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { sql(s"insert into $table values(0.13344406545919155429936259114971302408, 5)") checkSparkAnswerAndNumOfAggregates(s"SELECT b , AVG(a) FROM $table GROUP BY b", 2) - checkSparkAnswerAndNumOfAggregates(s"SELECT AVG(a) FROM $table", 2) + checkSparkAnswerAndNumOfAggregates(s"SELECT AVG(a) FROM $table", 0) checkSparkAnswerAndNumOfAggregates( s"SELECT b, MIN(a), MAX(a), COUNT(a), SUM(a), AVG(a) FROM $table GROUP BY b", 2) checkSparkAnswerAndNumOfAggregates( s"SELECT MIN(a), MAX(a), COUNT(a), SUM(a), AVG(a) FROM $table", - 2) + 0) } } } @@ -1461,7 +1818,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 +1828,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 +1841,9 @@ 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) + checkSparkAnswerAndNumOfAggregates("SELECT AVG(_9) FROM tbl", 0) } } } diff --git a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala index cd52f579559..f8cf72d5938 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala @@ -31,6 +31,7 @@ import org.apache.spark.sql.{CometTestBase, DataFrame, Dataset, Row} import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.functions.{col, count, sum} +import org.apache.spark.sql.types.DecimalType import org.apache.comet.CometConf @@ -72,7 +73,14 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper val shuffled = df .select($"_1") .repartition(10, col(c)) - checkShuffleAnswer(shuffled, 1, checkNativeOperators = true) + val nativeHashSupported = df.schema(c).dataType match { + case d: DecimalType => d.precision <= 18 + case _ => true + } + checkShuffleAnswer( + shuffled, + if (nativeHashSupported) 1 else 0, + checkNativeOperators = nativeHashSupported) } } } diff --git a/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala index 7cb6816298d..414ef81ced5 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala @@ -453,6 +453,56 @@ class CometWindowExecSuite extends CometTestBase { } } + test("window: high-precision decimal AVG and TRY_AVG fall back to Spark") { + withTempDir { dir => + Seq((1, 1, "0.6"), (1, 2, "0.6"), (2, 1, "0.6"), (2, 2, "0.5")) + .toDF("g", "ord", "raw_v") + .selectExpr( + "g", + "ord", + "CAST(raw_v AS DECIMAL(27,27)) AS v27", + "CAST(raw_v AS DECIMAL(28,28)) AS v28", + "CAST(raw_v AS DECIMAL(38,38)) AS v38") + .repartition(1) + .write + .mode("overwrite") + .parquet(dir.toString) + + spark.read.parquet(dir.toString).createOrReplaceTempView("high_precision_dec_avg") + val fallbackReason = + "AVG on DECIMAL with maximum-precision intermediate state is not supported" + def runningAverage(aggregate: String, column: String) = sql(s""" + SELECT g, ord, + $aggregate($column) OVER ( + PARTITION BY g + ORDER BY ord + ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW + ) AS running_avg + FROM high_precision_dec_avg + ORDER BY g, ord + """) + + for { + ansiEnabled <- Seq(false, true) + aggregate <- Seq("AVG", "TRY_AVG") + } { + withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString) { + val (_, cometPlan) = + checkSparkAnswerAndFallbackReason(runningAverage(aggregate, "v38"), fallbackReason) + assert(collect(cometPlan) { case window: SparkWindowExec => window }.nonEmpty) + assert(collect(cometPlan) { case window: CometWindowExec => window }.isEmpty) + } + } + + val (_, maxPrecisionPlan) = + checkSparkAnswerAndFallbackReason(runningAverage("AVG", "v28"), fallbackReason) + assert(collect(maxPrecisionPlan) { case window: SparkWindowExec => window }.nonEmpty) + + val (_, lowerPrecisionPlan) = checkSparkAnswerAndOperator(runningAverage("AVG", "v27")) + assertCometWindowExecExists(lowerPrecisionPlan) + } + } + test("window: decimal AVG fuzz with PARTITION BY and ORDER BY") { Seq((9, 1), (9, 4), (10, 2), (11, 3), (11, 6)).foreach { case (precision, scale) => withTempDir { dir => 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")