vortex_array/aggregate_fn/fns/mean/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexResult;
5use vortex_error::vortex_bail;
6use vortex_session::registry::CachedId;
7
8use crate::ArrayRef;
9use crate::ExecutionCtx;
10use crate::aggregate_fn::Accumulator;
11use crate::aggregate_fn::AggregateFnId;
12use crate::aggregate_fn::AggregateFnVTable;
13use crate::aggregate_fn::DynAccumulator;
14use crate::aggregate_fn::NumericalAggregateOpts;
15use crate::aggregate_fn::combined::BinaryCombined;
16use crate::aggregate_fn::combined::Combined;
17use crate::aggregate_fn::combined::CombinedOptions;
18use crate::aggregate_fn::combined::PairOptions;
19use crate::aggregate_fn::fns::count::Count;
20use crate::aggregate_fn::fns::sum::Sum;
21use crate::aggregate_fn::fns::sum::sum_decimal_dtype;
22use crate::builtins::ArrayBuiltins;
23use crate::dtype::DType;
24use crate::dtype::DecimalDType;
25use crate::dtype::MAX_PRECISION;
26use crate::dtype::MAX_SCALE;
27use crate::dtype::Nullability;
28use crate::dtype::PType;
29use crate::dtype::i256;
30use crate::scalar::DecimalValue;
31use crate::scalar::Scalar;
32use crate::scalar_fn::fns::operators::Operator;
33
34/// Compute the arithmetic mean of an array.
35///
36/// See [`Mean`] for details.
37pub fn mean(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Scalar> {
38    let mut acc = Accumulator::try_new(
39        Mean::combined(),
40        PairOptions(
41            NumericalAggregateOpts::default(),
42            NumericalAggregateOpts::default(),
43        ),
44        array.dtype().clone(),
45    )?;
46    acc.accumulate(array, ctx)?;
47    acc.finish()
48}
49
50/// Compute the arithmetic mean of an array.
51///
52/// Implemented as `Sum / Count` via [`BinaryCombined`].
53///
54/// Booleans and primitive numeric types are cast to f64. Decimals stay decimals.
55#[derive(Clone, Debug)]
56pub struct Mean;
57
58impl Mean {
59    pub fn combined() -> Combined<Self> {
60        Combined(Mean)
61    }
62}
63
64impl BinaryCombined for Mean {
65    type Left = Sum;
66    type Right = Count;
67
68    fn id(&self) -> AggregateFnId {
69        static ID: CachedId = CachedId::new("vortex.mean");
70        *ID
71    }
72
73    fn left(&self) -> Sum {
74        Sum
75    }
76
77    fn right(&self) -> Count {
78        Count
79    }
80
81    fn left_name(&self) -> &'static str {
82        "sum"
83    }
84
85    fn right_name(&self) -> &'static str {
86        "count"
87    }
88
89    fn return_dtype(&self, input_dtype: &DType) -> Option<DType> {
90        Some(mean_output_dtype(input_dtype)?.with_nullability(Nullability::Nullable))
91    }
92
93    fn finalize(&self, sum: ArrayRef, count: ArrayRef) -> VortexResult<ArrayRef> {
94        if let DType::Decimal(..) = sum.dtype() {
95            vortex_bail!("grouped mean over decimals is not yet supported");
96        }
97        let target = DType::Primitive(PType::F64, Nullability::Nullable);
98        let sum_cast = sum.cast(target.clone())?;
99        let count_cast = count.cast(target)?;
100        sum_cast.binary(count_cast, Operator::Div)
101    }
102
103    fn finalize_scalar(&self, left_scalar: Scalar, right_scalar: Scalar) -> VortexResult<Scalar> {
104        if let DType::Decimal(decimal_dtype, _) = *left_scalar.dtype() {
105            return finalize_decimal_scalar(&left_scalar, &right_scalar, decimal_dtype);
106        }
107
108        let target = DType::Primitive(PType::F64, Nullability::Nullable);
109        let sum_cast = left_scalar.cast(&target)?;
110        let count_cast = right_scalar.cast(&target)?;
111
112        let sum = sum_cast.as_primitive().typed_value::<f64>();
113        let count = count_cast.as_primitive().typed_value::<f64>();
114        let value = match (sum, count) {
115            (None, _) | (_, None) => return Ok(Scalar::null(target)), // Sum overflowed
116            // A count of zero yields 0/0 = NaN, matching the array `finalize` path: nulls are
117            // skipped during accumulation, so an all-null input is an empty mean, not null.
118            (Some(s), Some(c)) => s / c,
119        };
120        Ok(Scalar::primitive(value, Nullability::Nullable))
121    }
122
123    fn serialize(&self, _options: &CombinedOptions<Self>) -> VortexResult<Option<Vec<u8>>> {
124        unimplemented!("mean is not yet serializable");
125    }
126
127    fn coerce_args(
128        &self,
129        _options: &PairOptions<
130            <Sum as AggregateFnVTable>::Options,
131            <Count as AggregateFnVTable>::Options,
132        >,
133        input_dtype: &DType,
134    ) -> VortexResult<DType> {
135        // Advisory hint for query planners: where possible, cast input to the
136        // type we're going to compute the mean in.
137        Ok(coerced_input_dtype(input_dtype).unwrap_or_else(|| input_dtype.clone()))
138    }
139}
140
141/// Hint for callers: what to cast the input to before accumulation.
142///
143/// - Bool stays as bool — `Sum` has a native bool path and bool → f64 isn't
144///   currently a direct cast in vortex.
145/// - Primitive numerics → `f64` so the sum and finalize work without overflow.
146/// - Decimals stay as decimals
147fn coerced_input_dtype(input_dtype: &DType) -> Option<DType> {
148    match input_dtype {
149        DType::Bool(_) => Some(input_dtype.clone()),
150        DType::Primitive(_, n) => Some(DType::Primitive(PType::F64, *n)),
151        DType::Decimal(..) => Some(input_dtype.clone()),
152        _ => None,
153    }
154}
155
156fn mean_output_dtype(input_dtype: &DType) -> Option<DType> {
157    match input_dtype {
158        DType::Bool(_) | DType::Primitive(..) => {
159            Some(DType::Primitive(PType::F64, Nullability::Nullable))
160        }
161        DType::Decimal(decimal_dtype, _) => Some(DType::Decimal(
162            mean_decimal_dtype(&sum_decimal_dtype(decimal_dtype)),
163            Nullability::Nullable,
164        )),
165        _ => None,
166    }
167}
168
169/// mean() output decimal type mimicking Spark/DataFusion/MySQL: decimal(p+4, s+4)
170fn mean_decimal_dtype(sum: &DecimalDType) -> DecimalDType {
171    DecimalDType::new(
172        u8::min(MAX_PRECISION, sum.precision().saturating_sub(6)),
173        i8::min(MAX_SCALE, sum.scale() + 4),
174    )
175}
176
177fn finalize_decimal_scalar(
178    sum: &Scalar,
179    count: &Scalar,
180    sum_decimal: DecimalDType,
181) -> VortexResult<Scalar> {
182    let target_decimal_dtype = mean_decimal_dtype(&sum_decimal);
183    let target_dtype = DType::Decimal(target_decimal_dtype, Nullability::Nullable);
184
185    // overflow
186    let Some(sum_value) = sum.as_decimal().decimal_value() else {
187        return Ok(Scalar::null(target_dtype));
188    };
189    // empty input
190    let count = count.as_primitive().typed_value::<u64>().unwrap_or(0);
191    if count == 0 {
192        return Ok(Scalar::null(target_dtype));
193    }
194
195    let Ok(sum) = DecimalValue::rescale_i256(
196        sum_value.as_i256(),
197        sum_decimal.scale(),
198        target_decimal_dtype.scale(),
199    ) else {
200        return Ok(Scalar::null(target_dtype));
201    };
202    let mean = sum / i256::from_i128(i128::from(count));
203
204    let Ok(mean) = DecimalValue::try_from_i256(mean, target_decimal_dtype) else {
205        return Ok(Scalar::null(target_dtype));
206    };
207    Ok(Scalar::decimal(
208        mean,
209        target_decimal_dtype,
210        Nullability::Nullable,
211    ))
212}
213
214#[cfg(test)]
215mod tests {
216    use vortex_buffer::buffer;
217    use vortex_error::VortexResult;
218
219    use super::*;
220    use crate::IntoArray;
221    use crate::VortexSessionExecute;
222    use crate::array_session;
223    use crate::arrays::BoolArray;
224    use crate::arrays::ChunkedArray;
225    use crate::arrays::ConstantArray;
226    use crate::arrays::DecimalArray;
227    use crate::arrays::PrimitiveArray;
228    use crate::dtype::DecimalDType;
229    use crate::validity::Validity;
230
231    #[test]
232    fn mean_all_valid() -> VortexResult<()> {
233        let array = PrimitiveArray::new(buffer![1.0f64, 2.0, 3.0, 4.0, 5.0], Validity::NonNullable)
234            .into_array();
235        let mut ctx = array_session().create_execution_ctx();
236        let result = mean(&array, &mut ctx)?;
237        assert_eq!(result.as_primitive().as_::<f64>(), Some(3.0));
238        Ok(())
239    }
240
241    #[test]
242    fn mean_with_nulls() -> VortexResult<()> {
243        let array = PrimitiveArray::from_option_iter([Some(2.0f64), None, Some(4.0)]).into_array();
244        let mut ctx = array_session().create_execution_ctx();
245        let result = mean(&array, &mut ctx)?;
246        assert_eq!(result.as_primitive().as_::<f64>(), Some(3.0));
247        Ok(())
248    }
249
250    #[test]
251    fn mean_integers() -> VortexResult<()> {
252        let array = PrimitiveArray::new(buffer![10i32, 20, 30], Validity::NonNullable).into_array();
253        let mut ctx = array_session().create_execution_ctx();
254        let result = mean(&array, &mut ctx)?;
255        assert_eq!(result.as_primitive().as_::<f64>(), Some(20.0));
256        Ok(())
257    }
258
259    #[test]
260    fn mean_bool() -> VortexResult<()> {
261        let array: BoolArray = [true, false, true, true].into_iter().collect();
262        let mut ctx = array_session().create_execution_ctx();
263        let result = mean(&array.into_array(), &mut ctx)?;
264        assert_eq!(result.as_primitive().as_::<f64>(), Some(0.75));
265        Ok(())
266    }
267
268    #[test]
269    fn mean_constant_non_null() -> VortexResult<()> {
270        let array = ConstantArray::new(5.0f64, 4);
271        let mut ctx = array_session().create_execution_ctx();
272        let result = mean(&array.into_array(), &mut ctx)?;
273        assert_eq!(result.as_primitive().as_::<f64>(), Some(5.0));
274        Ok(())
275    }
276
277    #[test]
278    fn mean_chunked() -> VortexResult<()> {
279        let chunk1 = PrimitiveArray::from_option_iter([Some(1.0f64), None, Some(3.0)]);
280        let chunk2 = PrimitiveArray::from_option_iter([Some(5.0f64), None]);
281        let dtype = chunk1.dtype().clone();
282        let chunked = ChunkedArray::try_new(vec![chunk1.into_array(), chunk2.into_array()], dtype)?;
283        let mut ctx = array_session().create_execution_ctx();
284        let result = mean(&chunked.into_array(), &mut ctx)?;
285        assert_eq!(result.as_primitive().as_::<f64>(), Some(3.0));
286        Ok(())
287    }
288
289    #[test]
290    fn mean_skips_nans_by_default() -> VortexResult<()> {
291        // NaNs are excluded from both the sum and the count.
292        let array =
293            PrimitiveArray::new(buffer![1.0f64, f64::NAN, 3.0], Validity::NonNullable).into_array();
294        let mut ctx = array_session().create_execution_ctx();
295        let result = mean(&array, &mut ctx)?;
296        assert_eq!(result.as_primitive().as_::<f64>(), Some(2.0));
297        Ok(())
298    }
299
300    #[test]
301    fn mean_with_nan_not_skipping() -> VortexResult<()> {
302        let array =
303            PrimitiveArray::new(buffer![1.0f64, f64::NAN, 3.0], Validity::NonNullable).into_array();
304        let mut ctx = array_session().create_execution_ctx();
305        let keep_nans = NumericalAggregateOpts::include_nans();
306        let mut acc = Accumulator::try_new(
307            Mean::combined(),
308            PairOptions(keep_nans, keep_nans),
309            array.dtype().clone(),
310        )?;
311        acc.accumulate(&array, &mut ctx)?;
312        let result = acc.finish()?;
313        assert!(result.as_primitive().as_::<f64>().is_some_and(f64::is_nan));
314        Ok(())
315    }
316
317    #[test]
318    fn mean_all_null_returns_nan() -> VortexResult<()> {
319        let array = PrimitiveArray::from_option_iter::<f64, _>([None, None, None]).into_array();
320        let mut ctx = array_session().create_execution_ctx();
321        let result = mean(&array, &mut ctx)?;
322        assert!(result.as_primitive().as_::<f64>().is_some_and(f64::is_nan));
323        Ok(())
324    }
325
326    #[test]
327    fn mean_decimal() -> VortexResult<()> {
328        let dtype = DecimalDType::new(6, 2);
329        let array =
330            DecimalArray::new(buffer![100i32, 200, 300], dtype, Validity::NonNullable).into_array();
331        let mut ctx = array_session().create_execution_ctx();
332        let result = mean(&array, &mut ctx)?;
333        assert_eq!(
334            result.dtype(),
335            &DType::Decimal(DecimalDType::new(10, 6), Nullability::Nullable)
336        );
337        // mean(1.00, 2.00, 3.00) = 2.000000
338        assert_eq!(
339            result.as_decimal().decimal_value(),
340            Some(DecimalValue::I256(i256::from_i128(2_000_000)))
341        );
342        Ok(())
343    }
344
345    #[test]
346    fn mean_decimal_null() -> VortexResult<()> {
347        let dtype = DecimalDType::new(6, 2);
348        let validity = Validity::from_iter([true, false, true]);
349        let array = DecimalArray::new(buffer![150i32, 0, 450], dtype, validity).into_array();
350        let mut ctx = array_session().create_execution_ctx();
351        let result = mean(&array, &mut ctx)?;
352        // mean(1.50, 4.50) = 3.000000
353        assert_eq!(
354            result.as_decimal().decimal_value(),
355            Some(DecimalValue::I256(i256::from_i128(3_000_000)))
356        );
357        Ok(())
358    }
359
360    #[test]
361    fn mean_decimal_chunked() -> VortexResult<()> {
362        let dtype = DecimalDType::new(6, 2);
363        let validity = Validity::NonNullable;
364        let chunk1 = DecimalArray::new(buffer![100i32, 200], dtype, validity.clone()).into_array();
365        let chunk2 = DecimalArray::new(buffer![300i32, 400, 500], dtype, validity).into_array();
366        let dtype = chunk1.dtype().clone();
367        let chunked = ChunkedArray::try_new(vec![chunk1, chunk2], dtype)?;
368        let mut ctx = array_session().create_execution_ctx();
369        let result = mean(&chunked.into_array(), &mut ctx)?;
370        // mean(1.00, 2.00, 3.00, 4.00, 5.00) = 3.000000
371        assert_eq!(
372            result.as_decimal().decimal_value(),
373            Some(DecimalValue::I256(i256::from_i128(3_000_000)))
374        );
375        Ok(())
376    }
377
378    #[test]
379    fn mean_decimal_33() -> VortexResult<()> {
380        let dtype = DecimalDType::new(6, 2);
381        let buf = buffer![100i32, 0, 0];
382        let array = DecimalArray::new(buf, dtype, Validity::NonNullable).into_array();
383        let mut ctx = array_session().create_execution_ctx();
384        let result = mean(&array, &mut ctx)?;
385        // mean(1.00, 0.00, 0.00) = 1/3 => 0.333333
386        assert_eq!(
387            result.as_decimal().decimal_value(),
388            Some(DecimalValue::I256(i256::from_i128(333_333)))
389        );
390        Ok(())
391    }
392
393    #[test]
394    fn mean_multi_batch() -> VortexResult<()> {
395        let mut ctx = array_session().create_execution_ctx();
396        let dtype = DType::Primitive(PType::F64, Nullability::NonNullable);
397        let mut acc = Accumulator::try_new(
398            Mean::combined(),
399            PairOptions(
400                NumericalAggregateOpts::default(),
401                NumericalAggregateOpts::default(),
402            ),
403            dtype,
404        )?;
405
406        let batch1 =
407            PrimitiveArray::new(buffer![1.0f64, 2.0, 3.0], Validity::NonNullable).into_array();
408        acc.accumulate(&batch1, &mut ctx)?;
409
410        let batch2 = PrimitiveArray::new(buffer![4.0f64, 5.0], Validity::NonNullable).into_array();
411        acc.accumulate(&batch2, &mut ctx)?;
412
413        let result = acc.finish()?;
414        assert_eq!(result.as_primitive().as_::<f64>(), Some(3.0));
415        Ok(())
416    }
417}
vortex_array/aggregate_fn/fns/mean/mod.rs

vortex_array/aggregate_fn/fns/mean/
mod.rs
