vortex_array/arrays/primitive/compute/
min_max.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use itertools::Itertools;
5use vortex_dtype::Nullability::NonNullable;
6use vortex_dtype::{NativePType, match_each_native_ptype};
7use vortex_error::VortexResult;
8use vortex_mask::Mask;
9use vortex_scalar::{PValue, Scalar};
10
11use crate::arrays::{PrimitiveArray, PrimitiveVTable};
12use crate::compute::{MinMaxKernel, MinMaxKernelAdapter, MinMaxResult};
13use crate::register_kernel;
14
15impl MinMaxKernel for PrimitiveVTable {
16    fn min_max(&self, array: &PrimitiveArray) -> VortexResult<Option<MinMaxResult>> {
17        match_each_native_ptype!(array.ptype(), |T| {
18            compute_min_max_with_validity::<T>(array)
19        })
20    }
21}
22
23register_kernel!(MinMaxKernelAdapter(PrimitiveVTable).lift());
24
25#[inline]
26fn compute_min_max_with_validity<T>(array: &PrimitiveArray) -> VortexResult<Option<MinMaxResult>>
27where
28    T: NativePType,
29    PValue: From<T>,
30{
31    Ok(match array.validity_mask() {
32        Mask::AllTrue(_) => compute_min_max(array.as_slice::<T>().iter()),
33        Mask::AllFalse(_) => None,
34        Mask::Values(v) => compute_min_max(
35            array
36                .as_slice::<T>()
37                .iter()
38                .zip(v.bit_buffer().iter())
39                .filter_map(|(v, m)| m.then_some(v)),
40        ),
41    })
42}
43
44fn compute_min_max<'a, T>(iter: impl Iterator<Item = &'a T>) -> Option<MinMaxResult>
45where
46    T: NativePType,
47    PValue: From<T>,
48{
49    // `total_compare` function provides a total ordering (even for NaN values).
50    // However, we exclude NaNs from min max as they're not useful for any purpose where min/max would be used
51    match iter
52        .filter(|v| !v.is_nan())
53        .minmax_by(|a, b| a.total_compare(**b))
54    {
55        itertools::MinMaxResult::NoElements => None,
56        itertools::MinMaxResult::OneElement(&x) => {
57            let scalar = Scalar::primitive(x, NonNullable);
58            Some(MinMaxResult {
59                min: scalar.clone(),
60                max: scalar,
61            })
62        }
63        itertools::MinMaxResult::MinMax(&min, &max) => Some(MinMaxResult {
64            min: Scalar::primitive(min, NonNullable),
65            max: Scalar::primitive(max, NonNullable),
66        }),
67    }
68}
69
70#[cfg(test)]
71mod tests {
72    use vortex_buffer::buffer;
73
74    use crate::arrays::PrimitiveArray;
75    use crate::compute::min_max;
76    use crate::validity::Validity;
77
78    #[test]
79    fn min_max_nan() {
80        let array = PrimitiveArray::new(
81            buffer![f32::NAN, -f32::NAN, -1.0, 1.0],
82            Validity::NonNullable,
83        );
84        let min_max = min_max(array.as_ref()).unwrap().unwrap();
85        assert_eq!(f32::try_from(min_max.min).unwrap(), -1.0);
86        assert_eq!(f32::try_from(min_max.max).unwrap(), 1.0);
87    }
88
89    #[test]
90    fn min_max_inf() {
91        let array = PrimitiveArray::new(
92            buffer![f32::INFINITY, f32::NEG_INFINITY, -1.0, 1.0],
93            Validity::NonNullable,
94        );
95        let min_max = min_max(array.as_ref()).unwrap().unwrap();
96        assert_eq!(f32::try_from(min_max.min).unwrap(), f32::NEG_INFINITY);
97        assert_eq!(f32::try_from(min_max.max).unwrap(), f32::INFINITY);
98    }
99}