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

vortex_array/aggregate_fn/fns/mean/
mod.rs
