fix[rust]: handle NaNs in argsort (#4497)

pola-rs · Aug 19, 2022 · 5dc495c · 5dc495c
1 parent 90f1feb
commit 5dc495c
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 8 deletions.
diff --git a/polars/polars-arrow/src/kernels/rolling/mod.rs b/polars/polars-arrow/src/kernels/rolling/mod.rs
@@ -24,7 +24,7 @@ type Len = usize;
 
 pub fn compare_fn_nan_min<T>(a: &T, b: &T) -> Ordering
 where
-    T: PartialOrd + IsFloat + NativeType,
+    T: PartialOrd + IsFloat,
 {
     if T::is_float() {
         match (a.is_nan(), b.is_nan()) {
@@ -41,9 +41,9 @@ where
     }
 }
 
-fn compare_fn_nan_max<T>(a: &T, b: &T) -> Ordering
+pub fn compare_fn_nan_max<T>(a: &T, b: &T) -> Ordering
 where
-    T: PartialOrd + IsFloat + NativeType,
+    T: PartialOrd + IsFloat,
 {
     if T::is_float() {
         match (a.is_nan(), b.is_nan()) {

diff --git a/polars/polars-core/src/chunked_array/ops/sort/mod.rs b/polars/polars-core/src/chunked_array/ops/sort/mod.rs
@@ -11,6 +11,7 @@ use std::iter::FromIterator;
 use arrow::{bitmap::MutableBitmap, buffer::Buffer};
 use num::Float;
 use polars_arrow::array::default_arrays::FromDataUtf8;
+use polars_arrow::kernels::rolling::compare_fn_nan_max;
 use polars_arrow::prelude::{FromData, ValueSize};
 use polars_arrow::trusted_len::PushUnchecked;
 use rayon::prelude::*;
@@ -82,14 +83,14 @@ fn sort_branch<T, Fd, Fr>(
 #[cfg(feature = "private")]
 pub fn argsort_no_nulls<Idx, T>(slice: &mut [(Idx, T)], reverse: bool)
 where
-    T: PartialOrd + Send,
+    T: PartialOrd + Send + IsFloat,
     Idx: PartialOrd + Send,
 {
     argsort_branch(
         slice,
         reverse,
-        |(_, a), (_, b)| a.partial_cmp(b).unwrap(),
-        |(_, a), (_, b)| b.partial_cmp(a).unwrap(),
+        |(_, a), (_, b)| compare_fn_nan_max(a, b),
+        |(_, a), (_, b)| compare_fn_nan_max(b, a),
     );
 }
 

diff --git a/polars/polars-ops/src/series/ops/search_sorted.rs b/polars/polars-ops/src/series/ops/search_sorted.rs
@@ -1,6 +1,6 @@
 use std::cmp::Ordering;
 
-use polars_arrow::kernels::rolling::compare_fn_nan_min;
+use polars_arrow::kernels::rolling::compare_fn_nan_max;
 use polars_arrow::prelude::*;
 use polars_core::downcast_as_macro_arg_physical;
 use polars_core::export::num::NumCast;
@@ -27,7 +27,7 @@ where
         // - `mid < size`: `mid` is limited by `[left; right)` bound.
         let cmp = match unsafe { taker.get_unchecked(mid as usize) } {
             None => Ordering::Less,
-            Some(value) => compare_fn_nan_min(&value, &search_value),
+            Some(value) => compare_fn_nan_max(&value, &search_value),
         };
 
         // The reason why we use if/else control flow rather than match

diff --git a/py-polars/tests/test_sort.py b/py-polars/tests/test_sort.py
@@ -209,3 +209,20 @@ def test_sorted_fast_paths() -> None:
     assert rev.to_list() == [None, 3, 2, 1]
     assert rev.sort(reverse=True).to_list() == [None, 3, 2, 1]
     assert rev.sort().to_list() == [None, 1, 2, 3]
+
+
+def test_argsort_rank_nans() -> None:
+    assert (
+        pl.DataFrame(
+            {
+                "val": [1.0, float("NaN")],
+            }
+        )
+        .with_columns(
+            [
+                pl.col("val").rank().alias("rank"),
+                pl.col("val").argsort().alias("argsort"),
+            ]
+        )
+        .select(["rank", "argsort"])
+    ).to_dict(False) == {"rank": [1.0, 2.0], "argsort": [0, 1]}