polars_utils/
arg_min_max.rs

1use crate::float16::pf16;
2
3pub trait ArgMinMax {
4    fn argmin(&self) -> usize;
5
6    fn argmax(&self) -> usize;
7
8    /// Skips the values whose bit in the `validity` bitmap (starting at bit `offset`) is
9    /// not set. Returns `None` if no value is valid.
10    fn argmin_masked(&self, validity: &[u8], offset: usize) -> Option<usize>;
11
12    fn argmax_masked(&self, validity: &[u8], offset: usize) -> Option<usize>;
13}
14
15macro_rules! impl_argminmax {
16    ($T:ty) => {
17        impl ArgMinMax for $T {
18            fn argmin(&self) -> usize {
19                argminmax::ArgMinMax::argmin(self)
20            }
21
22            fn argmax(&self) -> usize {
23                argminmax::ArgMinMax::argmax(self)
24            }
25
26            fn argmin_masked(&self, validity: &[u8], offset: usize) -> Option<usize> {
27                argminmax::ArgMinMaxMasked::argmin_masked(self, validity, offset)
28            }
29
30            fn argmax_masked(&self, validity: &[u8], offset: usize) -> Option<usize> {
31                argminmax::ArgMinMaxMasked::argmax_masked(self, validity, offset)
32            }
33        }
34    };
35}
36
37impl_argminmax!(&[u8]);
38impl_argminmax!(&[u16]);
39impl_argminmax!(&[u32]);
40impl_argminmax!(&[u64]);
41impl_argminmax!(&[u128]);
42impl_argminmax!(&[i8]);
43impl_argminmax!(&[i16]);
44impl_argminmax!(&[i32]);
45impl_argminmax!(&[i64]);
46impl_argminmax!(&[i128]);
47impl_argminmax!(&[f32]);
48impl_argminmax!(&[f64]);
49
50impl ArgMinMax for &[pf16] {
51    fn argmin(&self) -> usize {
52        let transmuted: &&[half::f16] = unsafe { std::mem::transmute(self) };
53        argminmax::ArgMinMax::argmin(transmuted)
54    }
55
56    fn argmax(&self) -> usize {
57        let transmuted: &&[half::f16] = unsafe { std::mem::transmute(self) };
58        argminmax::ArgMinMax::argmax(transmuted)
59    }
60
61    fn argmin_masked(&self, validity: &[u8], offset: usize) -> Option<usize> {
62        let transmuted: &&[half::f16] = unsafe { std::mem::transmute(self) };
63        argminmax::ArgMinMaxMasked::argmin_masked(transmuted, validity, offset)
64    }
65
66    fn argmax_masked(&self, validity: &[u8], offset: usize) -> Option<usize> {
67        let transmuted: &&[half::f16] = unsafe { std::mem::transmute(self) };
68        argminmax::ArgMinMaxMasked::argmax_masked(transmuted, validity, offset)
69    }
70}
polars_utils/arg_min_max.rs

polars_utils/
arg_min_max.rs
