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

vortex_array/aggregate_fn/fns/mean/
mod.rs
