AVG DISTINCT support

Signed-off-by: Adam Gutglick <adamgsal@gmail.com>

Author: AdamGS

datafusion/functions-aggregate-common/src/aggregate/avg_distinct/decimal.rs

@@ -16,48 +16,61 @@
 // under the License.
 
 use arrow::{
-    array::{ArrayRef, ArrowNumericType},
-    datatypes::{
-        Decimal32Type, Decimal64Type, Decimal128Type, Decimal256Type, DecimalType, i256,
-    },
+    array::{ArrayRef, ArrowNativeTypeOp, ArrowNumericType},
+    compute::DecimalCast,
+    datatypes::{ArrowNativeType, DecimalType},
 };
-use datafusion_common::{Result, ScalarValue};
+use datafusion_common::{Result, ScalarValue, exec_datafusion_err, exec_err};
 use datafusion_expr_common::accumulator::Accumulator;
 use std::fmt::Debug;
+use std::marker::PhantomData;
 use std::mem::size_of_val;
 
 use crate::aggregate::sum_distinct::DistinctSumAccumulator;
 use crate::utils::DecimalAverager;
 
 /// Generic implementation of `AVG DISTINCT` for Decimal types.
 /// Handles both all Arrow decimal types (32, 64, 128 and 256 bits).
+///
+/// The distinct values are stored in the input type `I`; only the intermediate
+/// sum is computed in the (never narrower) sum type `S` so it cannot overflow
+/// `I`'s native type.
 #[derive(Debug)]
