Use DynComparator for lexsort and partition (#637)

src/compute/partition.rs

@@ -17,7 +17,8 @@
 
 //! Defines partition kernel for [`crate::array::Array`]
 
-use crate::compute::sort::{build_compare, Compare, SortColumn};
+use crate::array::ord::DynComparator;
+use crate::compute::sort::{build_compare, SortColumn};
 use crate::error::{ArrowError, Result};
 use std::cmp::Ordering;
 use std::iter::Iterator;
@@ -32,22 +33,22 @@ use std::ops::Range;
 /// The returned vec would be of size k where k is cardinality of the sorted values; Consecutive
 /// values will be connected: (a, b) and (b, c), where start = 0 and end = n for the first and last
 /// range.
-pub fn lexicographical_partition_ranges<'a>(
-    columns: &'a [SortColumn],
-) -> Result<impl Iterator<Item = Range<usize>> + 'a> {
+pub fn lexicographical_partition_ranges(
+    columns: &[SortColumn],
+) -> Result<impl Iterator<Item = Range<usize>>> {
     LexicographicalPartitionIterator::try_new(columns)
 }
 
-struct LexicographicalPartitionIterator<'a> {
-    comparator: Compare<'a>,
+struct LexicographicalPartitionIterator {
+    comparator: DynComparator,
     num_rows: usize,
     previous_partition_point: usize,
     partition_point: usize,
     value_indices: Vec<usize>,
 }
 
-impl<'a> LexicographicalPartitionIterator<'a> {
-    fn try_new(columns: &'a [SortColumn]) -> Result<Self> {
+impl LexicographicalPartitionIterator {
+    fn try_new(columns: &[SortColumn]) -> Result<Self> {
         if columns.is_empty() {
             return Err(ArrowError::InvalidArgumentError(
                 "Sort requires at least one column".to_string(),
@@ -87,7 +88,7 @@ impl<'a> LexicographicalPartitionIterator<'a> {
     }
 }
 
-impl<'a> Iterator for LexicographicalPartitionIterator<'a> {
+impl Iterator for LexicographicalPartitionIterator {
     type Item = Range<usize>;
 
     fn next(&mut self) -> Option<Self::Item> {

src/compute/sort/lex_sort.rs

@@ -9,8 +9,9 @@ use crate::{
 };
 
 use super::{sort_to_indices, SortOptions};
+use crate::array::ord::DynComparator;
 
-type IsValid<'a> = Box<dyn Fn(usize) -> bool + 'a>;
+type IsValid = Box<dyn Fn(usize) -> bool + Send + Sync>;
 
 /// One column to be used in lexicographical sort
 #[derive(Clone, Debug)]
@@ -71,15 +72,14 @@ pub fn lexsort<I: Index>(
 #[inline]
 fn build_is_valid(array: &dyn Array) -> IsValid {
     if let Some(validity) = array.validity() {
+        let validity = validity.clone();
         Box::new(move |x| unsafe { validity.get_bit_unchecked(x) })
     } else {
         Box::new(move |_| true)
     }
 }
 
-pub(crate) type Compare<'a> = Box<dyn Fn(usize, usize) -> Ordering + 'a>;
-
-pub(crate) fn build_compare(array: &dyn Array, sort_option: SortOptions) -> Result<Compare> {
+pub(crate) fn build_compare(array: &dyn Array, sort_option: SortOptions) -> Result<DynComparator> {
     let is_valid = build_is_valid(array);
     let comparator = ord::build_compare(array, array)?;
 
@@ -150,10 +150,10 @@ pub fn lexsort_to_indices<I: Index>(
     // map arrays to comparators
     let comparators = columns
         .iter()
-        .map(|column| -> Result<Compare> {
+        .map(|column| -> Result<DynComparator> {
             build_compare(column.values, column.options.unwrap_or_default())
         })
-        .collect::<Result<Vec<Compare>>>()?;
+        .collect::<Result<Vec<DynComparator>>>()?;
 
     let lex_comparator = |a_idx: &I, b_idx: &I| -> Ordering {
         let a_idx = a_idx.to_usize();

src/compute/sort/mod.rs

@@ -18,7 +18,7 @@ mod lex_sort;
 mod primitive;
 mod utf8;
 
-pub(crate) use lex_sort::{build_compare, Compare};
+pub(crate) use lex_sort::build_compare;
 pub use lex_sort::{lexsort, lexsort_to_indices, SortColumn};
 
 macro_rules! dyn_sort {