Skip to content

Commit

Permalink
feat: Sort kernel for RunArray (#3695)
Browse files Browse the repository at this point in the history
* Handle sliced array in run array iterator

* sort_to_indices for RunArray

* better loop

* sort for run array

* improve docs

* some minor tweaks

* doc fix

* format fix

* fix sort run to return all logical indices

* pr comment

* rename test function, pull sort run logic into a separate function

---------

Co-authored-by: ask <ask@local>
  • Loading branch information
askoa and ask committed Feb 23, 2023
1 parent e753dea commit ebe6f53
Show file tree
Hide file tree
Showing 4 changed files with 374 additions and 15 deletions.
28 changes: 22 additions & 6 deletions arrow-array/src/array/run_array.rs
Expand Up @@ -112,17 +112,37 @@ impl<R: RunEndIndexType> RunArray<R> {

/// Returns a reference to run_ends array
///
/// Note: any slicing of this array is not applied to the returned array
/// Note: any slicing of this [`RunArray`] array is not applied to the returned array
/// and must be handled separately
pub fn run_ends(&self) -> &PrimitiveArray<R> {
&self.run_ends
}

/// Returns a reference to values array
///
/// Note: any slicing of this [`RunArray`] array is not applied to the returned array
/// and must be handled separately
pub fn values(&self) -> &ArrayRef {
&self.values
}

/// Returns the physical index at which the array slice starts.
pub fn get_start_physical_index(&self) -> usize {
if self.offset() == 0 {
return 0;
}
self.get_zero_offset_physical_index(self.offset()).unwrap()
}

/// Returns the physical index at which the array slice ends.
pub fn get_end_physical_index(&self) -> usize {
if self.offset() + self.len() == Self::logical_len(&self.run_ends) {
return self.run_ends.len() - 1;
}
self.get_zero_offset_physical_index(self.offset() + self.len() - 1)
.unwrap()
}

/// Downcast this [`RunArray`] to a [`TypedRunArray`]
///
/// ```
Expand Down Expand Up @@ -230,11 +250,7 @@ impl<R: RunEndIndexType> RunArray<R> {
}

// Skip some physical indices based on offset.
let skip_value = if self.offset() > 0 {
self.get_zero_offset_physical_index(self.offset()).unwrap()
} else {
0
};
let skip_value = self.get_start_physical_index();

let mut physical_indices = vec![0; indices_len];

Expand Down
9 changes: 2 additions & 7 deletions arrow-array/src/run_iterator.rs
Expand Up @@ -57,13 +57,8 @@ where
{
/// create a new iterator
pub fn new(array: TypedRunArray<'a, R, V>) -> Self {
let current_front_physical: usize =
array.run_array().get_physical_index(0).unwrap();
let current_back_physical: usize = array
.run_array()
.get_physical_index(array.len() - 1)
.unwrap()
+ 1;
let current_front_physical = array.run_array().get_start_physical_index();
let current_back_physical = array.run_array().get_end_physical_index() + 1;
RunArrayIter {
array,
current_front_logical: array.offset(),
Expand Down

0 comments on commit ebe6f53

Please sign in to comment.