-pub struct DecimalDistinctAvgAccumulator<T: DecimalType + Debug> {
-    sum_accumulator: DistinctSumAccumulator<T>,
+pub struct DecimalDistinctAvgAccumulator<
+    I: DecimalType + Debug,
+    S: DecimalType + Debug = I,
+> {
+    sum_accumulator: DistinctSumAccumulator<I>,
     sum_scale: i8,
     target_precision: u8,
     target_scale: i8,
+    _sum_type: PhantomData<S>,
 }
 
-impl<T: DecimalType + Debug> DecimalDistinctAvgAccumulator<T> {
+impl<I: DecimalType + Debug, S: DecimalType + Debug> DecimalDistinctAvgAccumulator<I, S> {
     pub fn with_decimal_params(
         sum_scale: i8,
         target_precision: u8,
         target_scale: i8,
     ) -> Self {
-        let data_type = T::TYPE_CONSTRUCTOR(T::MAX_PRECISION, sum_scale);
+        let data_type = I::TYPE_CONSTRUCTOR(I::MAX_PRECISION, sum_scale);
 
         Self {
             sum_accumulator: DistinctSumAccumulator::new(&data_type),
             sum_scale,
             target_precision,
             target_scale,
+            _sum_type: PhantomData,
         }
     }
 }
 
-impl<T: DecimalType + ArrowNumericType + Debug> Accumulator
-    for DecimalDistinctAvgAccumulator<T>
+impl<I, S> Accumulator for DecimalDistinctAvgAccumulator<I, S>
+where
+    I: DecimalType + ArrowNumericType + Debug,
+    S: DecimalType + ArrowNumericType + Debug,
+    I::Native: Into<S::Native> + DecimalCast,
+    S::Native: DecimalCast,
 {
     fn state(&mut self) -> Result<Vec<ScalarValue>> {
         self.sum_accumulator.state()
@@ -72,78 +85,43 @@ impl<T: DecimalType + ArrowNumericType + Debug> Accumulator
     }
 
     fn evaluate(&mut self) -> Result<ScalarValue> {
-        if self.sum_accumulator.distinct_count() == 0 {
-            return ScalarValue::new_primitive::<T>(
-                None,
-                &T::TYPE_CONSTRUCTOR(self.target_precision, self.target_scale),
-            );
+        let out_type = I::TYPE_CONSTRUCTOR(self.target_precision, self.target_scale);
+        let count = self.sum_accumulator.distinct_count();
+        if count == 0 {
+            return ScalarValue::new_primitive::<I>(None, &out_type);
         }
 
-        let sum_scalar = self.sum_accumulator.evaluate()?;
-
-        match sum_scalar {
-            ScalarValue::Decimal32(Some(sum), _, _) => {
-                let decimal_averager = DecimalAverager::<Decimal32Type>::try_new(
-                    self.sum_scale,
-                    self.target_precision,
-                    self.target_scale,
-                )?;
-                let avg = decimal_averager
-                    .avg(sum, self.sum_accumulator.distinct_count() as i32)?;
-                Ok(ScalarValue::Decimal32(
-                    Some(avg),
-                    self.target_precision,
-                    self.target_scale,
-                ))
-            }
-            ScalarValue::Decimal64(Some(sum), _, _) => {
-                let decimal_averager = DecimalAverager::<Decimal64Type>::try_new(
-                    self.sum_scale,
-                    self.target_precision,
-                    self.target_scale,
-                )?;
-                let avg = decimal_averager
-                    .avg(sum, self.sum_accumulator.distinct_count() as i64)?;
-                Ok(ScalarValue::Decimal64(
-                    Some(avg),
-                    self.target_precision,
-                    self.target_scale,
-                ))
-            }
-            ScalarValue::Decimal128(Some(sum), _, _) => {
-                let decimal_averager = DecimalAverager::<Decimal128Type>::try_new(
-                    self.sum_scale,
-                    self.target_precision,
-                    self.target_scale,
-                )?;
-                let avg = decimal_averager
-                    .avg(sum, self.sum_accumulator.distinct_count() as i128)?;
-                Ok(ScalarValue::Decimal128(
-                    Some(avg),
-                    self.target_precision,
-                    self.target_scale,
-                ))
-            }
-            ScalarValue::Decimal256(Some(sum), _, _) => {
-                let decimal_averager = DecimalAverager::<Decimal256Type>::try_new(
-                    self.sum_scale,
-                    self.target_precision,
-                    self.target_scale,
-                )?;
-                // `distinct_count` returns `u64`, but `avg` expects `i256`
-                // first convert `u64` to `i128`, then convert `i128` to `i256` to avoid overflow
-                let distinct_cnt: i128 = self.sum_accumulator.distinct_count() as i128;
-                let count: i256 = i256::from_i128(distinct_cnt);
-                let avg = decimal_averager.avg(sum, count)?;
-                Ok(ScalarValue::Decimal256(
-                    Some(avg),
-                    self.target_precision,
-                    self.target_scale,
-                ))
-            }
-
-            _ => unreachable!("Unsupported decimal type: {:?}", sum_scalar),
+        // Sum the distinct input values in the wider `S` so the total cannot
+        // overflow the input's native width (mirrors the non-distinct path).
+        let mut sum = S::Native::usize_as(0);
+        for value in self.sum_accumulator.distinct_values() {
+            sum = sum.add_wrapping(value.into());
         }
+
+        let Some(count) = S::Native::from_usize(count) else {
+            return exec_err!(
+                "Arithmetic overflow in avg: the distinct count {count} cannot \
+                 be represented in the sum type"
+            );
+        };
+
+        let averager = DecimalAverager::<S>::try_new(
+            self.sum_scale,
+            self.target_precision,
+            self.target_scale,
+        )?;
+        // Narrowing the average back to the (never wider) output type cannot
+        // fail in practice: `DecimalAverager::avg` validates the average
+        // against the output precision, whose bound fits the output's native
+        // type by construction
+        let avg =
+            I::Native::from_decimal(averager.avg(sum, count)?).ok_or_else(|| {
+                exec_datafusion_err!(
+                    "Arithmetic overflow in avg: the computed average does not fit \
+                 the output type"
+                )
+            })?;
+        ScalarValue::new_primitive::<I>(Some(avg), &out_type)
     }
 
     fn size(&self) -> usize {
@@ -160,6 +138,9 @@ mod tests {
     use arrow::array::{
         Decimal32Array, Decimal64Array, Decimal128Array, Decimal256Array,
     };
+    use arrow::datatypes::{
+        Decimal32Type, Decimal64Type, Decimal128Type, Decimal256Type, i256,
+    };
     use std::sync::Arc;
 
     #[test]
@@ -279,4 +260,94 @@ mod tests {
 
         Ok(())
     }
+
+    // The overflow regression tests below use odd-count ranges symmetric
+    // around a center value, so the exact sum is `count * center` and the
+    // average is exactly `center`.
+
+    #[test]
+    fn test_decimal32_distinct_avg_widens_to_decimal64() -> Result<()> {
+        // 42951 distinct values centered on 50000:
+        // sum = 42951 * 50000 = 2,147,550,000 > i32::MAX
+        let array = Decimal32Array::from_iter_values(28525..=71475)
+            .with_precision_and_scale(5, 0)?;
+
+        let mut accumulator = DecimalDistinctAvgAccumulator::<
+            Decimal32Type,
+            Decimal64Type,
+        >::with_decimal_params(0, 9, 4);
+        accumulator.update_batch(&[Arc::new(array)])?;
+
+        assert_eq!(
+            accumulator.evaluate()?,
+            ScalarValue::Decimal32(Some(500_000_000), 9, 4)
+        );
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_decimal32_distinct_avg_widens_to_decimal128() -> Result<()> {
+        // 21477 distinct values centered on 99999:
+        // sum = 21477 * 99999 = 2,147,678,523 > i32::MAX
+        let array = Decimal32Array::from_iter_values(89261..=110737)
+            .with_precision_and_scale(9, 0)?;
+
+        let mut accumulator = DecimalDistinctAvgAccumulator::<
+            Decimal32Type,
+            Decimal128Type,
+        >::with_decimal_params(0, 9, 4);
+        accumulator.update_batch(&[Arc::new(array)])?;
+
+        assert_eq!(
+            accumulator.evaluate()?,
+            ScalarValue::Decimal32(Some(999_990_000), 9, 4)
+        );
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_decimal64_distinct_avg_widens_to_decimal128() -> Result<()> {
+        // 92235 distinct values centered on 10^14 - 1:
+        // sum = 92235 * (10^14 - 1) ~= 9.22e18 > i64::MAX
+        let center: i64 = 100_000_000_000_000 - 1;
+        let array = Decimal64Array::from_iter_values(center - 46117..=center + 46117)
+            .with_precision_and_scale(18, 0)?;
+
+        let mut accumulator = DecimalDistinctAvgAccumulator::<
+            Decimal64Type,
+            Decimal128Type,
+        >::with_decimal_params(0, 18, 4);
+        accumulator.update_batch(&[Arc::new(array)])?;
+
+        assert_eq!(
+            accumulator.evaluate()?,
+            ScalarValue::Decimal64(Some(999_999_999_999_990_000), 18, 4)
+        );
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_decimal128_distinct_avg_widens_to_decimal256() -> Result<()> {
+        // 21477 distinct values ending at 10^34 - 1, centered on 10^34 - 10739:
+        // sum = 21477 * (10^34 - 10739) ~= 2.15e38 > i128::MAX
+        let center: i128 = 10_i128.pow(34) - 10739;
+        let array = Decimal128Array::from_iter_values(center - 10738..=center + 10738)
+            .with_precision_and_scale(34, 0)?;
+
+        let mut accumulator = DecimalDistinctAvgAccumulator::<
+            Decimal128Type,
+            Decimal256Type,
+        >::with_decimal_params(0, 38, 4);
+        accumulator.update_batch(&[Arc::new(array)])?;
+
+        assert_eq!(
+            accumulator.evaluate()?,
+            ScalarValue::Decimal128(Some(center * 10_000), 38, 4)
+        );
+
+        Ok(())
+    }
 }

datafusion/functions-aggregate-common/src/aggregate/sum_distinct/numeric.rs

@@ -50,6 +50,13 @@ impl<T: ArrowPrimitiveType> DistinctSumAccumulator<T> {
     pub fn distinct_count(&self) -> usize {
         self.values.values.len()
     }
+
+    /// Iterates the distinct values collected so far. `AVG(DISTINCT)` re-sums
+    /// them in a wider type instead of using [`Self::evaluate`]'s input-typed
+    /// sum.
+    pub(crate) fn distinct_values(&self) -> impl Iterator<Item = T::Native> + '_ {
+        self.values.values.iter().map(|v| v.0)
+    }
 }
 
 impl<T: ArrowPrimitiveType + Debug> Accumulator for DistinctSumAccumulator<T> {

datafusion/functions-aggregate/src/average.rs

@@ -138,7 +138,7 @@ impl Default for Avg {
 /// of reach. `Decimal256` input near max precision is the exception: no wider
 /// type exists, so its sum keeps only whatever headroom `Decimal256(76, _)` has
 /// left, as before this budget was introduced.
-const AVG_SUM_HEADROOM_DIGITS: u8 = 4 + 9;
+const AVG_SUM_HEADROOM_DIGITS: u8 = 13;
 
 /// The narrowest decimal that can accumulate `avg`'s sum over `data_type`, never
 /// narrower than `data_type` itself. Other types accumulate as themselves.
@@ -261,39 +261,19 @@ impl AggregateUDFImpl for Avg {
                 // Numeric types are converted to Float64 via `coerce_avg_type` during logical plan creation
                 (Float64, _) => Ok(Box::new(Float64DistinctAvgAccumulator::default())),
 
-                (
-                    Decimal32(_, scale),
-                    Decimal32(target_precision, target_scale),
-                ) => Ok(Box::new(DecimalDistinctAvgAccumulator::<Decimal32Type>::with_decimal_params(
-                    *scale,
-                    *target_precision,
-                    *target_scale,
-                ))),
-                                (
-                    Decimal64(_, scale),
-                    Decimal64(target_precision, target_scale),
-                ) => Ok(Box::new(DecimalDistinctAvgAccumulator::<Decimal64Type>::with_decimal_params(
-                    *scale,
-                    *target_precision,
-                    *target_scale,
-                ))),
-                (
-                    Decimal128(_, scale),
-                    Decimal128(target_precision, target_scale),
-                ) => Ok(Box::new(DecimalDistinctAvgAccumulator::<Decimal128Type>::with_decimal_params(
-                    *scale,
-                    *target_precision,
-                    *target_scale,
-                ))),
-
-                (
-                    Decimal256(_, scale),
-                    Decimal256(target_precision, target_scale),
-                ) => Ok(Box::new(DecimalDistinctAvgAccumulator::<Decimal256Type>::with_decimal_params(
-                    *scale,
-                    *target_precision,
-                    *target_scale,
-                ))),
+                (Decimal32(..), Decimal32(..))
+                | (Decimal64(..), Decimal64(..))
+                | (Decimal128(..), Decimal128(..))
+                | (Decimal256(..), Decimal256(..)) => {
+                    let sum_data_type = avg_sum_data_type(data_type);
+                    decimal_avg_dispatch!(
+                        data_type,
+                        &sum_data_type,
+                        decimal_distinct_avg_accumulator,
+                        &sum_data_type,
+                        acc_args.return_type()
+                    )
+                }
 
                 (dt, return_type) => exec_err!(
                     "AVG(DISTINCT) for ({} --> {}) not supported",
@@ -550,6 +530,28 @@ where
     )))
 }
 
+fn decimal_distinct_avg_accumulator<I, S>(
+    sum_data_type: &DataType,
+    return_data_type: &DataType,
+) -> Result<Box<dyn Accumulator>>
+where
+    I: DecimalType + ArrowNumericType + Debug + Send + Sync,
+    S: DecimalType + ArrowNumericType + Debug + Send + Sync,
+    I::Native: Into<S::Native> + DecimalCast,
+    S::Native: DecimalCast,
+{
+    let (_, sum_scale) = decimal_parts(sum_data_type)?;
+    let (target_precision, target_scale) = decimal_parts(return_data_type)?;
+
+    Ok(Box::new(
+        DecimalDistinctAvgAccumulator::<I, S>::with_decimal_params(
+            sum_scale,
+            target_precision,
+            target_scale,
+        ),
+    ))
+}
+
 fn decimal_avg_groups_accumulator<I, S>(
     sum_data_type: &DataType,
     return_data_type: &DataType,