fix: preserve Spark AVG buffer and overflow semantics

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

docs/source/user-guide/latest/compatibility/operators.md

@@ -60,6 +60,25 @@ 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.
+
+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 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).
+
 ## Window Functions
 
 Comet runs `WindowExec` natively and it is enabled by default (`spark.comet.exec.window.enabled`). A broad set of
@@ -86,6 +105,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

native/spark-expr/src/agg_funcs/avg.rs

@@ -140,24 +140,24 @@ impl AggregateUDFImpl for Avg {
 /// An accumulator to compute the average
 #[derive(Debug, Default)]
 pub struct AvgAccumulator {
-    sum: Option<f64>,
+    /// 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<Vec<ScalarValue>> {
         Ok(vec![
-            ScalarValue::Float64(self.sum),
+            ScalarValue::Float64(Some(self.sum)),
             ScalarValue::from(self.count),
         ])
     }
 
     fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
         let values = values[0].as_primitive::<Float64Type>();
         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::<Float64Type>()) {
-            let v = self.sum.get_or_insert(0.);
-            *v += x;
+            self.sum += x;
         }
         Ok(())
     }
 
     fn evaluate(&mut self) -> Result<ScalarValue> {
         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::<T>()
     }
 }
+
+#[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(())
+    }
+}