diff --git a/src/compute/sort/lex_sort.rs b/src/compute/sort/lex_sort.rs index 114093f49a8..c598cfd3abd 100644 --- a/src/compute/sort/lex_sort.rs +++ b/src/compute/sort/lex_sort.rs @@ -198,7 +198,12 @@ pub fn lexsort_to_indices_impl( 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 index bf0b146a8fe..24ead89f977 100644 --- 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, expected: Vec>) { .map(|x| x.slice(0, 2)) .collect::>(); assert_eq!(sorted, expected); + + let sorted = lexsort::(&input, Some(4)).unwrap(); + let expected = expected + .into_iter() + .map(|x| x.slice(0, 4)) + .collect::>(); + assert_eq!(sorted, expected); } #[test]