polars_ops/series/ops/
arg_min_max.rs

1use polars_core::chunked_array::arg_min_max::{
2    arg_max_binary, arg_max_bool, arg_max_numeric, arg_max_str, arg_min_binary, arg_min_bool,
3    arg_min_numeric, arg_min_str,
4};
5#[cfg(feature = "dtype-categorical")]
6use polars_core::chunked_array::arg_min_max::{arg_max_cat, arg_min_cat};
7#[cfg(feature = "dtype-categorical")]
8use polars_core::with_match_categorical_physical_type;
9
10use super::*;
11
12/// Argmin/ Argmax
13pub trait ArgAgg {
14    /// Get the index of the minimal value
15    fn arg_min(&self) -> Option<usize>;
16    /// Get the index of the maximal value
17    fn arg_max(&self) -> Option<usize>;
18}
19
20macro_rules! with_match_physical_numeric_polars_type {(
21    $key_type:expr, | $_:tt $T:ident | $($body:tt)*
22) => ({
23    macro_rules! __with_ty__ {( $_ $T:ident ) => ( $($body)* )}
24    use DataType::*;
25    match $key_type {
26        #[cfg(feature = "dtype-i8")]
27        Int8 => __with_ty__! { Int8Type },
28        #[cfg(feature = "dtype-i16")]
29        Int16 => __with_ty__! { Int16Type },
30        Int32 => __with_ty__! { Int32Type },
31        Int64 => __with_ty__! { Int64Type },
32        #[cfg(feature = "dtype-i128")]
33        Int128 => __with_ty__! { Int128Type },
34        #[cfg(feature = "dtype-u8")]
35        UInt8 => __with_ty__! { UInt8Type },
36        #[cfg(feature = "dtype-u16")]
37        UInt16 => __with_ty__! { UInt16Type },
38        UInt32 => __with_ty__! { UInt32Type },
39        UInt64 => __with_ty__! { UInt64Type },
40        #[cfg(feature = "dtype-u128")]
41        UInt128 => __with_ty__! { UInt128Type },
42        #[cfg(feature = "dtype-f16")]
43        Float16 => __with_ty__! { Float16Type },
44        Float32 => __with_ty__! { Float32Type },
45        Float64 => __with_ty__! { Float64Type },
46        dt => panic!("not implemented for dtype {:?}", dt),
47    }
48})}
49
50impl ArgAgg for Series {
51    fn arg_min(&self) -> Option<usize> {
52        use DataType::*;
53        let phys_s = self.to_physical_repr();
54        match self.dtype() {
55            #[cfg(feature = "dtype-categorical")]
56            Categorical(cats, _) => {
57                with_match_categorical_physical_type!(cats.physical(), |$C| {
58                    arg_min_cat(self.cat::<$C>().unwrap())
59                })
60            },
61            #[cfg(feature = "dtype-categorical")]
62            Enum(_, _) => phys_s.arg_min(),
63            #[cfg(feature = "dtype-decimal")]
64            Decimal(_, _) => phys_s.arg_min(),
65            Date | Datetime(_, _) | Duration(_) | Time => phys_s.arg_min(),
66            String => arg_min_str(self.str().unwrap()),
67            Binary => arg_min_binary(self.binary().unwrap()),
68            Boolean => arg_min_bool(self.bool().unwrap()),
69            dt if dt.is_primitive_numeric() => {
70                with_match_physical_numeric_polars_type!(phys_s.dtype(), |$T| {
71                    let ca: &ChunkedArray<$T> = phys_s.as_ref().as_ref().as_ref();
72                    arg_min_numeric(ca)
73                })
74            },
75            _ => None,
76        }
77    }
78
79    fn arg_max(&self) -> Option<usize> {
80        use DataType::*;
81        let phys_s = self.to_physical_repr();
82        match self.dtype() {
83            #[cfg(feature = "dtype-categorical")]
84            Categorical(cats, _) => {
85                with_match_categorical_physical_type!(cats.physical(), |$C| {
86                    arg_max_cat(self.cat::<$C>().unwrap())
87                })
88            },
89            #[cfg(feature = "dtype-categorical")]
90            Enum(_, _) => phys_s.arg_max(),
91            #[cfg(feature = "dtype-decimal")]
92            Decimal(_, _) => phys_s.arg_max(),
93            Date | Datetime(_, _) | Duration(_) | Time => phys_s.arg_max(),
94            String => arg_max_str(self.str().unwrap()),
95            Binary => arg_max_binary(self.binary().unwrap()),
96            Boolean => arg_max_bool(self.bool().unwrap()),
97            dt if dt.is_primitive_numeric() => {
98                with_match_physical_numeric_polars_type!(phys_s.dtype(), |$T| {
99                    let ca: &ChunkedArray<$T> = phys_s.as_ref().as_ref().as_ref();
100                    arg_max_numeric(ca)
101                })
102            },
103            _ => None,
104        }
105    }
106}
polars_ops/series/ops/arg_min_max.rs

polars_ops/series/ops/
arg_min_max.rs
