Sort improve

benches/growable.rs

@@ -14,7 +14,7 @@ fn add_benchmark(c: &mut Criterion) {
     let i32_array = create_primitive_array::<i32>(1026 * 10, DataType::Int32, 0.0);
     c.bench_function("growable::primitive::non_null::non_null", |b| {
         b.iter(|| {
-            let mut a = GrowablePrimitive::new(&[&i32_array], false, 1026 * 10);
+            let mut a = GrowablePrimitive::new(vec![&i32_array], false, 1026 * 10);
             values
                 .clone()
                 .into_iter()
@@ -25,7 +25,7 @@ fn add_benchmark(c: &mut Criterion) {
     let i32_array = create_primitive_array::<i32>(1026 * 10, DataType::Int32, 0.0);
     c.bench_function("growable::primitive::non_null::null", |b| {
         b.iter(|| {
-            let mut a = GrowablePrimitive::new(&[&i32_array], true, 1026 * 10);
+            let mut a = GrowablePrimitive::new(vec![&i32_array], true, 1026 * 10);
             values.clone().into_iter().for_each(|start| {
                 if start % 2 == 0 {
                     a.extend_validity(10);
@@ -41,7 +41,7 @@ fn add_benchmark(c: &mut Criterion) {
     let values = values.collect::<Vec<_>>();
     c.bench_function("growable::primitive::null::non_null", |b| {
         b.iter(|| {
-            let mut a = GrowablePrimitive::new(&[&i32_array], false, 1026 * 10);
+            let mut a = GrowablePrimitive::new(vec![&i32_array], false, 1026 * 10);
             values
                 .clone()
                 .into_iter()
@@ -50,7 +50,7 @@ fn add_benchmark(c: &mut Criterion) {
     });
     c.bench_function("growable::primitive::null::null", |b| {
         b.iter(|| {
-            let mut a = GrowablePrimitive::new(&[&i32_array], true, 1026 * 10);
+            let mut a = GrowablePrimitive::new(vec![&i32_array], true, 1026 * 10);
             values.clone().into_iter().for_each(|start| {
                 if start % 2 == 0 {
                     a.extend_validity(10);

benches/sort_kernel.rs

@@ -19,7 +19,7 @@
 extern crate criterion;
 use criterion::Criterion;
 
-use arrow2::compute::sort::{lexsort, sort, SortColumn, SortOptions};
+use arrow2::compute::sort::{lexsort, sort, sort_to_indices, SortColumn, SortOptions};
 use arrow2::util::bench_util::*;
 use arrow2::{array::*, datatypes::*};
 
@@ -42,6 +42,15 @@ fn bench_sort(arr_a: &dyn Array) {
     sort(criterion::black_box(arr_a), &SortOptions::default(), None).unwrap();
 }
 
+fn bench_sort_limit(arr_a: &dyn Array) {
+    let _: PrimitiveArray<u32> = sort_to_indices(
+        criterion::black_box(arr_a),
+        &SortOptions::default(),
+        Some(100),
+    )
+    .unwrap();
+}
+
 fn add_benchmark(c: &mut Criterion) {
     (10..=20).step_by(2).for_each(|log2_size| {
         let size = 2usize.pow(log2_size);
@@ -51,6 +60,10 @@ fn add_benchmark(c: &mut Criterion) {
             b.iter(|| bench_sort(&arr_a))
         });
 
+        c.bench_function(&format!("sort-limit 2^{} f32", log2_size), |b| {
+            b.iter(|| bench_sort_limit(&arr_a))
+        });
+
         let arr_b = create_primitive_array_with_seed::<f32>(size, DataType::Float32, 0.0, 43);
         c.bench_function(&format!("lexsort 2^{} f32", log2_size), |b| {
             b.iter(|| bench_lexsort(&arr_a, &arr_b))

benches/take_kernels.rs

@@ -35,7 +35,7 @@ fn create_random_index(size: usize, null_density: f32) -> PrimitiveArray<i32> {
     (0..size)
         .map(|_| {
             if rng.gen::<f32>() > null_density {
-                let value = rng.gen_range::<i32, _, _>(0i32, size as i32);
+                let value = rng.gen_range::<i32, _>(0i32..size as i32);
                 Some(value)
             } else {
                 None

src/array/specification.rs

@@ -3,14 +3,19 @@ use std::convert::TryFrom;
 use num::Num;
 
 use crate::{
-    buffer::Buffer,
+    buffer::{Buffer, MutableBuffer},
     types::{NativeType, NaturalDataType},
 };
 
 /// Trait describing any type that can be used to index a slot of an array.
 pub trait Index: NativeType + NaturalDataType {
     fn to_usize(&self) -> usize;
     fn from_usize(index: usize) -> Option<Self>;
+    fn is_usize() -> bool {
+        false
+    }
+
+    fn buffer_from_range(start: usize, end: usize) -> Option<MutableBuffer<Self>>;
 }
 
 /// Trait describing types that can be used as offsets as per Arrow specification.
@@ -71,6 +76,19 @@ impl Index for i32 {
     fn from_usize(value: usize) -> Option<Self> {
         Self::try_from(value).ok()
     }
+
+    fn buffer_from_range(start: usize, end: usize) -> Option<MutableBuffer<Self>> {
+        let start = Self::from_usize(start);
+        let end = Self::from_usize(end);
+        match (start, end) {
+            (Some(start), Some(end)) => unsafe {
+                Some(MutableBuffer::<Self>::from_trusted_len_iter_unchecked(
+                    start..end,
+                ))
+            },
+            _ => None,
+        }
+    }
 }
 
 impl Index for i64 {
@@ -83,6 +101,19 @@ impl Index for i64 {
     fn from_usize(value: usize) -> Option<Self> {
         Self::try_from(value).ok()
     }
+
+    fn buffer_from_range(start: usize, end: usize) -> Option<MutableBuffer<Self>> {
+        let start = Self::from_usize(start);
+        let end = Self::from_usize(end);
+        match (start, end) {
+            (Some(start), Some(end)) => unsafe {
+                Some(MutableBuffer::<Self>::from_trusted_len_iter_unchecked(
+                    start..end,
+                ))
+            },
+            _ => None,
+        }
+    }
 }
 
 impl Index for u32 {
@@ -95,6 +126,23 @@ impl Index for u32 {
     fn from_usize(value: usize) -> Option<Self> {
         Self::try_from(value).ok()
     }
+
+    fn is_usize() -> bool {
+        std::mem::size_of::<Self>() == std::mem::size_of::<usize>()
+    }
+
+    fn buffer_from_range(start: usize, end: usize) -> Option<MutableBuffer<Self>> {
+        let start = Self::from_usize(start);
+        let end = Self::from_usize(end);
+        match (start, end) {
+            (Some(start), Some(end)) => unsafe {
+                Some(MutableBuffer::<Self>::from_trusted_len_iter_unchecked(
+                    start..end,
+                ))
+            },
+            _ => None,
+        }
+    }
 }
 
 impl Index for u64 {
@@ -107,6 +155,23 @@ impl Index for u64 {
     fn from_usize(value: usize) -> Option<Self> {
         Self::try_from(value).ok()
     }
+
+    fn is_usize() -> bool {
+        std::mem::size_of::<Self>() == std::mem::size_of::<usize>()
+    }
+
+    fn buffer_from_range(start: usize, end: usize) -> Option<MutableBuffer<Self>> {
+        let start = Self::from_usize(start);
+        let end = Self::from_usize(end);
+        match (start, end) {
+            (Some(start), Some(end)) => unsafe {
+                Some(MutableBuffer::<Self>::from_trusted_len_iter_unchecked(
+                    start..end,
+                ))
+            },
+            _ => None,
+        }
+    }
 }
 
 #[inline]

src/buffer/mutable.rs

@@ -93,6 +93,15 @@ impl<T: NativeType> MutableBuffer<T> {
         }
     }
 
+    /// Allocates a new [MutableBuffer] with `len` and capacity to be at least `len` where
+    /// all bytes are not initialized
+    #[inline]
+    pub unsafe fn from_len(len: usize) -> Self {
+        let mut buffer = MutableBuffer::with_capacity(len);
+        buffer.set_len(len);
+        buffer
+    }
+
     /// Ensures that this buffer has at least `self.len + additional` bytes. This re-allocates iff
     /// `self.len + additional > capacity`.
     /// # Example

src/compute/merge_sort/mod.rs

@@ -77,7 +77,7 @@ use crate::error::Result;
 /// This representation is useful when building arrays in memory as it allows to memcopy slices of arrays.
 /// This is particularly useful in merge-sort because sorted arrays (passed to the merge-sort) are more likely
 /// to have contiguous blocks of sorted elements (than by random).
-type MergeSlice = (usize, usize, usize);
+pub type MergeSlice = (usize, usize, usize);
 
 /// Takes N arrays together through `slices` under the assumption that the slices have
 /// a total coverage of the arrays.
@@ -234,7 +234,7 @@ fn recursive_merge_sort(slices: &[&[MergeSlice]], comparator: &Comparator) -> Ve
 
 // An iterator adapter that merge-sorts two iterators of `MergeSlice` into a single `MergeSlice`
 // such that the resulting `MergeSlice`s are ordered according to `comparator`.
-struct MergeSortSlices<'a, L, R>
+pub struct MergeSortSlices<'a, L, R>
 where
     L: Iterator<Item = &'a MergeSlice>,
     R: Iterator<Item = &'a MergeSlice>,
@@ -291,6 +291,29 @@ where
             None => self.right = None,
         }
     }
+
+    /// Collect the MergeSortSlices to be a vec for reusing
+    #[warn(dead_code)]
+    pub fn to_vec(self, limit: Option<usize>) -> Vec<MergeSlice> {
+        match limit {
+            Some(limit) => {
+                let mut v = Vec::with_capacity(limit);
+                let mut current_len = 0;
+                for (index, start, len) in self {
+                    if len + current_len >= limit {
+                        v.push((index, start, limit - current_len));
+                        break;
+                    } else {
+                        v.push((index, start, len));
+                    }
+                    current_len += len;
+                }
+
+                v
+            }
+            None => self.into_iter().collect(),
+        }
+    }
 }
 
 impl<'a, L, R> Iterator for MergeSortSlices<'a, L, R>
@@ -426,7 +449,11 @@ where
 /// Given two iterators of slices representing two sets of sorted [`Array`]s, and a `comparator` bound to those [`Array`]s,
 /// returns a new iterator of slices denoting how to `take` slices from each of the arrays such that the resulting
 /// array is sorted according to `comparator`
-fn merge_sort_slices<'a, L: Iterator<Item = &'a MergeSlice>, R: Iterator<Item = &'a MergeSlice>>(
+pub fn merge_sort_slices<
+    'a,
+    L: Iterator<Item = &'a MergeSlice>,
+    R: Iterator<Item = &'a MergeSlice>,
+>(
     lhs: L,
     rhs: R,
     comparator: &'a Comparator,
@@ -439,7 +466,7 @@ type Comparator<'a> = Box<dyn Fn(usize, usize, usize, usize) -> Ordering + 'a>;
 type IsValid<'a> = Box<dyn Fn(usize) -> bool + 'a>;
 
 /// returns a comparison function between any two arrays of each pair of arrays, according to `SortOptions`.
-fn build_comparator<'a>(
+pub fn build_comparator<'a>(
     pairs: &'a [(&'a [&'a dyn Array], &SortOptions)],
 ) -> Result<Comparator<'a>> {
     // prepare the comparison function of _values_ between all pairs of arrays
@@ -558,6 +585,22 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn test_merge_slices_to_vec() -> Result<()> {
+        let a0: &dyn Array = &Int32Array::from_slice(&[0, 2, 4, 6, 8]);
+        let a1: &dyn Array = &Int32Array::from_slice(&[1, 3, 5, 7, 9]);
+
+        let options = SortOptions::default();
+        let arrays = vec![a0, a1];
+        let pairs = vec![(arrays.as_ref(), &options)];
+        let comparator = build_comparator(&pairs)?;
+
+        let slices = merge_sort_slices(once(&(0, 0, 5)), once(&(1, 0, 5)), &comparator);
+        let vec = slices.to_vec(Some(5));
+        assert_eq!(vec, [(0, 0, 1), (1, 0, 1), (0, 1, 1), (1, 1, 1), (0, 2, 1)]);
+        Ok(())
+    }
+
     #[test]
     fn test_merge_4_i32() -> Result<()> {
         let a0: &dyn Array = &Int32Array::from_slice(&[0, 1]);

src/compute/sort/boolean.rs

@@ -26,7 +26,7 @@ pub fn sort_boolean<I: Index>(
     if !descending {
         valids.sort_by(|a, b| a.1.cmp(&b.1));
     } else {
-        valids.sort_by(|a, b| a.1.cmp(&b.1).reverse());
+        valids.sort_by(|a, b| b.1.cmp(&a.1));
         // reverse to keep a stable ordering
         nulls.reverse();
     }
@@ -45,6 +45,7 @@ pub fn sort_boolean<I: Index>(
     // un-efficient; there are much more performant ways of sorting nulls above, anyways.
     if let Some(limit) = limit {
         values.truncate(limit);
+        values.shrink_to_fit();
     }
 
     PrimitiveArray::<I>::from_data(I::DATA_TYPE, values.into(), None)