diff --git a/src/compute/partition.rs b/src/compute/partition.rs index 15d46fd5de5..dc04da81db1 100644 --- a/src/compute/partition.rs +++ b/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> + 'a> { +pub fn lexicographical_partition_ranges( + columns: &[SortColumn], +) -> Result>> { 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, } -impl<'a> LexicographicalPartitionIterator<'a> { - fn try_new(columns: &'a [SortColumn]) -> Result { +impl LexicographicalPartitionIterator { + fn try_new(columns: &[SortColumn]) -> Result { 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; fn next(&mut self) -> Option { diff --git a/src/compute/sort/lex_sort.rs b/src/compute/sort/lex_sort.rs index 1b0182e1202..77856902c4c 100644 --- a/src/compute/sort/lex_sort.rs +++ b/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 bool + 'a>; +type IsValid = Box bool + Send + Sync>; /// One column to be used in lexicographical sort #[derive(Clone, Debug)] @@ -71,15 +72,14 @@ pub fn lexsort( #[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 Ordering + 'a>; - -pub(crate) fn build_compare(array: &dyn Array, sort_option: SortOptions) -> Result { +pub(crate) fn build_compare(array: &dyn Array, sort_option: SortOptions) -> Result { let is_valid = build_is_valid(array); let comparator = ord::build_compare(array, array)?; @@ -150,10 +150,10 @@ pub fn lexsort_to_indices( // map arrays to comparators let comparators = columns .iter() - .map(|column| -> Result { + .map(|column| -> Result { build_compare(column.values, column.options.unwrap_or_default()) }) - .collect::>>()?; + .collect::>>()?; let lex_comparator = |a_idx: &I, b_idx: &I| -> Ordering { let a_idx = a_idx.to_usize(); diff --git a/src/compute/sort/mod.rs b/src/compute/sort/mod.rs index af41e27efe3..43fd2070cb5 100644 --- a/src/compute/sort/mod.rs +++ b/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 {