Fixed lexsort limit equal or greater than row_count (#1021)

jorgecarleitao · May 28, 2022 · 89005e9 · 89005e9
1 parent e976c1e
commit 89005e9
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 1 deletion.
diff --git a/src/compute/sort/lex_sort.rs b/src/compute/sort/lex_sort.rs
@@ -198,7 +198,12 @@ pub fn lexsort_to_indices_impl<I: Index>(
 
     if let Some(limit) = limit {
         let limit = limit.min(row_count);
-        let (before, _, _) = values.select_nth_unstable_by(limit, lex_comparator);
+        let before = if limit < row_count {
+            let (before, _, _) = values.select_nth_unstable_by(limit, lex_comparator);
+            before
+        } else {
+            &mut values[..]
+        };
         before.sort_unstable_by(lex_comparator);
         values.truncate(limit);
         values.shrink_to_fit();

diff --git a/tests/it/compute/sort/lex_sort.rs b/tests/it/compute/sort/lex_sort.rs
@@ -5,6 +5,13 @@ fn test_lex_sort_arrays(input: Vec<SortColumn>, expected: Vec<Box<dyn Array>>) {
     let sorted = lexsort::<i32>(&input, None).unwrap();
     assert_eq!(sorted, expected);
 
+    let sorted = lexsort::<i32>(&input, Some(4)).unwrap();
+    let expected = expected
+        .into_iter()
+        .map(|x| x.slice(0, 4))
+        .collect::<Vec<_>>();
+    assert_eq!(sorted, expected);
+
     let sorted = lexsort::<i32>(&input, Some(2)).unwrap();
     let expected = expected
         .into_iter()