Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

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

Merged
merged 9 commits into from
Apr 14, 2023
Merged

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

Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
39053d2
perf: use argminmax for faster arg_min and arg_max
jvdd Apr 8, 2023
8b0dfa9
chore: formatting
jvdd Apr 8, 2023
f45e1a1
fix: use arg_max_bool
jvdd Apr 8, 2023
6c19b6a
refactor: use downcast_iter instead of macro
jvdd Apr 9, 2023
b68e191
chore: use argminmax latest release + improve code readability
jvdd Apr 10, 2023
e5beb0e
fix: handle chunks with nulls in numeric arg_min and arg_max
jvdd Apr 10, 2023
4ebff93
chore: fix linting error
jvdd Apr 10, 2023
24dac92
fix: do not use unstable SIMD by default
jvdd Apr 14, 2023
d8267c3
Merge branch 'main' into argminmax_integration
jvdd Apr 14, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions polars/polars-ops/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 = { version = "0.6" }
arrow.workspace = true
base64 = { version = "0.21", optional = true }
either.workspace = true
Expand Down
171 changes: 136 additions & 35 deletions polars/polars-ops/src/series/ops/arg_min_max.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
use argminmax::ArgMinMax;
use arrow::array::Array;
use arrow::bitmap::utils::{BitChunkIterExact, BitChunksExact};
use arrow::bitmap::Bitmap;
use polars_core::series::IsSorted;
Expand All @@ -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();
Expand All @@ -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)
}
})
}
Expand All @@ -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();
Expand All @@ -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)
}
})
}
Expand All @@ -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() {
Expand Down Expand Up @@ -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),
Expand All @@ -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),
Expand All @@ -222,32 +218,137 @@ 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: PolarsNumericType,
for<'b> &'b [T::Native]: ArgMinMax,
{
match ca.is_sorted_flag2() {
IsSorted::Ascending => Some(0),
IsSorted::Descending => Some(ca.len() - 1),
IsSorted::Not => {
ca.downcast_iter()
.fold((None, None, 0), |acc, arr| {
if arr.len() == 0 {
return acc;
}
let chunk_min_idx: Option<usize>;
let chunk_min_val: Option<T::Native>;
if arr.null_count() > 0 {
// When there are nulls, we should compare Option<T::Native>
chunk_min_val = None; // because None < Some(_)
chunk_min_idx = arr
.into_iter()
.enumerate()
.reduce(|acc, (idx, val)| if acc.1 > val { (idx, val) } else { acc })
.map(|tpl| tpl.0);
Comment on lines +239 to +244
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Basically we should return the index of the first null value

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good point, we can speed this up. I can do this later. 👍

} else {
// When no nulls & array not empty => we can use fast argminmax
let min_idx: usize = arr.values().as_slice().argmin();
chunk_min_idx = Some(min_idx);
chunk_min_val = Some(arr.value(min_idx));
}

let new_offset: usize = acc.2 + arr.len();
match acc {
(Some(_), Some(_), offset) => {
if chunk_min_val < acc.1 {
match chunk_min_idx {
Some(idx) => (Some(idx + offset), chunk_min_val, new_offset),
None => (acc.0, acc.1, new_offset),
}
} else {
(acc.0, acc.1, new_offset)
}
}
(None, None, offset) => match chunk_min_idx {
Some(idx) => (Some(idx + offset), chunk_min_val, new_offset),
None => (None, None, new_offset),
},
_ => unreachable!(),
}
})
.0
}
}
}

fn arg_max_numeric<'a, T>(ca: &'a ChunkedArray<T>) -> Option<usize>
where
T: PolarsNumericType,
for<'b> &'b [T::Native]: ArgMinMax,
{
match ca.is_sorted_flag2() {
IsSorted::Ascending => Some(ca.len() - 1),
IsSorted::Descending => Some(0),
IsSorted::Not => {
ca.downcast_iter()
.fold((None, None, 0), |acc, arr| {
if arr.len() == 0 {
return acc;
}
let chunk_max_idx: Option<usize>;
let chunk_max_val: Option<T::Native>;
if arr.null_count() > 0 {
// When there are nulls, we should compare Option<T::Native>
chunk_max_idx = arr
.into_iter()
.enumerate()
.reduce(|acc, (idx, val)| if acc.1 < val { (idx, val) } else { acc })
.map(|tpl| tpl.0);
chunk_max_val = match chunk_max_idx {
Some(idx) => Some(arr.value(idx)),
None => None,
};
} else {
// When no nulls & array not empty => we can use fast argminmax
let max_idx: usize = arr.values().as_slice().argmax();
chunk_max_idx = Some(max_idx);
chunk_max_val = Some(arr.value(max_idx));
}

let new_offset: usize = acc.2 + arr.len();
match acc {
(Some(_), Some(_), offset) => {
if chunk_max_val > acc.1 {
match chunk_max_idx {
Some(idx) => (Some(idx + offset), chunk_max_val, new_offset),
_ => unreachable!(), // because None < Some(_)
}
} else {
(acc.0, acc.1, new_offset)
}
}
(None, None, offset) => match chunk_max_idx {
Some(idx) => (Some(idx + offset), chunk_max_val, new_offset),
None => (None, None, new_offset),
},
_ => 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
}
}
4 changes: 2 additions & 2 deletions polars/polars-ops/src/series/ops/is_first.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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() {
Expand Down
12 changes: 11 additions & 1 deletion py-polars/Cargo.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.