vortex_array/arrays/primitive/compute/
min_max.rs

1use itertools::Itertools;
2use vortex_dtype::{DType, NativePType, match_each_native_ptype};
3use vortex_error::VortexResult;
4use vortex_mask::Mask;
5use vortex_scalar::{Scalar, ScalarValue};
6
7use crate::Array;
8use crate::arrays::{PrimitiveArray, PrimitiveEncoding};
9use crate::compute::{MinMaxFn, MinMaxResult};
10use crate::variants::PrimitiveArrayTrait;
11
12impl MinMaxFn<&PrimitiveArray> for PrimitiveEncoding {
13    fn min_max(&self, array: &PrimitiveArray) -> VortexResult<Option<MinMaxResult>> {
14        match_each_native_ptype!(array.ptype(), |$T| {
15            compute_min_max_with_validity::<$T>(array)
16        })
17    }
18}
19
20#[inline]
21fn compute_min_max_with_validity<T>(array: &PrimitiveArray) -> VortexResult<Option<MinMaxResult>>
22where
23    T: Into<ScalarValue> + NativePType,
24{
25    Ok(match array.validity_mask()? {
26        Mask::AllTrue(_) => compute_min_max(array.as_slice::<T>().iter(), array.dtype()),
27        Mask::AllFalse(_) => None,
28        Mask::Values(v) => compute_min_max(
29            array
30                .as_slice::<T>()
31                .iter()
32                .zip(v.boolean_buffer().iter())
33                .filter_map(|(v, m)| m.then_some(v)),
34            array.dtype(),
35        ),
36    })
37}
38
39fn compute_min_max<'a, T>(iter: impl Iterator<Item = &'a T>, dtype: &DType) -> Option<MinMaxResult>
40where
41    T: Into<ScalarValue> + NativePType + Copy,
42{
43    // this `compare` function provides a total ordering (even for NaN values)
44    match iter.minmax_by(|a, b| a.total_compare(**b)) {
45        itertools::MinMaxResult::NoElements => None,
46        itertools::MinMaxResult::OneElement(&x) => {
47            let scalar = Scalar::new(dtype.clone(), x.into());
48            Some(MinMaxResult {
49                min: scalar.clone(),
50                max: scalar,
51            })
52        }
53        itertools::MinMaxResult::MinMax(&min, &max) => Some(MinMaxResult {
54            min: Scalar::new(dtype.clone(), min.into()),
55            max: Scalar::new(dtype.clone(), max.into()),
56        }),
57    }
58}