From 89005e9602c8561de52b5df21fb6bbaff3d6cde1 Mon Sep 17 00:00:00 2001
From: baishen <baishen2009@gmail.com>
Date: Sun, 29 May 2022 02:55:10 +0800
Subject: [PATCH] Fixed lexsort limit equal or greater than row_count (#1021)

---
 src/compute/sort/lex_sort.rs      | 7 ++++++-
 tests/it/compute/sort/lex_sort.rs | 7 +++++++
 2 files changed, 13 insertions(+), 1 deletion(-)

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<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
index bf0b146a8fe..a24f116dc8c 100644
--- 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()