vortex_array/arrays/primitive/compute/
min_max.rs1use 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 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}