fix lexsort limit equal or greater than row_count

jorgecarleitao · May 28, 2022 · 42d300c · 42d300c
1 parent e976c1e
commit 42d300c
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
@@ -11,6 +11,13 @@ fn test_lex_sort_arrays(input: Vec<SortColumn>, expected: Vec<Box<dyn Array>>) {
         .map(|x| x.slice(0, 2))
         .collect::<Vec<_>>();
     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);
 }
 
 #[test]