polars_compute/min_max/
scalar.rs

1use arrow::array::{
2    Array, BinaryArray, BinaryViewArray, BooleanArray, PrimitiveArray, Utf8Array, Utf8ViewArray,
3};
4use arrow::types::{NativeType, Offset};
5use polars_utils::min_max::MinMax;
6
7use super::MinMaxKernel;
8
9fn min_max_ignore_nan<T: NativeType>((cur_min, cur_max): (T, T), (min, max): (T, T)) -> (T, T) {
10    (
11        MinMax::min_ignore_nan(cur_min, min),
12        MinMax::max_ignore_nan(cur_max, max),
13    )
14}
15
16fn min_max_propagate_nan<T: NativeType>((cur_min, cur_max): (T, T), (min, max): (T, T)) -> (T, T) {
17    (
18        MinMax::min_propagate_nan(cur_min, min),
19        MinMax::max_propagate_nan(cur_max, max),
20    )
21}
22
23fn reduce_vals<T, F>(v: &PrimitiveArray<T>, f: F) -> Option<T>
24where
25    T: NativeType,
26    F: Fn(T, T) -> T,
27{
28    if v.null_count() == 0 {
29        v.values_iter().copied().reduce(f)
30    } else {
31        v.non_null_values_iter().reduce(f)
32    }
33}
34
35fn reduce_tuple_vals<T, F>(v: &PrimitiveArray<T>, f: F) -> Option<(T, T)>
36where
37    T: NativeType,
38    F: Fn((T, T), (T, T)) -> (T, T),
39{
40    if v.null_count() == 0 {
41        v.values_iter().copied().map(|v| (v, v)).reduce(f)
42    } else {
43        v.non_null_values_iter().map(|v| (v, v)).reduce(f)
44    }
45}
46
47impl<T: NativeType + MinMax + super::NotSimdPrimitive> MinMaxKernel for PrimitiveArray<T> {
48    type Scalar<'a> = T;
49
50    fn min_ignore_nan_kernel(&self) -> Option<Self::Scalar<'_>> {
51        reduce_vals(self, MinMax::min_ignore_nan)
52    }
53
54    fn max_ignore_nan_kernel(&self) -> Option<Self::Scalar<'_>> {
55        reduce_vals(self, MinMax::max_ignore_nan)
56    }
57
58    fn min_max_ignore_nan_kernel(&self) -> Option<(Self::Scalar<'_>, Self::Scalar<'_>)> {
59        reduce_tuple_vals(self, min_max_ignore_nan)
60    }
61
62    fn min_propagate_nan_kernel(&self) -> Option<Self::Scalar<'_>> {
63        reduce_vals(self, MinMax::min_propagate_nan)
64    }
65
66    fn max_propagate_nan_kernel(&self) -> Option<Self::Scalar<'_>> {
67        reduce_vals(self, MinMax::max_propagate_nan)
68    }
69
70    fn min_max_propagate_nan_kernel(&self) -> Option<(Self::Scalar<'_>, Self::Scalar<'_>)> {
71        reduce_tuple_vals(self, min_max_propagate_nan)
72    }
73}
74
75impl<T: NativeType + MinMax + super::NotSimdPrimitive> MinMaxKernel for [T] {
76    type Scalar<'a> = T;
77
78    fn min_ignore_nan_kernel(&self) -> Option<Self::Scalar<'_>> {
79        self.iter().copied().reduce(MinMax::min_ignore_nan)
80    }
81
82    fn max_ignore_nan_kernel(&self) -> Option<Self::Scalar<'_>> {
83        self.iter().copied().reduce(MinMax::max_ignore_nan)
84    }
85
86    fn min_max_ignore_nan_kernel(&self) -> Option<(Self::Scalar<'_>, Self::Scalar<'_>)> {
87        self.iter()
88            .copied()
89            .map(|v| (v, v))
90            .reduce(min_max_ignore_nan)
91    }
92
93    fn min_propagate_nan_kernel(&self) -> Option<Self::Scalar<'_>> {
94        self.iter().copied().reduce(MinMax::min_propagate_nan)
95    }
96
97    fn max_propagate_nan_kernel(&self) -> Option<Self::Scalar<'_>> {
98        self.iter().copied().reduce(MinMax::max_propagate_nan)
99    }
100
101    fn min_max_propagate_nan_kernel(&self) -> Option<(Self::Scalar<'_>, Self::Scalar<'_>)> {
102        self.iter()
103            .copied()
104            .map(|v| (v, v))
105            .reduce(min_max_propagate_nan)
106    }
107}
108
109impl MinMaxKernel for BooleanArray {
110    type Scalar<'a> = bool;
111
112    fn min_ignore_nan_kernel(&self) -> Option<Self::Scalar<'_>> {
113        if self.len() - self.null_count() == 0 {
114            return None;
115        }
116
117        let unset_bits = self.values().unset_bits();
118        Some(unset_bits == 0)
119    }
120
121    fn max_ignore_nan_kernel(&self) -> Option<Self::Scalar<'_>> {
122        if self.len() - self.null_count() == 0 {
123            return None;
124        }
125
126        let set_bits = self.values().set_bits();
127        Some(set_bits > 0)
128    }
129
130    #[inline(always)]
131    fn min_propagate_nan_kernel(&self) -> Option<Self::Scalar<'_>> {
132        self.min_ignore_nan_kernel()
133    }
134
135    #[inline(always)]
136    fn max_propagate_nan_kernel(&self) -> Option<Self::Scalar<'_>> {
137        self.max_ignore_nan_kernel()
138    }
139}
140
141impl MinMaxKernel for BinaryViewArray {
142    type Scalar<'a> = &'a [u8];
143
144    fn min_ignore_nan_kernel(&self) -> Option<Self::Scalar<'_>> {
145        if self.null_count() == 0 {
146            self.values_iter().reduce(MinMax::min_ignore_nan)
147        } else {
148            self.non_null_values_iter().reduce(MinMax::min_ignore_nan)
149        }
150    }
151
152    fn max_ignore_nan_kernel(&self) -> Option<Self::Scalar<'_>> {
153        if self.null_count() == 0 {
154            self.values_iter().reduce(MinMax::max_ignore_nan)
155        } else {
156            self.non_null_values_iter().reduce(MinMax::max_ignore_nan)
157        }
158    }
159
160    #[inline(always)]
161    fn min_propagate_nan_kernel(&self) -> Option<Self::Scalar<'_>> {
162        self.min_ignore_nan_kernel()
163    }
164
165    #[inline(always)]
166    fn max_propagate_nan_kernel(&self) -> Option<Self::Scalar<'_>> {
167        self.max_ignore_nan_kernel()
168    }
169}
170
171impl MinMaxKernel for Utf8ViewArray {
172    type Scalar<'a> = &'a str;
173
174    #[inline(always)]
175    fn min_ignore_nan_kernel(&self) -> Option<Self::Scalar<'_>> {
176        self.to_binview().min_ignore_nan_kernel().map(|s| unsafe {
177            // SAFETY: the lifetime is the same, and it is valid UTF-8.
178            #[allow(clippy::transmute_bytes_to_str)]
179            std::mem::transmute::<&[u8], &str>(s)
180        })
181    }
182
183    #[inline(always)]
184    fn max_ignore_nan_kernel(&self) -> Option<Self::Scalar<'_>> {
185        self.to_binview().max_ignore_nan_kernel().map(|s| unsafe {
186            // SAFETY: the lifetime is the same, and it is valid UTF-8.
187            #[allow(clippy::transmute_bytes_to_str)]
188            std::mem::transmute::<&[u8], &str>(s)
189        })
190    }
191
192    #[inline(always)]
193    fn min_propagate_nan_kernel(&self) -> Option<Self::Scalar<'_>> {
194        self.min_ignore_nan_kernel()
195    }
196
197    #[inline(always)]
198    fn max_propagate_nan_kernel(&self) -> Option<Self::Scalar<'_>> {
199        self.max_ignore_nan_kernel()
200    }
201}
202
203impl<O: Offset> MinMaxKernel for BinaryArray<O> {
204    type Scalar<'a> = &'a [u8];
205
206    fn min_ignore_nan_kernel(&self) -> Option<Self::Scalar<'_>> {
207        if self.null_count() == 0 {
208            self.values_iter().reduce(MinMax::min_ignore_nan)
209        } else {
210            self.non_null_values_iter().reduce(MinMax::min_ignore_nan)
211        }
212    }
213
214    fn max_ignore_nan_kernel(&self) -> Option<Self::Scalar<'_>> {
215        if self.null_count() == 0 {
216            self.values_iter().reduce(MinMax::max_ignore_nan)
217        } else {
218            self.non_null_values_iter().reduce(MinMax::max_ignore_nan)
219        }
220    }
221
222    #[inline(always)]
223    fn min_propagate_nan_kernel(&self) -> Option<Self::Scalar<'_>> {
224        self.min_ignore_nan_kernel()
225    }
226
227    #[inline(always)]
228    fn max_propagate_nan_kernel(&self) -> Option<Self::Scalar<'_>> {
229        self.max_ignore_nan_kernel()
230    }
231}
232
233impl<O: Offset> MinMaxKernel for Utf8Array<O> {
234    type Scalar<'a> = &'a str;
235
236    #[inline(always)]
237    fn min_ignore_nan_kernel(&self) -> Option<Self::Scalar<'_>> {
238        self.to_binary().min_ignore_nan_kernel().map(|s| unsafe {
239            // SAFETY: the lifetime is the same, and it is valid UTF-8.
240            #[allow(clippy::transmute_bytes_to_str)]
241            std::mem::transmute::<&[u8], &str>(s)
242        })
243    }
244
245    #[inline(always)]
246    fn max_ignore_nan_kernel(&self) -> Option<Self::Scalar<'_>> {
247        self.to_binary().max_ignore_nan_kernel().map(|s| unsafe {
248            // SAFETY: the lifetime is the same, and it is valid UTF-8.
249            #[allow(clippy::transmute_bytes_to_str)]
250            std::mem::transmute::<&[u8], &str>(s)
251        })
252    }
253
254    #[inline(always)]
255    fn min_propagate_nan_kernel(&self) -> Option<Self::Scalar<'_>> {
256        self.min_ignore_nan_kernel()
257    }
258
259    #[inline(always)]
260    fn max_propagate_nan_kernel(&self) -> Option<Self::Scalar<'_>> {
261        self.max_ignore_nan_kernel()
262    }
263}