[arrow-rs] branch master updated: Fix primitive sort when input contains more nulls than the given sort limit (#954)

[alamb]

This is an automated email from the ASF dual-hosted git repository.

alamb pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git


The following commit(s) were added to refs/heads/master by this push:
     new 02f3ec8  Fix primitive sort when input contains more nulls than the 
given sort limit (#954)
02f3ec8 is described below

commit 02f3ec8c0549105a67eca79c1ab597c37f399e5e
Author: Jörn Horstmann <g...@jhorstmann.net>
AuthorDate: Thu Nov 18 23:26:24 2021 +0100

    Fix primitive sort when input contains more nulls than the given sort limit 
(#954)
---
 arrow/src/compute/kernels/sort.rs | 27 +++++++++++++++++++--------
 1 file changed, 19 insertions(+), 8 deletions(-)

diff --git a/arrow/src/compute/kernels/sort.rs 
b/arrow/src/compute/kernels/sort.rs
index 88c7785..6a72224 100644
--- a/arrow/src/compute/kernels/sort.rs
+++ b/arrow/src/compute/kernels/sort.rs
@@ -1032,13 +1032,11 @@ fn sort_valids<T, U>(
 ) where
     T: ?Sized + Copy,
 {
-    let nulls_len = nulls.len();
+    let valids_len = valids.len();
     if !descending {
-        sort_unstable_by(valids, len.saturating_sub(nulls_len), |a, b| 
cmp(a.1, b.1));
+        sort_unstable_by(valids, len.min(valids_len), |a, b| cmp(a.1, b.1));
     } else {
-        sort_unstable_by(valids, len.saturating_sub(nulls_len), |a, b| {
-            cmp(a.1, b.1).reverse()
-        });
+        sort_unstable_by(valids, len.min(valids_len), |a, b| cmp(a.1, 
b.1).reverse());
         // reverse to keep a stable ordering
         nulls.reverse();
     }
@@ -1050,13 +1048,13 @@ fn sort_valids_array<T>(
     nulls: &mut [T],
     len: usize,
 ) {
-    let nulls_len = nulls.len();
+    let valids_len = valids.len();
     if !descending {
-        sort_unstable_by(valids, len.saturating_sub(nulls_len), |a, b| {
+        sort_unstable_by(valids, len.min(valids_len), |a, b| {
             cmp_array(a.1.as_ref(), b.1.as_ref())
         });
     } else {
-        sort_unstable_by(valids, len.saturating_sub(nulls_len), |a, b| {
+        sort_unstable_by(valids, len.min(valids_len), |a, b| {
             cmp_array(a.1.as_ref(), b.1.as_ref()).reverse()
         });
         // reverse to keep a stable ordering
@@ -1556,6 +1554,19 @@ mod tests {
     }
 
     #[test]
+    fn test_sort_to_indices_primitive_more_nulls_than_limit() {
+        test_sort_to_indices_primitive_arrays::<Int32Type>(
+            vec![None, None, Some(3), None, Some(1), None, Some(2)],
+            Some(SortOptions {
+                descending: false,
+                nulls_first: false,
+            }),
+            Some(2),
+            vec![4, 6],
+        );
+    }
+
+    #[test]
     fn test_sort_boolean() {
         // boolean
         test_sort_to_indices_boolean_arrays(