From 422649b6e7590b1743729c7afaf9921977eade5e Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Wed, 23 Sep 2026 22:21:01 +0000 Subject: [PATCH 1/4] fix: preserve Spark AVG buffer and overflow semantics --- .../user-guide/latest/compatibility/index.md | 6 +- .../latest/compatibility/operators.md | 18 + native/spark-expr/src/agg_funcs/avg.rs | 36 +- .../spark-expr/src/agg_funcs/avg_decimal.rs | 332 +++++++++++------- .../org/apache/comet/serde/aggregates.scala | 8 +- .../spark/sql/comet/CometWindowExec.scala | 43 ++- .../apache/spark/sql/comet/operators.scala | 28 +- .../comet/exec/CometAggregateSuite.scala | 243 ++++++++++--- .../comet/exec/CometWindowExecSuite.scala | 44 +++ .../comet/rules/CometExecRuleSuite.scala | 18 +- .../CometCelebornShufflePlanningSuite.scala | 9 +- 11 files changed, 567 insertions(+), 218 deletions(-) diff --git a/docs/source/user-guide/latest/compatibility/index.md b/docs/source/user-guide/latest/compatibility/index.md index 017140b14dd..7c44219aef5 100644 --- a/docs/source/user-guide/latest/compatibility/index.md +++ b/docs/source/user-guide/latest/compatibility/index.md @@ -121,13 +121,17 @@ divergence: ## Known result-value divergences -The following native paths silently return values that differ from Spark for edge-case inputs. +The following native paths can return different values or raise errors that Spark does not raise +for edge-case inputs. Most also have entries in the per-category expression pages linked above; they are collected here so users hunting an unexpected value have a single place to check: - Native `RANGE` window frames with an explicit `PRECEDING` / `FOLLOWING` offset diverge from Spark when the boundary arithmetic overflows for `DATE` or `DECIMAL` `ORDER BY` columns ([#5022](https://github.com/apache/datafusion-comet/issues/5022)). +- Grouped ANSI decimal `AVG` can evaluate an overflowing group in a batch even when Spark would + stop before that group at `LIMIT`. It remains natively enabled with this error-timing difference; + disabling `spark.comet.exec.aggregate.enabled` restores Spark's aggregate evaluation. - Ungrouped decimal `SUM` keeps an unbounded intermediate and checks the result precision only when a partial is written out or the sum is evaluated, which matches Spark's whole-stage codegen path. Without codegen Spark buffers the aggregate in an `UnsafeRow` and latches as soon as a diff --git a/docs/source/user-guide/latest/compatibility/operators.md b/docs/source/user-guide/latest/compatibility/operators.md index 657a32de532..7c20cb2bc75 100644 --- a/docs/source/user-guide/latest/compatibility/operators.md +++ b/docs/source/user-guide/latest/compatibility/operators.md @@ -60,6 +60,22 @@ 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. +Narrower ungrouped decimal averages remain eligible for native execution. + +Grouped decimal `AVG` remains eligible for native execution. Native finalization can raise an +ANSI overflow for a group that Spark would never consume after `LIMIT`; this error-timing difference +is documented under [known divergences](index.md#known-result-value-divergences). +High-precision decimal `AVG` with object aggregation has a separate overflow difference tracked in +[#5509](https://github.com/apache/datafusion-comet/issues/5509). + ## Window Functions Comet runs `WindowExec` natively and it is enabled by default (`spark.comet.exec.window.enabled`). A broad set of @@ -86,6 +102,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)). - `RANGE` frame bounded by `CURRENT ROW` when an `ORDER BY` key is an array of arrays or structs, or a struct diff --git a/native/spark-expr/src/agg_funcs/avg.rs b/native/spark-expr/src/agg_funcs/avg.rs index bb5371568b5..2b2422ce775 100644 --- a/native/spark-expr/src/agg_funcs/avg.rs +++ b/native/spark-expr/src/agg_funcs/avg.rs @@ -140,14 +140,15 @@ impl AggregateUDFImpl for Avg { /// An accumulator to compute the average #[derive(Debug, Default)] pub struct AvgAccumulator { - sum: Option, + /// Spark's partial buffer uses zero for empty input; count determines result nulls. + sum: f64, count: i64, } impl Accumulator for AvgAccumulator { fn state(&mut self) -> Result> { Ok(vec![ - ScalarValue::Float64(self.sum), + ScalarValue::Float64(Some(self.sum)), ScalarValue::from(self.count), ]) } @@ -155,9 +156,8 @@ impl Accumulator for AvgAccumulator { fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { let values = values[0].as_primitive::(); self.count += (values.len() - values.null_count()) as i64; - let v = self.sum.get_or_insert(0.); if let Some(x) = sum(values) { - *v += x; + self.sum += x; } Ok(()) } @@ -168,21 +168,16 @@ impl Accumulator for AvgAccumulator { // sums are summed - no overflow checking in all Eval Modes if let Some(x) = sum(states[0].as_primitive::()) { - let v = self.sum.get_or_insert(0.); - *v += x; + self.sum += x; } Ok(()) } fn evaluate(&mut self) -> Result { if self.count == 0 { - // If all input are nulls, count will be 0, and we will get null after the division. - // This is consistent with Spark Average implementation. Ok(ScalarValue::Float64(None)) } else { - Ok(ScalarValue::Float64( - self.sum.map(|f| f / self.count as f64), - )) + Ok(ScalarValue::Float64(Some(self.sum / self.count as f64))) } } @@ -352,3 +347,22 @@ 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<()> { + let mut partial = AvgAccumulator::default(); + let empty = vec![ScalarValue::Float64(Some(0.0)), ScalarValue::Int64(Some(0))]; + assert_eq!(partial.state()?, empty); + assert_eq!(partial.evaluate()?, ScalarValue::Float64(None)); + + partial.update_batch(&[Arc::new(Float64Array::from(vec![None, None]))])?; + assert_eq!(partial.state()?, empty); + assert_eq!(partial.evaluate()?, ScalarValue::Float64(None)); + Ok(()) + } +} diff --git a/native/spark-expr/src/agg_funcs/avg_decimal.rs b/native/spark-expr/src/agg_funcs/avg_decimal.rs index c58dd213dec..74a7d572d0a 100644 --- a/native/spark-expr/src/agg_funcs/avg_decimal.rs +++ b/native/spark-expr/src/agg_funcs/avg_decimal.rs @@ -185,9 +185,11 @@ impl AggregateUDFImpl for AvgDecimal { fn default_value(&self, _data_type: &DataType) -> Result { match &self.result_data_type { - Decimal128(target_precision, target_scale) => { - Ok(make_decimal128(None, *target_precision, *target_scale)) - } + Decimal128(target_precision, target_scale) => Ok(ScalarValue::Decimal128( + None, + *target_precision, + *target_scale, + )), _ => not_impl_err!( "The result_data_type of AvgDecimal should be Decimal128 but got{}", self.result_data_type @@ -213,10 +215,9 @@ impl AggregateUDFImpl for AvgDecimal { /// An accumulator to compute the average for decimals #[derive(Debug)] struct AvgDecimalAccumulator { + /// Some(0) with a zero count is empty; None is a sticky overflow marker. sum: Option, count: i64, - is_empty: bool, - is_not_null: bool, sum_scale: i8, sum_precision: u8, target_precision: u8, @@ -237,10 +238,8 @@ impl AvgDecimalAccumulator { registry: Arc, ) -> Self { Self { - sum: None, + sum: Some(0), count: 0, - is_empty: true, - is_not_null: true, sum_scale, sum_precision, target_precision, @@ -264,38 +263,6 @@ impl AvgDecimalAccumulator { } datafusion::common::DataFusionError::from(error) } - - fn update_single(&mut self, values: &Decimal128Array, idx: usize) -> Result<()> { - let v = unsafe { values.value_unchecked(idx) }; - let (new_sum, is_overflow) = match self.sum { - Some(sum) => sum.overflowing_add(v), - None => (v, false), - }; - - if is_overflow || !Decimal128Type::is_valid_decimal_precision(new_sum, self.sum_precision) { - // Overflow: set to null. Error will be thrown during evaluate in ANSI mode. - // This matches Spark's DecimalAddNoOverflowCheck behavior. - self.is_not_null = false; - return Ok(()); - } - - self.sum = Some(new_sum); - - if let Some(new_count) = self.count.checked_add(1) { - self.count = new_count; - } else { - // Count overflow: set to null. Error will be thrown during evaluate in ANSI mode. - self.is_not_null = false; - return Ok(()); - } - - self.is_not_null = true; - Ok(()) - } -} - -fn make_decimal128(value: Option, precision: u8, scale: i8) -> ScalarValue { - ScalarValue::Decimal128(value, precision, scale) } impl Accumulator for AvgDecimalAccumulator { @@ -307,99 +274,63 @@ impl Accumulator for AvgDecimalAccumulator { } fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { - if !self.is_empty && !self.is_not_null { - // This means there's a overflow in decimal, so we will just skip the rest - // of the computation + if self.sum.is_none() { return Ok(()); } - let values = &values[0]; - let data = values.as_primitive::(); - self.is_empty = self.is_empty && values.len() == values.null_count(); - if values.null_count() == 0 { - for i in 0..data.len() { - self.update_single(data, i)?; - } - } else { - for i in 0..data.len() { - if data.is_null(i) { - continue; - } - self.update_single(data, i)?; - } + for value in values[0].as_primitive::().iter().flatten() { + self.sum = self + .sum + .and_then(|sum| sum.checked_add(value)) + .filter(|sum| Decimal128Type::is_valid_decimal_precision(*sum, self.sum_precision)); + self.count += 1; } Ok(()) } + /// Merge partial sum/count arrays, preserving null sums as overflow markers. + /// Empty partials carry zero sums; grouped overflow may also carry a null count. fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { let partial_sums = states[0].as_primitive::(); let partial_counts = states[1].as_primitive::(); - - // Update is_empty: if any partial state has data, we're not empty - if self.is_empty { - self.is_empty = partial_counts.len() == partial_counts.null_count(); - } - - // counts are summed self.count += sum(partial_counts).unwrap_or_default(); - // sums are summed - if let Some(x) = sum(partial_sums) { - let v = self.sum.get_or_insert(0); - let (result, overflowed) = v.overflowing_add(x); - - if overflowed || !Decimal128Type::is_valid_decimal_precision(result, self.sum_precision) - { - // Overflow during merge: set to null, error will be thrown during evaluate in ANSI mode - self.is_not_null = false; - self.sum = None; - } else { - *v = result; - } + if partial_sums.null_count() > 0 { + self.sum = None; + } else { + self.sum = self + .sum + .and_then(|value| value.checked_add(sum(partial_sums).unwrap_or_default())) + .filter(|value| { + Decimal128Type::is_valid_decimal_precision(*value, self.sum_precision) + }); } Ok(()) } 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 - { + if self.sum.is_none() && 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)); + if self.count == 0 { + return Ok(ScalarValue::Decimal128( + None, + self.target_precision, + self.target_scale, + )); } let scaler = 10_i128.pow(self.target_scale.saturating_sub(self.sum_scale) as u32); let target_min = MIN_DECIMAL128_FOR_EACH_PRECISION[self.target_precision as usize]; let target_max = MAX_DECIMAL128_FOR_EACH_PRECISION[self.target_precision as usize]; - let result = self .sum - .map(|v| avg(v, self.count as i128, target_min, target_max, scaler)); - - match result { - Some(value) => Ok(make_decimal128( - value, - self.target_precision, - self.target_scale, - )), - _ => Ok(make_decimal128( - None, - self.target_precision, - self.target_scale, - )), - } + .and_then(|sum| avg(sum, self.count as i128, target_min, target_max, scaler)); + Ok(ScalarValue::Decimal128( + result, + self.target_precision, + self.target_scale, + )) } fn size(&self) -> usize { @@ -481,7 +412,7 @@ impl AvgDecimalGroupsAccumulator { } #[inline] - fn update_single(&mut self, group_index: usize, value: i128) -> Result<()> { + fn update_single(&mut self, group_index: usize, value: i128) { let (new_sum, is_overflow) = self.sums[group_index].overflowing_add(value); self.counts[group_index] += 1; self.sums[group_index] = new_sum; @@ -493,7 +424,6 @@ impl AvgDecimalGroupsAccumulator { // This matches Spark's DecimalAddNoOverflowCheck behavior. self.is_not_null.set_bit(group_index, false); } - Ok(()) } } @@ -524,7 +454,7 @@ impl GroupsAccumulator for AvgDecimalGroupsAccumulator { let iter = group_indices.iter().zip(data.iter()); if opt_filter.is_none() && values.null_count() == 0 { for (&group_index, &value) in iter { - self.update_single(group_index, value)?; + self.update_single(group_index, value); } } else { for (idx, (&group_index, &value)) in iter.enumerate() { @@ -536,12 +466,15 @@ impl GroupsAccumulator for AvgDecimalGroupsAccumulator { if values.is_null(idx) { continue; } - self.update_single(group_index, value)?; + self.update_single(group_index, value); } } Ok(()) } + /// Merge parallel sum/count arrays into the indexed groups. Sum nulls also cover + /// null counts: native state shares their validity buffer and Spark counts are valid. + /// Keep overflow in the group state and defer ANSI errors until evaluation. fn merge_batch( &mut self, values: &[ArrayRef], @@ -578,23 +511,14 @@ impl GroupsAccumulator for AvgDecimalGroupsAccumulator { if is_overflow || !Decimal128Type::is_valid_decimal_precision(new_sum, self.sum_precision) { - if self.eval_mode == EvalMode::Ansi { - let error = decimal_sum_overflow_error("avg"); - return Err(self.wrap_error_with_context(error)); - } + // Preserve overflow in intermediate/spilled state without throwing. Like + // Spark's decimal AVG, report ANSI overflow only when evaluate emits the + // affected group; a prefix emission may not need this group yet. self.is_not_null.set_bit(group_index, false); } else { self.sums[group_index] = new_sum; } } - if partial_counts.null_count() != 0 { - for (index, &group_index) in group_indices.iter().enumerate() { - if partial_counts.is_null(index) { - self.is_not_null.set_bit(group_index, false); - } - } - } - Ok(()) } @@ -612,9 +536,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)); } @@ -699,3 +623,159 @@ 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( + 38, + 38, + 38, + 38, + mode, + None, + crate::create_query_context_map(), + ) + } + + #[test] + fn empty_partial_state_matches_spark() -> Result<()> { + let mut partial = accumulator(EvalMode::Ansi); + let empty = vec![ + ScalarValue::Decimal128(Some(0), 38, 38), + ScalarValue::Int64(Some(0)), + ]; + assert_eq!(partial.state()?, empty); + let values = Decimal128Array::from(vec![None, None]).with_precision_and_scale(38, 38)?; + partial.update_batch(&[Arc::new(values)])?; + assert_eq!(partial.state()?, empty); + + let state = partial + .state()? + .into_iter() + .map(|value| value.to_array_of_size(1)) + .collect::>>()?; + let mut final_acc = accumulator(EvalMode::Ansi); + final_acc.merge_batch(&state)?; + assert_eq!(final_acc.evaluate()?, ScalarValue::Decimal128(None, 38, 38)); + + let values = Decimal128Array::from(vec![Some(100), None, Some(300)]) + .with_precision_and_scale(38, 38)?; + final_acc.update_batch(&[Arc::new(values)])?; + final_acc.merge_batch(&state)?; + assert_eq!( + final_acc.evaluate()?, + ScalarValue::Decimal128(Some(200), 38, 38) + ); + Ok(()) + } + + #[test] + fn scalar_overflow_stays_null() -> Result<()> { + for mode in [EvalMode::Ansi, EvalMode::Legacy] { + let max = 10_i128.pow(38) - 1; + let mut updated = accumulator(mode); + // A cancelling value cannot revive a sum that already overflowed. + let values = + Decimal128Array::from(vec![max, max, -max]).with_precision_and_scale(38, 38)?; + updated.update_batch(&[Arc::new(values)])?; + + let mut merged = accumulator(mode); + let sums = Decimal128Array::from(vec![6 * 10_i128.pow(37); 2]) + .with_precision_and_scale(38, 38)?; + merged.merge_batch(&[Arc::new(sums), Arc::new(Int64Array::from(vec![1, 1]))])?; + for mut acc in [updated, merged] { + assert_eq!(acc.state()?[0], ScalarValue::Decimal128(None, 38, 38)); + if mode == EvalMode::Ansi { + let error = acc.evaluate().unwrap_err().to_string(); + assert!(error.contains("ARITHMETIC_OVERFLOW"), "{error}"); + } else { + assert_eq!(acc.evaluate()?, ScalarValue::Decimal128(None, 38, 38)); + } + } + } + Ok(()) + } + + /// Grouped merge overflow survives a shuffle that zeroes null count payloads and + /// later updates. ANSI evaluation fails only when the affected group is emitted. + #[test] + fn grouped_merge_overflow_survives_roundtrip() -> Result<()> { + for mode in [EvalMode::Ansi, EvalMode::Legacy] { + 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(); + let large = 6 * 10_i128.pow(37); + // Group 0 is empty, group 1 is valid, and group 2 overflows during merge. + let sums = Decimal128Array::from(vec![0, 100, large, large]) + .with_precision_and_scale(38, 38)?; + partial.merge_batch( + &[Arc::new(sums), Arc::new(Int64Array::from(vec![0, 1, 1, 1]))], + &[0, 1, 2, 2], + 3, + )?; + let state = partial.state(EmitTo::All)?; + // Rebuilding logical values models a shuffle: null payloads become zero. + let sums = state[0] + .as_primitive::() + .iter() + .collect::() + .with_precision_and_scale(38, 38)?; + let counts = state[1] + .as_primitive::() + .iter() + .collect::(); + assert!(sums.is_null(2)); + assert!(counts.is_null(2)); + assert_eq!(counts.value(2), 0); + let state: Vec = vec![Arc::new(sums), Arc::new(counts)]; + + // A scalar final also treats a null sum/count as overflow, not empty input. + let mut scalar = accumulator(mode); + scalar.merge_batch(&[state[0].slice(2, 1), state[1].slice(2, 1)])?; + assert_eq!(scalar.state()?[0], ScalarValue::Decimal128(None, 38, 38)); + if mode == EvalMode::Ansi { + let error = scalar.evaluate().unwrap_err().to_string(); + assert!(error.contains("ARITHMETIC_OVERFLOW"), "{error}"); + } else { + assert_eq!(scalar.evaluate()?, ScalarValue::Decimal128(None, 38, 38)); + } + + let mut final_acc = new_acc(); + final_acc.merge_batch(&state, &[0, 1, 2], 3)?; + let prefix = final_acc.evaluate(EmitTo::First(2))?; + assert_eq!( + prefix + .as_primitive::() + .iter() + .collect::>(), + vec![None, Some(100)] + ); + // Prefix emission shifts the overflowing group to index 0. + let values = Decimal128Array::from(vec![-large]).with_precision_and_scale(38, 38)?; + final_acc.update_batch(&[Arc::new(values)], &[0], None, 1)?; + if mode == EvalMode::Ansi { + let error = final_acc.evaluate(EmitTo::All).unwrap_err().to_string(); + assert!(error.contains("ARITHMETIC_OVERFLOW"), "{error}"); + } else { + let result = final_acc.evaluate(EmitTo::All)?; + assert_eq!(result.len(), 1); + assert!(result.is_null(0)); + } + } + Ok(()) + } +} 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 1ed8af82d73..937da61c7a5 100644 --- a/spark/src/main/scala/org/apache/comet/serde/aggregates.scala +++ b/spark/src/main/scala/org/apache/comet/serde/aggregates.scala @@ -223,12 +223,12 @@ object CometCount extends CometAggregateExpressionSerde[Count] { object CometAverage extends CometAggregateExpressionSerde[Average] { - // Keep the default native-to-Spark restriction until #5420: an untouched native AVG emits - // (null, 0), but Spark's merge needs (0.0, 0). + // Non-decimal AVG emits Spark's (sum: double, count: long) buffer, including (0.0, 0) for + // empty partials. Decimal AVG nulls count differently on overflow and remains unsafe to mix. + override def supportsNativePartialToSparkFinal(fn: Average): Boolean = + !fn.child.dataType.isInstanceOf[DecimalType] 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] override def getUnsupportedReasons(): Seq[String] = Seq( 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 20de01f8c30..d8db3e5c7ea 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 @@ -355,24 +355,31 @@ object CometWindowExec extends CometOperatorSerde[WindowExec] { case _: SpecifiedWindowFrame => false case _ => true } - if (!isEverExpanding) { - windowExpr.windowFunction match { - case agg: AggregateExpression => - agg.aggregateFunction match { - case s: Sum if s.dataType.isInstanceOf[DecimalType] => - withFallbackReason( - windowExpr, - "SUM on DECIMAL with a sliding window frame is not supported") - return None - case a: Average if a.dataType.isInstanceOf[DecimalType] => - withFallbackReason( - windowExpr, - "AVG on DECIMAL with a sliding window frame is not supported") - return None - case _ => - } - case _ => - } + windowExpr.windowFunction match { + case AggregateExpression(a: Average, _, _, _, _) if isEverExpanding => + a.sumDataType match { + case decimal: DecimalType if decimal.precision == DecimalType.MAX_PRECISION => + // Spark can retain an out-of-precision window sum until AVG divides by count. + // Comet records overflow immediately, so the two accumulators can disagree. + withFallbackReason( + windowExpr, + "AVG on DECIMAL with maximum-precision intermediate state is not supported") + return None + case _ => + } + case AggregateExpression(s: Sum, _, _, _, _) + if !isEverExpanding && s.dataType.isInstanceOf[DecimalType] => + withFallbackReason( + windowExpr, + "SUM on DECIMAL with a sliding window frame is not supported") + return None + case AggregateExpression(a: Average, _, _, _, _) + if !isEverExpanding && a.dataType.isInstanceOf[DecimalType] => + withFallbackReason( + windowExpr, + "AVG on DECIMAL with a sliding window frame is not supported") + return None + case _ => } // Comet's native window planner ships RANGE frame offsets as 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 44aee93d3fd..b068436d78c 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, AttributeSeq, AttributeSet, CodegenObjectFactoryMode, Expression, ExpressionSet, Generator, LeafExpression, Literal, NamedExpression, SortOrder, XXH64} -import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, AggregateMode, CollectList, CollectSet, Final, ImperativeAggregate, Mode, Partial, PartialMerge, Percentile, Sum} +import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, AggregateMode, Average, CollectList, CollectSet, Final, ImperativeAggregate, Mode, Partial, PartialMerge, Percentile, Sum} import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} import org.apache.spark.sql.catalyst.plans._ @@ -1979,6 +1979,28 @@ trait CometBaseAggregate { case _ => false }) + 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 global buffer can retain a wider sum until division; native overflow is sticky. + // Both stages must fall back because decimal AVG buffers cannot cross engines. + Unsupported( + Some( + "Ungrouped AVG on DECIMAL with maximum-precision intermediate state is not supported")) + } else { + Compatible() + } + } + def doConvert( aggregate: BaseAggregateExec, builder: Operator.Builder, @@ -2367,7 +2389,7 @@ object CometHashAggregateExec "Ungrouped decimal SUM at maximum precision without codegen cannot match Spark's " + "latching UnsafeRow buffer")) } - Compatible() + aggregateSupportLevel(op) } override def convert( @@ -2435,7 +2457,7 @@ object CometObjectHashAggregateExec "Grouped decimal SUM at maximum precision cannot match Spark's unbounded object " + "aggregation buffer")) } - Compatible() + aggregateSupportLevel(op) } override def convert( 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 f0b37eef884..4b552467a05 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -32,11 +32,10 @@ import org.apache.spark.sql.catalyst.expressions.aggregate.{Final, Partial, Part import org.apache.spark.sql.catalyst.optimizer.EliminateSorts 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.LocalTableScanExec -import org.apache.spark.sql.execution.SQLExecution +import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometShuffleExchangeExec} +import org.apache.spark.sql.execution.{LocalTableScanExec, SQLExecution} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AdaptiveSparkPlanHelper, ShuffleQueryStageExec} -import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, ObjectHashAggregateExec} +import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec, ObjectHashAggregateExec} import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeExec} import org.apache.spark.sql.functions.{avg, col, collect_list, collect_set, count_distinct, expr, sort_array, sum} import org.apache.spark.sql.internal.SQLConf @@ -545,8 +544,8 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { 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. + // The filter leaves three input partitions empty. Decimal AVG remains unsafe to + // mix between engines because overflow nulls its count differently from Spark. 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")))) @@ -593,7 +592,8 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { } for (adaptive <- Seq(false, true)) { - test(s"COUNT and AVG fall back together across a Spark shuffle (AQE=$adaptive)") { + test( + s"COUNT and repaired AVG preserve native partials across a Spark shuffle (AQE=$adaptive)") { withTempDir { dir => val path = s"${dir.getAbsolutePath}/data" withSQLConf(CometConf.COMET_ENABLED.key -> "false") { @@ -614,29 +614,29 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { 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. + // Three partitions have no surviving rows, so AVG has no update_batch call. + // Its repaired (0.0, 0) state and COUNT's zero buffer are both safe for Spark Final. + // Keep Final enabled so the Spark shuffle determines the execution boundary. 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: CometHashAggregateExec if agg.modes == Seq(Partial) => agg + } + assert(partials.size == 1) + val finals = collect(plan) { case agg: BaseAggregateExec - if agg.aggregateExpressions.map(_.mode).distinct == Seq(Partial) => + if agg.aggregateExpressions.map(_.mode).distinct == Seq(Final) => agg } - assert(partials.size == 1) - assert(partials.head.getTagValue(CometExecRule.COMET_UNSAFE_PARTIAL).isDefined) + assert(finals.size == 1) 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. + // The fully native aggregate path remains enabled as well. val native = sql("SELECT COUNT(*), AVG(v) FROM count_avg_fallback WHERE id = 1") val initialNativePlan = stripAQEPlan(native.queryExecution.executedPlan) @@ -688,7 +688,7 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { 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. + // Integer keys isolate this from the separate wide-decimal shuffle fix in #6005. // 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" @@ -924,16 +924,15 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } - test("mixed engine sum/avg falls back when Spark Final would consume native AVG") { + test("mixed engine sum/avg: Comet partial + Spark final matches Spark") { 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") { - checkSparkAnswerAndNumOfAggregates( - "SELECT _4, SUM(_1), SUM(_2), SUM(_3), AVG(_1), AVG(_2), AVG(_3) FROM tbl GROUP BY _4", - 0) + checkSparkAnswer( + "SELECT _4, SUM(_1), SUM(_2), SUM(_3), AVG(_1), AVG(_2), AVG(_3) FROM tbl GROUP BY _4") } } } @@ -1150,6 +1149,47 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } + test("avg empty partial buffers cross engines in both directions") { + 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 -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "jvm") { + withTempPath { path => + spark + .range(0L, 8L, 1L, 4) + .selectExpr("id", "cast(id + 10 as decimal(7, 2)) as price") + .write + .parquet(path.getCanonicalPath) + withParquetTable(path.getCanonicalPath, "avg_empty_partials") { + assert(sql("SELECT * FROM avg_empty_partials").rdd.getNumPartitions == 4) + // One surviving partition and no surviving rows exercise empty buffers independently. + // Spark rewrites the narrow decimal AVG to AVG(UnscaledValue(price)). + for { + (disabled, nativeMode) <- Seq( + CometConf.COMET_ENABLE_FINAL_HASH_AGGREGATE -> Partial, + CometConf.COMET_ENABLE_PARTIAL_HASH_AGGREGATE -> Final) + predicate <- Seq("id = 1", "id < 0") + } { + withSQLConf(disabled.key -> "false") { + val df = sql(s"SELECT AVG(id), AVG(price) FROM avg_empty_partials WHERE $predicate") + checkSparkAnswer(df) + assert(collect(df.queryExecution.executedPlan) { case a: CometHashAggregateExec => + a.modes + }.flatten == Seq(nativeMode)) + assert(collect(df.queryExecution.executedPlan) { case a: HashAggregateExec => + a.aggregateExpressions.map(_.mode).distinct + }.flatten == Seq(if (nativeMode == Partial) Final else Partial)) + } + } + } + } + } + } + test("count, avg with null") { Seq(false, true).foreach { dictionary => withSQLConf("parquet.enable.dictionary" -> dictionary.toString) { @@ -1203,10 +1243,7 @@ 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) - val avgStages = if (nativeShuffleEnabled) 2 else 0 - checkSparkAnswerAndNumOfAggregates( - "SELECT _2, AVG(_1) FROM tbl GROUP BY _2", - avgStages) + checkSparkAnswerAndNumOfAggregates("SELECT _2, AVG(_1) FROM tbl GROUP BY _2", n) } } } @@ -2330,20 +2367,20 @@ 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)) { 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" withTable(table) { - sql(s"create table $table(a decimal(38, 2), b INT) using parquet") + sql(s"create table $table(a decimal(27, 2), b INT) using parquet") // no rows at all - checkSparkAnswer(s"select avg(a) from $table") + checkSparkAnswerAndNumOfAggregates(s"select avg(a) from $table", 2) 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", 2) checkSparkAnswer(s"select b, avg(a) from $table group by b") } } @@ -2372,6 +2409,127 @@ 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 => + Seq((1, "0.6"), (2, "0.6")) + .toDF("ord", "raw_v") + .selectExpr( + "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) + withParquetTable(spark.read.parquet(dir.toString).coalesce(1), "global_avg") { + // Precision 27 stays native; 28 starts fallback. At 38, the intermediate 1.2 + // overflows in Comet although Spark can divide it and return 0.6. + for ((ansi, aggregates, nativeCount) <- Seq( + (false, "AVG(v27)", 2), + (false, "AVG(v28)", 0), + (true, "AVG(v38)", 0), + (true, "TRY_AVG(v38)", 0), + (false, "AVG(v38), COUNT(DISTINCT ord)", 0), + (false, "AVG(v38), sort_array(collect_list(ord))", 0))) { + withSQLConf(SQLConf.ANSI_ENABLED.key -> ansi.toString) { + checkSparkAnswerAndNumOfAggregates( + s"SELECT $aggregates FROM global_avg", + nativeCount) + } + } + // Merging individually valid partials can overflow too; keep the shuffled final safe. + spark.read + .parquet(dir.toString) + .repartition(2, col("ord")) + .createOrReplaceTempView("global_avg") + checkSparkAnswerAndNumOfAggregates("SELECT AVG(v38), TRY_AVG(v38) FROM global_avg", 0) + } + } + } + } + + test("high-precision DISTINCT decimal AVG preserves Spark shuffle semantics") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_ENABLED.key -> "true", + // Native wide-decimal hash routing is tracked separately in #6005. + CometConf.COMET_SHUFFLE_MODE.key -> "jvm") { + withTempDir { dir => + 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) + withParquetTable( + spark.read.parquet(dir.toString).repartition(2, col("ord")), + "distinct_avg") { + // Spark hashes all three distinct values to one partition. Materializing its + // partial sum overflows, unlike the adjacent scalar stages exercised above. + val df = sql("SELECT AVG(DISTINCT v) FROM distinct_avg") + checkSparkAnswer(df) + checkAnswer(df, Seq(Row(null))) + } + } + } + } + + test("grouped decimal AVG preserves overflow across columnar shuffle") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + CometConf.COMET_SHUFFLE_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "jvm") { + withTempDir { dir => + 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(raw_v AS DECIMAL(38,38)) AS v") + .coalesce(1) + .write + .mode("overwrite") + .parquet(dir.toString) + withParquetTable( + spark.read.parquet(dir.toString).repartition(2, col("part")), + "grouped_avg") { + val expected = + Seq(Row(1, null), Row(2, new java.math.BigDecimal("0.15").setScale(38)), Row(3, null)) + for ((aggregate, ansi) <- Seq(("AVG", false), ("TRY_AVG", true), ("AVG", true))) { + withSQLConf(SQLConf.ANSI_ENABLED.key -> ansi.toString) { + val df = sql(s"SELECT k, $aggregate(v) FROM grouped_avg GROUP BY k") + // The overflowing partial exports a null count whose payload becomes zero + // during row conversion. It must remain distinct from the all-null group. + if (aggregate == "AVG" && ansi) { + val (sparkError, cometError) = checkSparkAnswerMaybeThrows(df) + assert(sparkError.exists(_.getMessage.contains("ARITHMETIC_OVERFLOW"))) + assert(cometError.exists(_.getMessage.contains("ARITHMETIC_OVERFLOW"))) + checkSparkAnswerAndNumOfAggregates( + "SELECT k, AVG(v) FROM grouped_avg WHERE k > 1 GROUP BY k", + 2) + } else { + checkSparkAnswer(df) + checkAnswer(df, expected) + } + val plan = df.queryExecution.executedPlan + assert(collect(plan) { case agg: CometHashAggregateExec => agg }.size == 2) + assert(collect(plan) { + case exchange: CometShuffleExchangeExec + if exchange.shuffleType == CometColumnarShuffle => + exchange + }.nonEmpty) + } + } + } + } + } + } + test("final decimal avg") { withSQLConf( CometConf.COMET_SHUFFLE_ENABLED.key -> "true", @@ -2393,27 +2551,27 @@ 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) } } } } } - test("AVG stays in Spark across a Spark shuffle") { + test("test partial avg") { 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", 0) + checkSparkAnswerAndNumOfAggregates("SELECT _2 , AVG(_1) FROM tbl GROUP BY _2", 1) } } } @@ -2450,9 +2608,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") { - // 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 + // 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", @@ -2460,11 +2618,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", @@ -2472,11 +2631,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/CometWindowExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala index 700748b6edc..f07eb3fe932 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala @@ -701,6 +701,50 @@ 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")) + .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") + .coalesce(1) + .write + .mode("overwrite") + .parquet(dir.toString) + withParquetTable(spark.read.parquet(dir.toString), "window_avg") { + withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") { + // Pin the 27/28 boundary and both error modes when the running sum exceeds 38 digits. + for ((aggregate, precision) <- Seq( + ("AVG", 27), + ("AVG", 28), + ("AVG", 38), + ("TRY_AVG", 38))) { + val df = sql(s""" + SELECT g, ord, $aggregate(v$precision) OVER ( + PARTITION BY g ORDER BY ord + ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) AS running_avg + FROM window_avg ORDER BY g, ord + """) + if (precision == 27) { + val (_, plan) = checkSparkAnswerAndOperator(df) + assertCometWindowExecExists(plan) + } else { + val (_, plan) = checkSparkAnswerAndFallbackReason( + df, + "AVG on DECIMAL with maximum-precision intermediate state is not supported") + assert(collect(plan) { case window: SparkWindowExec => window }.nonEmpty) + assert(collect(plan) { case window: CometWindowExec => window }.isEmpty) + } + } + } + } + } + } + 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 6e63f998302..7c48aa19c60 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -1153,7 +1153,7 @@ class CometExecRuleSuite extends CometTestBase { } } - test("CometExecRule should not allow AVG Comet partial and Spark final before buffer repair") { + test("CometExecRule should allow AVG mixed Comet partial and Spark final") { withTempView("test_data") { createTestDataFrame.createOrReplaceTempView("test_data") val sparkPlan = @@ -1163,9 +1163,8 @@ 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) - // 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) + assert(countOperators(transformedPlan, classOf[HashAggregateExec]) == 1) // final + assert(countOperators(transformedPlan, classOf[CometHashAggregateExec]) == 1) // partial } } } @@ -1217,7 +1216,8 @@ class CometExecRuleSuite extends CometTestBase { 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 "") + val aggregates = "AVG(CAST(id AS DECIMAL(20, 2)))" + + (if (distinct) ", SUM(DISTINCT id)" else "") for (fallback <- Seq("disabled hash partitioning", "prior shuffle fallback", "none")) { withSQLConf( @@ -1293,7 +1293,9 @@ class CometExecRuleSuite extends CometTestBase { withTempView("test_data") { createTestDataFrame.createOrReplaceTempView("test_data") val plan = applyCometExecRule( - createSparkPlan(spark, "SELECT AVG(id) FROM test_data GROUP BY (id % 3)")) + createSparkPlan( + spark, + "SELECT AVG(CAST(id AS DECIMAL(20, 2))) FROM test_data GROUP BY (id % 3)")) val partial = plan.collectFirst { case agg: CometHashAggregateExec if agg.modes == Seq(Partial) => agg }.get @@ -1327,7 +1329,9 @@ class CometExecRuleSuite extends CometTestBase { withTempView("test_data") { createTestDataFrame.createOrReplaceTempView("test_data") val plan = applyCometExecRule( - createSparkPlan(spark, "SELECT AVG(id) FROM test_data GROUP BY (id % 3)")) + createSparkPlan( + spark, + "SELECT AVG(CAST(id AS DECIMAL(20, 2))) FROM test_data GROUP BY (id % 3)")) val partial = plan.collectFirst { case agg: CometHashAggregateExec if agg.modes == Seq(Partial) => agg }.get 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 0e2b47e15ba..841c15636c5 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 @@ -604,15 +604,14 @@ class CometCelebornShufflePlanningSuite extends CometTestBase { val nativeAggregates = collect(executedPlan) { case aggregate: CometHashAggregateExec => aggregate } - if (function == "count") { - // COUNT's non-null Long buffer is safe for Spark Final to consume. + if (function == "count" || function == "avg") { + // COUNT and repaired non-decimal AVG emit Spark-compatible partial buffers. 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 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. + // state as BinaryType. Both halves must agree when the exchange falls back, + // not just when an aggregate operator itself is unsupported. assert(nativeAggregates.isEmpty, s"$executedPlan") } } From a2f1fe6dd39f1010fa082b13d6e66a35c91b7e97 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Thu, 24 Sep 2026 02:19:50 +0000 Subject: [PATCH 2/4] test: pin decimal AVG results and shuffle writers --- .../latest/compatibility/operators.md | 6 +- .../comet/exec/CometAggregateSuite.scala | 61 +++++++++++-------- 2 files changed, 41 insertions(+), 26 deletions(-) diff --git a/docs/source/user-guide/latest/compatibility/operators.md b/docs/source/user-guide/latest/compatibility/operators.md index 7c20cb2bc75..d006b642dcf 100644 --- a/docs/source/user-guide/latest/compatibility/operators.md +++ b/docs/source/user-guide/latest/compatibility/operators.md @@ -67,8 +67,10 @@ has reached precision 38 (`p >= 28`). Spark's generated aggregation code can ret 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. -Narrower ungrouped decimal averages remain eligible for native execution. +This fallback applies to the whole aggregate operator: a single wide-decimal `AVG` or `TRY_AVG` +also moves sibling aggregates such as `MIN`, `MAX`, `COUNT`, and `SUM` to Spark. Both partial and +final aggregate operators stay in Spark, including when a shuffle separates them. Narrower +ungrouped decimal averages remain eligible for native execution. Grouped decimal `AVG` remains eligible for native execution. Native finalization can raise an ANSI overflow for a group that Spark would never consume after `LIMIT`; this error-timing difference 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 4b552467a05..23f3e68a57b 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -32,7 +32,7 @@ import org.apache.spark.sql.catalyst.expressions.aggregate.{Final, Partial, Part import org.apache.spark.sql.catalyst.optimizer.EliminateSorts 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.{CometColumnarShuffle, CometShuffleExchangeExec} +import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometNativeShuffle, CometShuffleExchangeExec} import org.apache.spark.sql.execution.{LocalTableScanExec, SQLExecution} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AdaptiveSparkPlanHelper, ShuffleQueryStageExec} import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec, ObjectHashAggregateExec} @@ -2427,19 +2427,21 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { .mode("overwrite") .parquet(dir.toString) withParquetTable(spark.read.parquet(dir.toString).coalesce(1), "global_avg") { + val expectedAverage = new java.math.BigDecimal("0.6").setScale(38) // Precision 27 stays native; 28 starts fallback. At 38, the intermediate 1.2 // overflows in Comet although Spark can divide it and return 0.6. - for ((ansi, aggregates, nativeCount) <- Seq( - (false, "AVG(v27)", 2), - (false, "AVG(v28)", 0), - (true, "AVG(v38)", 0), - (true, "TRY_AVG(v38)", 0), - (false, "AVG(v38), COUNT(DISTINCT ord)", 0), - (false, "AVG(v38), sort_array(collect_list(ord))", 0))) { + for ((ansi, aggregates, nativeCount, expected) <- Seq( + (false, "AVG(v27)", 2, Row(expectedAverage.setScale(31))), + (false, "AVG(v28)", 0, Row(expectedAverage.setScale(32))), + (true, "AVG(v38)", 0, Row(expectedAverage)), + (true, "TRY_AVG(v38)", 0, Row(expectedAverage)), + (false, "AVG(v38), COUNT(DISTINCT ord)", 0, Row(expectedAverage, 2L)), + // Object aggregation overflows when materializing its partial buffer. + (false, "AVG(v38), sort_array(collect_list(ord))", 0, Row(null, Seq(1, 2))))) { withSQLConf(SQLConf.ANSI_ENABLED.key -> ansi.toString) { - checkSparkAnswerAndNumOfAggregates( - s"SELECT $aggregates FROM global_avg", - nativeCount) + val df = sql(s"SELECT $aggregates FROM global_avg") + checkAnswer(df, Seq(expected)) + assert(getNumCometHashAggregate(df) == nativeCount) } } // Merging individually valid partials can overflow too; keep the shuffled final safe. @@ -2447,7 +2449,9 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { .parquet(dir.toString) .repartition(2, col("ord")) .createOrReplaceTempView("global_avg") - checkSparkAnswerAndNumOfAggregates("SELECT AVG(v38), TRY_AVG(v38) FROM global_avg", 0) + val shuffled = sql("SELECT AVG(v38), TRY_AVG(v38) FROM global_avg") + checkAnswer(shuffled, Seq(Row(expectedAverage, expectedAverage))) + assert(getNumCometHashAggregate(shuffled) == 0) } } } @@ -2481,12 +2485,11 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } - test("grouped decimal AVG preserves overflow across columnar shuffle") { + test("grouped decimal AVG preserves overflow across JVM and native shuffle") { withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", SQLConf.SHUFFLE_PARTITIONS.key -> "2", - CometConf.COMET_SHUFFLE_ENABLED.key -> "true", - CometConf.COMET_SHUFFLE_MODE.key -> "jvm") { + CometConf.COMET_SHUFFLE_ENABLED.key -> "true") { withTempDir { dir => 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") @@ -2500,11 +2503,18 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { "grouped_avg") { val expected = Seq(Row(1, null), Row(2, new java.math.BigDecimal("0.15").setScale(38)), Row(3, null)) - for ((aggregate, ansi) <- Seq(("AVG", false), ("TRY_AVG", true), ("AVG", true))) { - withSQLConf(SQLConf.ANSI_ENABLED.key -> ansi.toString) { + for { + (shuffleMode, shuffleType) <- Seq( + ("jvm", CometColumnarShuffle), + ("native", CometNativeShuffle)) + (aggregate, ansi) <- Seq(("AVG", false), ("TRY_AVG", true), ("AVG", true)) + } { + withSQLConf( + CometConf.COMET_SHUFFLE_MODE.key -> shuffleMode, + SQLConf.ANSI_ENABLED.key -> ansi.toString) { val df = sql(s"SELECT k, $aggregate(v) FROM grouped_avg GROUP BY k") - // The overflowing partial exports a null count whose payload becomes zero - // during row conversion. It must remain distinct from the all-null group. + // The overflowing partial exports a null count whose payload may become zero + // across shuffle. It must remain distinct from the all-null group. if (aggregate == "AVG" && ansi) { val (sparkError, cometError) = checkSparkAnswerMaybeThrows(df) assert(sparkError.exists(_.getMessage.contains("ARITHMETIC_OVERFLOW"))) @@ -2518,11 +2528,14 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { } val plan = df.queryExecution.executedPlan assert(collect(plan) { case agg: CometHashAggregateExec => agg }.size == 2) - assert(collect(plan) { - case exchange: CometShuffleExchangeExec - if exchange.shuffleType == CometColumnarShuffle => - exchange - }.nonEmpty) + val bufferShuffles = collect(plan) { + case exchange: CometShuffleExchangeExec if (exchange.child match { + case agg: CometHashAggregateExec => agg.modes == Seq(Partial) + case _ => false + }) => + exchange.shuffleType + } + assert(bufferShuffles == Seq(shuffleType)) } } } From da3014dbc298e755e33c62e3d9bbc4f446694b00 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Sun, 4 Oct 2026 17:40:42 +0000 Subject: [PATCH 3/4] fix: address decimal AVG object buffers and wide division review --- .../latest/compatibility/operators.md | 7 +- .../spark-expr/src/agg_funcs/avg_decimal.rs | 148 +++++++++++++++--- .../apache/spark/sql/comet/operators.scala | 28 ++-- .../comet/exec/CometAggregateSuite.scala | 90 ++++++++++- 4 files changed, 239 insertions(+), 34 deletions(-) diff --git a/docs/source/user-guide/latest/compatibility/operators.md b/docs/source/user-guide/latest/compatibility/operators.md index d006b642dcf..b4c3a69049a 100644 --- a/docs/source/user-guide/latest/compatibility/operators.md +++ b/docs/source/user-guide/latest/compatibility/operators.md @@ -72,11 +72,12 @@ also moves sibling aggregates such as `MIN`, `MAX`, `COUNT`, and `SUM` to Spark. final aggregate operators stay in Spark, including when a shuffle separates them. Narrower ungrouped decimal averages remain eligible for native execution. -Grouped decimal `AVG` remains eligible for native execution. Native finalization can raise an +Grouped decimal `AVG` remains eligible for native execution with hash aggregation. With Spark's +object aggregation, decimal `AVG` and `TRY_AVG` fall back when their intermediate sum has precision +38, because Spark's object buffer can recover after an intermediate overflow. As with the ungrouped +case, the whole aggregate operator falls back. Native finalization can raise an ANSI overflow for a group that Spark would never consume after `LIMIT`; this error-timing difference is documented under [known divergences](index.md#known-result-value-divergences). -High-precision decimal `AVG` with object aggregation has a separate overflow difference tracked in -[#5509](https://github.com/apache/datafusion-comet/issues/5509). ## Window Functions diff --git a/native/spark-expr/src/agg_funcs/avg_decimal.rs b/native/spark-expr/src/agg_funcs/avg_decimal.rs index 74a7d572d0a..265547ff208 100644 --- a/native/spark-expr/src/agg_funcs/avg_decimal.rs +++ b/native/spark-expr/src/agg_funcs/avg_decimal.rs @@ -21,7 +21,7 @@ use arrow::array::{ types::{Decimal128Type, DecimalType, Int64Type}, Array, ArrayRef, Decimal128Array, Int64Array, PrimitiveArray, }; -use arrow::datatypes::{DataType, Field, FieldRef}; +use arrow::datatypes::{format_decimal_str, i256, DataType, Field, FieldRef}; use arrow::{array::BooleanBufferBuilder, buffer::NullBuffer, compute::sum}; use datafusion::common::{not_impl_err, Result, ScalarValue}; use datafusion::logical_expr::{ @@ -323,9 +323,20 @@ impl Accumulator for AvgDecimalAccumulator { let scaler = 10_i128.pow(self.target_scale.saturating_sub(self.sum_scale) as u32); let target_min = MIN_DECIMAL128_FOR_EACH_PRECISION[self.target_precision as usize]; let target_max = MAX_DECIMAL128_FOR_EACH_PRECISION[self.target_precision as usize]; - let result = self - .sum - .and_then(|sum| avg(sum, self.count as i128, target_min, target_max, scaler)); + let result = match self.sum { + Some(sum) => match avg(sum, self.count as i128, target_min, target_max, scaler) { + Ok(value) => Some(value), + Err(value) if self.eval_mode == EvalMode::Ansi => { + return Err(self.wrap_error_with_context(avg_result_overflow_error( + value, + self.target_precision, + self.target_scale, + ))); + } + Err(_) => None, + }, + None => None, + }; Ok(ScalarValue::Decimal128( result, self.target_precision, @@ -549,12 +560,15 @@ impl GroupsAccumulator for AvgDecimalGroupsAccumulator { } match avg(sum, count as i128, target_min, target_max, scaler) { - Some(value) => { - builder.append_value(value); - } - _ => { - builder.append_null(); + Ok(value) => builder.append_value(value), + Err(value) if self.eval_mode == EvalMode::Ansi => { + return Err(self.wrap_error_with_context(avg_result_overflow_error( + value, + self.target_precision, + self.target_scale, + ))); } + Err(_) => builder.append_null(), } } let array: PrimitiveArray = builder.finish(); @@ -595,7 +609,7 @@ impl GroupsAccumulator for AvgDecimalGroupsAccumulator { } /// Returns the `sum`/`count` as a i128 Decimal128 with -/// target_scale and target_precision and return None if overflows. +/// target_scale and target_precision, returning the rounded wide value on overflow. /// /// * sum: The total sum value stored as Decimal128 with sum_scale /// * count: total count, stored as a i128 (*NOT* a Decimal128 value) @@ -603,24 +617,52 @@ impl GroupsAccumulator for AvgDecimalGroupsAccumulator { /// * target_max: The maximum output value possible to represent with the target precision /// * scaler: scale factor for avg #[inline(always)] -fn avg(sum: i128, count: i128, target_min: i128, target_max: i128, scaler: i128) -> Option { - if let Some(value) = sum.checked_mul(scaler) { - // `sum / count` with ROUND_HALF_UP +fn avg( + sum: i128, + count: i128, + target_min: i128, + target_max: i128, + scaler: i128, +) -> std::result::Result { + // The scaled numerator can exceed i128 while both the sum and final average fit. + // Use the inexpensive path when it fits, and widen only that intermediate otherwise. + let rounded = if let Some(value) = sum.checked_mul(scaler) { let (div, rem) = value.div_rem(&count); let half = div_ceil(count, 2); - let half_neg = half.neg_wrapping(); - let new_value = match value >= 0 { + let rounded = match value >= 0 { true if rem >= half => div.add_wrapping(1), - false if rem <= half_neg => div.sub_wrapping(1), + false if rem <= half.neg_wrapping() => div.sub_wrapping(1), _ => div, }; - if new_value >= target_min && new_value <= target_max { - Some(new_value) + i256::from_i128(rounded) + } else { + // A Decimal128 numerator times the Decimal128 scale factor fits within i256. + let value = i256::from_i128(sum) * i256::from_i128(scaler); + let divisor = i256::from_i128(count); + let (div, rem) = (value / divisor, value % divisor); + let half = i256::from_i128(div_ceil(count, 2)); + if rem >= half { + div + i256::ONE + } else if rem <= -half { + div - i256::ONE } else { - None + div } + }; + if rounded >= i256::from_i128(target_min) && rounded <= i256::from_i128(target_max) { + Ok(rounded.to_i128().expect("bounded Decimal128 average")) } else { - None + Err(rounded) + } +} + +fn avg_result_overflow_error(value: i256, precision: u8, scale: i8) -> crate::SparkError { + let unscaled = value.to_string(); + let digits = unscaled.trim_start_matches('-').len(); + crate::SparkError::NumericValueOutOfRange { + value: format_decimal_str(&unscaled, digits, scale), + precision, + scale, } } @@ -640,6 +682,72 @@ mod tests { ) } + #[test] + fn decimal_average_scales_widely_and_reports_result_overflow() -> Result<()> { + for mode in [EvalMode::Legacy, EvalMode::Try, EvalMode::Ansi] { + for sign in [-1_i128, 1] { + for (value, count, expected) in [ + (9 * 10_i128.pow(33), 2, Some(sign * 9 * 10_i128.pow(37))), + (10_i128.pow(35), 1, None), + ] { + let values: ArrayRef = Arc::new( + Decimal128Array::from(vec![sign * value; count]) + .with_precision_and_scale(38, 18)?, + ); + let mut scalar = AvgDecimalAccumulator::new( + 18, + 38, + 38, + 22, + mode, + None, + crate::create_query_context_map(), + ); + scalar.update_batch(&[Arc::clone(&values)])?; + let mut grouped = AvgDecimalGroupsAccumulator::new( + &DataType::Decimal128(38, 22), + &DataType::Decimal128(38, 18), + 38, + 22, + 38, + 18, + mode, + None, + crate::create_query_context_map(), + ); + grouped.update_batch(&[values], &vec![0; count], None, 1)?; + let scalar_result = scalar.evaluate(); + let grouped_result = grouped.evaluate(EmitTo::All); + if expected.is_none() && mode == EvalMode::Ansi { + for error in [scalar_result.unwrap_err(), grouped_result.unwrap_err()] { + assert!( + error.to_string().contains("NUMERIC_VALUE_OUT_OF_RANGE"), + "{error}" + ); + } + } else { + assert_eq!(scalar_result?, ScalarValue::Decimal128(expected, 38, 22)); + assert_eq!( + grouped_result? + .as_primitive::() + .iter() + .collect::>(), + vec![expected] + ); + } + } + } + } + // HALF_UP rounding must agree for positive and negative values on the widened path. + let sum = 2 * 10_i128.pow(34) + 2; + let min = MIN_DECIMAL128_FOR_EACH_PRECISION[38]; + let max = MAX_DECIMAL128_FOR_EACH_PRECISION[38]; + let expected = 5 * 10_i128.pow(33) + 1; + assert_eq!(avg(sum, 40_000, min, max, 10_000), Ok(expected)); + assert_eq!(avg(-sum, 40_000, min, max, 10_000), Ok(-expected)); + Ok(()) + } + #[test] fn empty_partial_state_matches_spark() -> Result<()> { let mut partial = accumulator(EvalMode::Ansi); 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 b068436d78c..5458940ecc6 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 @@ -1979,16 +1979,18 @@ trait CometBaseAggregate { case _ => false }) + protected def hasMaxPrecisionDecimalAvg(op: BaseAggregateExec): Boolean = + op.aggregateExpressions.exists(_.aggregateFunction match { + case avg: Average => + avg.sumDataType match { + case decimal: DecimalType => decimal.precision == DecimalType.MAX_PRECISION + case _ => false + } + case _ => false + }) + 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 - }) + val unsupportedAverage = op.groupingExpressions.isEmpty && hasMaxPrecisionDecimalAvg(op) if (unsupportedAverage) { // Spark's global buffer can retain a wider sum until division; native overflow is sticky. @@ -2450,13 +2452,19 @@ object CometObjectHashAggregateExec } // Spark's object aggregation buffers the intermediate decimal sum unbounded, while Comet's // grouped accumulator latches to null once a running sum leaves the precision, so the - // grouped case declines. Decimal AVG has the same gap and is tracked separately. + // grouped case declines. Decimal AVG has the same intermediate-buffer gap. if (op.groupingExpressions.nonEmpty && hasMaxPrecisionDecimalSum(op)) { return Unsupported( Some( "Grouped decimal SUM at maximum precision cannot match Spark's unbounded object " + "aggregation buffer")) } + if (op.groupingExpressions.nonEmpty && hasMaxPrecisionDecimalAvg(op)) { + return Unsupported( + Some( + "Grouped decimal AVG at maximum precision cannot match Spark's unbounded object " + + "aggregation buffer")) + } aggregateSupportLevel(op) } 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 23f3e68a57b..561ac34ac7f 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -2485,13 +2485,101 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } + test("grouped decimal AVG widens division and reports final result overflow") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "native") { + withTempPath { dir => + Seq((1, "9000000000000000"), (1, "9000000000000000"), (2, "100000000000000000")) + .toDF("k", "raw_v") + .selectExpr("k", "CAST(raw_v AS DECIMAL(38,18)) AS v") + .coalesce(1) + .write + .parquet(dir.toString) + withParquetTable(dir.toString, "avg_wide_division") { + for (ansi <- Seq(false, true)) { + withSQLConf(SQLConf.ANSI_ENABLED.key -> ansi.toString) { + val valid = sql("SELECT k, AVG(v) FROM avg_wide_division WHERE k = 1 GROUP BY k") + checkSparkAnswer(valid) + assert(collect(valid.queryExecution.executedPlan) { + case agg: CometHashAggregateExec => agg + }.size == 2) + checkAnswer( + valid, + Seq(Row(1, new java.math.BigDecimal("9000000000000000").setScale(22)))) + val overflow = sql("SELECT k, AVG(v) FROM avg_wide_division WHERE k = 2 GROUP BY k") + if (ansi) { + val (sparkError, cometError) = checkSparkAnswerMaybeThrows(overflow) + assert(sparkError.exists(_.getMessage.contains("NUMERIC_VALUE_OUT_OF_RANGE"))) + assert(cometError.exists(_.getMessage.contains("NUMERIC_VALUE_OUT_OF_RANGE"))) + } else { + checkSparkAnswer(overflow) + assert(collect(overflow.queryExecution.executedPlan) { + case agg: CometHashAggregateExec => agg + }.size == 2) + checkAnswer(overflow, Seq(Row(2, null))) + } + val attempted = sql("SELECT k, TRY_AVG(v) FROM avg_wide_division GROUP BY k") + checkSparkAnswer(attempted) + assert(collect(attempted.queryExecution.executedPlan) { + case agg: CometHashAggregateExec => agg + }.size == 2) + checkAnswer( + attempted, + Seq( + Row(1, new java.math.BigDecimal("9000000000000000").setScale(22)), + Row(2, null))) + } + } + } + } + } + } + + test("decimal AVG with object hash aggregate falls back at maximum precision") { + for (ansi <- Seq(false, true)) { + withSQLConf( + SQLConf.ANSI_ENABLED.key -> ansi.toString, + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.USE_OBJECT_HASH_AGG.key -> "true", + CometConf.COMET_SHUFFLE_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "native") { + withMaxPrecisionDecimalTable(Seq((1, "0.6"), (1, "0.6"), (1, "-0.4")), "avg_object") { + val df = sql("SELECT k, AVG(v), sort_array(collect_list(k)) FROM avg_object GROUP BY k") + assert(df.queryExecution.sparkPlan.exists(_.isInstanceOf[ObjectHashAggregateExec])) + checkSparkAnswerAndFallbackReason( + df, + "Grouped decimal AVG at maximum precision cannot match Spark's unbounded object " + + "aggregation buffer") + checkAnswer( + df, + Seq( + Row( + 1, + new java.math.BigDecimal("0.26666666666666666666666666666666666667"), + Seq(1, 1, 1)))) + assert(collect(df.queryExecution.executedPlan) { case agg: CometHashAggregateExec => + agg + }.isEmpty) + } + } + } + } + test("grouped decimal AVG preserves overflow across JVM and native shuffle") { withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", SQLConf.SHUFFLE_PARTITIONS.key -> "2", CometConf.COMET_SHUFFLE_ENABLED.key -> "true") { withTempDir { dir => - Seq((1, 0, "0.6"), (1, 0, "0.6"), (2, 2, "0.1"), (2, 2, "0.2"), (3, 4, null)) + Seq( + (1, 0, "0.6"), + (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(raw_v AS DECIMAL(38,38)) AS v") .coalesce(1) From 86f25149fb21a472a67d7fa3c0d929fc2931f6c5 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Sun, 4 Oct 2026 19:28:27 +0000 Subject: [PATCH 4/4] Keep AQE fallback fixtures on unsafe decimal AVG buffers --- .../org/apache/comet/rules/CometExecRuleSuite.scala | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) 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 7c48aa19c60..82c9e6ae790 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -200,7 +200,9 @@ class CometExecRuleSuite extends CometTestBase { withTempView("test_data") { createTestDataFrame.createOrReplaceTempView("test_data") val original = - createSparkPlan(spark, "SELECT AVG(id) FROM test_data GROUP BY (id % 3)") + createSparkPlan( + spark, + "SELECT AVG(CAST(id AS DECIMAL(20, 2))) FROM test_data GROUP BY (id % 3)") val partial = original.collectFirst { case agg: HashAggregateExec if agg.aggregateExpressions.forall(_.mode == Partial) => agg }.get @@ -215,7 +217,7 @@ class CometExecRuleSuite extends CometTestBase { } // The Partial initially converts, but its upper hash shuffle cannot. Repair restores - // Spark's AVG buffer producer and revisits the already wrapped round-robin input stage. + // Spark's decimal AVG buffer producer and revisits the wrapped round-robin input stage. val result = applyCometExecRule(staged) assert(result.collect { case agg: HashAggregateExec => agg }.size == 2) assert(!result.exists(_.isInstanceOf[CometHashAggregateExec])) @@ -245,11 +247,13 @@ class CometExecRuleSuite extends CometTestBase { val data = Seq((1, Some(10)), (1, Some(20)), (2, None), (2, Some(40)), (3, None)) withParquetTable(data, "test_data") { val df = sql(""" - |SELECT k, AVG(v) FROM + |SELECT k, AVG(CAST(v AS DECIMAL(20, 2))) FROM | (SELECT /*+ REPARTITION(4) */ _1 AS k, _2 AS v FROM test_data) |GROUP BY k |""".stripMargin) - QueryTest.checkAnswer(df, Seq(Row(1, 15.0), Row(2, 40.0), Row(3, null))) + QueryTest.checkAnswer( + df, + Seq(Row(1, BigDecimal("15.000000")), Row(2, BigDecimal("40.000000")), Row(3, null))) val plan = df.queryExecution.executedPlan assert(plan.asInstanceOf[AdaptiveSparkPlanExec].isFinalPlan) val nativeShuffles = collect(plan) { case s: CometShuffleExchangeExec => s }