perf(rust, python): SIMD accelerated arg_min/arg_max (via argminmax)

polars/polars-ops/Cargo.toml

@@ -10,6 +10,7 @@ description = "More operations on polars data structures"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+argminmax = { git = "https://github.com/jvdd/argminmax" }
 arrow.workspace = true
 base64 = { version = "0.21", optional = true }
 either.workspace = true

polars/polars-ops/src/series/ops/arg_min_max.rs

@@ -1,7 +1,9 @@
+use argminmax::ArgMinMax;
+use arrow::array::PrimitiveArray;
 use arrow::bitmap::utils::{BitChunkIterExact, BitChunksExact};
 use arrow::bitmap::Bitmap;
 use polars_core::series::IsSorted;
-use polars_core::with_match_physical_numeric_polars_type;
+use polars_core::{with_match_physical_numeric_polars_type, with_match_physical_numeric_type};
 
 use super::*;
 
@@ -20,7 +22,7 @@ impl ArgAgg for Series {
         match s.dtype() {
             Utf8 => {
                 let ca = s.utf8().unwrap();
-                arg_min(ca)
+                arg_min_str(ca)
             }
             Boolean => {
                 let ca = s.bool().unwrap();
@@ -29,10 +31,12 @@ impl ArgAgg for Series {
             dt if dt.is_numeric() => {
                 with_match_physical_numeric_polars_type!(s.dtype(), |$T| {
                     let ca: &ChunkedArray<$T> = s.as_ref().as_ref().as_ref();
-                    if let Ok(vals) = ca.cont_slice() {
-                        arg_min_slice(vals, ca.is_sorted_flag2())
+                    if ca.is_empty() { // because argminmax assumes not empty
+                        None
+                    } else if let Ok(vals) = ca.cont_slice() {
+                        arg_min_numeric_slice(vals, ca.is_sorted_flag2())
                     } else {
-                        arg_min(ca)
+                        arg_min_numeric(ca)
                     }
                 })
             }
@@ -46,7 +50,7 @@ impl ArgAgg for Series {
         match s.dtype() {
             Utf8 => {
                 let ca = s.utf8().unwrap();
-                arg_max(ca)
+                arg_max_str(ca)
             }
             Boolean => {
                 let ca = s.bool().unwrap();
@@ -55,10 +59,12 @@ impl ArgAgg for Series {
             dt if dt.is_numeric() => {
                 with_match_physical_numeric_polars_type!(s.dtype(), |$T| {
                     let ca: &ChunkedArray<$T> = s.as_ref().as_ref().as_ref();
-                    if let Ok(vals) = ca.cont_slice() {
-                        arg_max_slice(vals, ca.is_sorted_flag2())
+                    if ca.is_empty() { // because argminmax assumes not empty
+                        None
+                    } else if let Ok(vals) = ca.cont_slice() {
+                        arg_max_numeric_slice(vals, ca.is_sorted_flag2())
                     } else {
-                        arg_max(ca)
+                        arg_max_numeric(ca)
                     }
                 })
             }
@@ -67,7 +73,7 @@ impl ArgAgg for Series {
     }
 }
 
-fn arg_max_bool(ca: &BooleanChunked) -> Option<usize> {
+pub(crate) fn arg_max_bool(ca: &BooleanChunked) -> Option<usize> {
     if ca.is_empty() {
         None
     } else if ca.null_count() == ca.len() {
@@ -188,12 +194,7 @@ fn first_unset_bit(mask: &Bitmap) -> usize {
     }
 }
 
-fn arg_min<'a, T>(ca: &'a ChunkedArray<T>) -> Option<usize>
-where
-    T: PolarsDataType,
-    &'a ChunkedArray<T>: IntoIterator,
-    <&'a ChunkedArray<T> as IntoIterator>::Item: PartialOrd,
-{
+fn arg_min_str(ca: &Utf8Chunked) -> Option<usize> {
     match ca.is_sorted_flag2() {
         IsSorted::Ascending => Some(0),
         IsSorted::Descending => Some(ca.len() - 1),
@@ -205,12 +206,7 @@ where
     }
 }
 
-pub(crate) fn arg_max<'a, T>(ca: &'a ChunkedArray<T>) -> Option<usize>
-where
-    T: PolarsDataType,
-    &'a ChunkedArray<T>: IntoIterator,
-    <&'a ChunkedArray<T> as IntoIterator>::Item: PartialOrd,
-{
+fn arg_max_str(ca: &Utf8Chunked) -> Option<usize> {
     match ca.is_sorted_flag2() {
         IsSorted::Ascending => Some(ca.len() - 1),
         IsSorted::Descending => Some(0),
@@ -222,32 +218,94 @@ where
     }
 }
 
-fn arg_min_slice<T>(vals: &[T], is_sorted: IsSorted) -> Option<usize>
+fn arg_min_numeric<'a, T>(ca: &'a ChunkedArray<T>) -> Option<usize>
+where
+    T: PolarsDataType,
+    &'a ChunkedArray<T>: IntoIterator,
+    <&'a ChunkedArray<T> as IntoIterator>::Item: PartialOrd,
+{
+    match ca.is_sorted_flag2() {
+        IsSorted::Ascending => Some(0),
+        IsSorted::Descending => Some(ca.len() - 1),
+        IsSorted::Not => {
+            with_match_physical_numeric_type!(ca.dtype(), |$TN| {
+                ca.chunks().iter().fold((None, None, 0), |acc, chunk| {
+                    if chunk.is_empty() {
+                        return acc;
+                    }
+                    let arr: &PrimitiveArray<$TN> = chunk.as_any().downcast_ref().unwrap();
+                    let chunk_min_idx: usize = arr.values().as_slice().argmin();
+                    let chunk_min_val: $TN = arr.value(chunk_min_idx);
+                    match acc {
+                        (None, None, offset) => (Some(chunk_min_idx + offset), Some(chunk_min_val), offset + chunk.len()),
+                        (Some(acc_min_idx), Some(acc_min_val), offset) => {
+                            if chunk_min_val < acc_min_val {
+                                (Some(chunk_min_idx + offset), Some(chunk_min_val), offset + chunk.len())
+                            } else {
+                                (Some(acc_min_idx), Some(acc_min_val), offset + chunk.len())
+                            }
+                        }
+                        _ => unreachable!(),
+                    }
+                }).0
+            })
+        }
+    }
+}
+
+fn arg_max_numeric<'a, T>(ca: &'a ChunkedArray<T>) -> Option<usize>
+where
+    T: PolarsDataType,
+    &'a ChunkedArray<T>: IntoIterator,
+    <&'a ChunkedArray<T> as IntoIterator>::Item: PartialOrd,
+{
+    match ca.is_sorted_flag2() {
+        IsSorted::Ascending => Some(ca.len() - 1),
+        IsSorted::Descending => Some(0),
+        IsSorted::Not => {
+            with_match_physical_numeric_type!(ca.dtype(), |$TN| {
+                ca.chunks().iter().fold((None, None, 0), |acc, chunk| {
+                    if chunk.is_empty() {
+                        return acc;
+                    }
+                    let arr: &PrimitiveArray<$TN> = chunk.as_any().downcast_ref().unwrap();
+                    let chunk_max_idx: usize = arr.values().as_slice().argmax();
+                    let chunk_max_val: $TN = arr.value(chunk_max_idx);
+                    match acc {
+                        (None, None, offset) => (Some(chunk_max_idx + offset), Some(chunk_max_val), offset + chunk.len()),
+                        (Some(acc_max_idx), Some(acc_max_val), offset) => {
+                            if chunk_max_val > acc_max_val {
+                                (Some(chunk_max_idx + offset), Some(chunk_max_val), offset + chunk.len())
+                            } else {
+                                (Some(acc_max_idx), Some(acc_max_val), offset + chunk.len())
+                            }
+                        }
+                        _ => unreachable!(),
+                    }
+                }).0
+            })
+        }
+    }
+}
+
+fn arg_min_numeric_slice<T>(vals: &[T], is_sorted: IsSorted) -> Option<usize>
 where
-    T: PartialOrd,
+    for<'a> &'a [T]: ArgMinMax,
 {
     match is_sorted {
         IsSorted::Ascending => Some(0),
         IsSorted::Descending => Some(vals.len() - 1),
-        IsSorted::Not => vals
-            .iter()
-            .enumerate()
-            .reduce(|acc, (idx, val)| if acc.1 > val { (idx, val) } else { acc })
-            .map(|tpl| tpl.0),
+        IsSorted::Not => Some(vals.argmin()), // assumes not empty
     }
 }
 
-fn arg_max_slice<T>(vals: &[T], is_sorted: IsSorted) -> Option<usize>
+fn arg_max_numeric_slice<T>(vals: &[T], is_sorted: IsSorted) -> Option<usize>
 where
-    T: PartialOrd,
+    for<'a> &'a [T]: ArgMinMax,
 {
     match is_sorted {
         IsSorted::Ascending => Some(vals.len() - 1),
         IsSorted::Descending => Some(0),
-        IsSorted::Not => vals
-            .iter()
-            .enumerate()
-            .reduce(|acc, (idx, val)| if acc.1 < val { (idx, val) } else { acc })
-            .map(|tpl| tpl.0),
+        IsSorted::Not => Some(vals.argmax()), // assumes not empty
     }
 }

polars/polars-ops/src/series/ops/is_first.rs

@@ -6,7 +6,7 @@ use polars_arrow::utils::CustomIterTools;
 use polars_core::prelude::*;
 use polars_core::with_match_physical_integer_polars_type;
 
-use crate::series::ops::arg_min_max::arg_max;
+use crate::series::ops::arg_min_max::arg_max_bool;
 
 fn is_first_numeric<T>(ca: &ChunkedArray<T>) -> BooleanChunked
 where
@@ -47,7 +47,7 @@ fn is_first_bin(ca: &BinaryChunked) -> BooleanChunked {
 fn is_first_boolean(ca: &BooleanChunked) -> BooleanChunked {
     let mut out = MutableBitmap::with_capacity(ca.len());
     out.extend_constant(ca.len(), false);
-    if let Some(index) = arg_max(ca) {
+    if let Some(index) = arg_max_bool(ca) {
         out.set(index, true)
     }
     if let Some(index) = ca.first_non_null() {