sundy-li commented on a change in pull request #9602: URL: https://github.com/apache/arrow/pull/9602#discussion_r593899532
########## File path: rust/arrow/src/compute/kernels/sort.rs ########## @@ -686,90 +815,124 @@ pub fn lexsort_to_indices(columns: &[SortColumn]) -> Result<UInt32Array> { }; let mut value_indices = (0..row_count).collect::<Vec<usize>>(); - value_indices.sort_by(lex_comparator); + let mut len = value_indices.len(); + + if let Some(limit) = limit { + len = limit.min(len); + } + sort_by(&mut value_indices, len, lex_comparator); Ok(UInt32Array::from( - value_indices - .into_iter() - .map(|i| i as u32) + (&value_indices)[0..len] + .iter() + .map(|i| *i as u32) .collect::<Vec<u32>>(), )) } +pub fn partial_sort<T, F>(v: &mut [T], limit: usize, mut is_less: F) +where + F: FnMut(&T, &T) -> Ordering, +{ + let (before, _mid, _after) = v.select_nth_unstable_by(limit, &mut is_less); + before.sort_unstable_by(is_less); Review comment: There are some reasons for using unstable_sort. 1. `select_nth_unstable_by` is unstable, so using `before.sort_stable_by`cadon't ensure it's stable. Currently we do not have select_nth_stable_by 2. `sort_unstable` is faster than sort_stable, refer to [doc](https://doc.rust-lang.org/std/primitive.slice.html#method.sort_unstable) ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: us...@infra.apache.org