fix: preserve Spark AVG buffer and overflow semantics

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

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

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

@@ -146,8 +146,11 @@ pub struct AvgAccumulator {
 
 impl Accumulator for AvgAccumulator {
     fn state(&mut self) -> Result<Vec<ScalarValue>> {
+        // 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::<T>()
     }
 }
+
+#[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::<Result<Vec<_>>>()?;
+                final_acc.merge_batch(&state)?;
+            }
+            assert_eq!(
+                final_acc.evaluate()?,
+                ScalarValue::Float64(nonempty.then_some(3.0))
+            );
+        }
+        Ok(())
+    }
+}