This is an automated email from the ASF dual-hosted git repository.
tustvold pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/master by this push:
new 9131c30a4 feat: take kernel for RunArray (#3622)
9131c30a4 is described below
commit 9131c30a40cfb14d3e9454eac0f9b80e000b152b
Author: askoa <[email protected]>
AuthorDate: Mon Feb 6 06:03:02 2023 -0500
feat: take kernel for RunArray (#3622)
* Rebase to master branch
* Add take_run kernel and include benchmarks for take_run and primitive run
accessor.
* fix ci issues
* fix ci issues
* fix clippy issues
* fix clippy
* Alternative approach to find physical indices for given logical indices
* Remove unused code, refactor benchmarks
* minor fixes
* some refactor
* add some comments
* doc fixes
* change benchmkar parameters
* add test for run_iterator
* add comments
* Fix some PR suggestions
* incorporte pr suggestion
---------
Co-authored-by: ask <ask@local>
Co-authored-by: Raphael Taylor-Davies <[email protected]>
---
arrow-array/src/array/run_array.rs | 195 ++++++++++++++++-----
arrow-array/src/cast.rs | 126 +++++++++++++
arrow-array/src/run_iterator.rs | 77 +++++++-
arrow-select/src/take.rs | 98 ++++++++++-
arrow/Cargo.toml | 17 ++
arrow/benches/primitive_run_accessor.rs | 57 ++++++
arrow/benches/primitive_run_take.rs | 79 +++++++++
arrow/benches/string_run_builder.rs | 22 +--
...tring_run_builder.rs => string_run_iterator.rs} | 62 ++++---
arrow/benches/take_kernels.rs | 7 +
arrow/src/util/bench_util.rs | 68 +++++++
11 files changed, 716 insertions(+), 92 deletions(-)
diff --git a/arrow-array/src/array/run_array.rs
b/arrow-array/src/array/run_array.rs
index 8cc1f676b..2e378c90f 100644
--- a/arrow-array/src/array/run_array.rs
+++ b/arrow-array/src/array/run_array.rs
@@ -143,6 +143,96 @@ impl<R: RunEndIndexType> RunArray<R> {
values,
})
}
+
+ /// Returns index to the physical array for the given index to the logical
array.
+ /// Performs a binary search on the run_ends array for the input index.
+ #[inline]
+ pub fn get_physical_index(&self, logical_index: usize) -> Option<usize> {
+ if logical_index >= self.len() {
+ return None;
+ }
+ let mut st: usize = 0;
+ let mut en: usize = self.run_ends().len();
+ while st + 1 < en {
+ let mid: usize = (st + en) / 2;
+ if logical_index
+ < unsafe {
+ // Safety:
+ // The value of mid will always be between 1 and len - 1,
+ // where len is length of run ends array.
+ // This is based on the fact that `st` starts with 0 and
+ // `en` starts with len. The condition `st + 1 < en`
ensures
+ // `st` and `en` differs atleast by two. So the value of
`mid`
+ // will never be either `st` or `en`
+ self.run_ends().value_unchecked(mid - 1).as_usize()
+ }
+ {
+ en = mid
+ } else {
+ st = mid
+ }
+ }
+ Some(st)
+ }
+
+ /// Returns the physical indices of the input logical indices. Returns
error if any of the logical
+ /// index cannot be converted to physical index. The logical indices are
sorted and iterated along
+ /// with run_ends array to find matching physical index. The approach used
here was chosen over
+ /// finding physical index for each logical index using binary search
using the function
+ /// `get_physical_index`. Running benchmarks on both approaches showed
that the approach used here
+ /// scaled well for larger inputs.
+ /// See
<https://github.com/apache/arrow-rs/pull/3622#issuecomment-1407753727> for more
details.
+ #[inline]
+ pub fn get_physical_indices<I>(
+ &self,
+ logical_indices: &[I],
+ ) -> Result<Vec<usize>, ArrowError>
+ where
+ I: ArrowNativeType,
+ {
+ let indices_len = logical_indices.len();
+
+ // `ordered_indices` store index into `logical_indices` and can be used
+ // to iterate `logical_indices` in sorted order.
+ let mut ordered_indices: Vec<usize> = (0..indices_len).collect();
+
+ // Instead of sorting `logical_idices` directly, sort the
`ordered_indices`
+ // whose values are index of `logical_indices`
+ ordered_indices.sort_unstable_by(|lhs, rhs| {
+ logical_indices[*lhs]
+ .partial_cmp(&logical_indices[*rhs])
+ .unwrap()
+ });
+
+ let mut physical_indices = vec![0; indices_len];
+
+ let mut ordered_index = 0_usize;
+ for (physical_index, run_end) in
self.run_ends.values().iter().enumerate() {
+ // Get the run end index of current physical index
+ let run_end_value = run_end.as_usize();
+
+ // All the `logical_indices` that are less than current run end
index
+ // belongs to current physical index.
+ while ordered_index < indices_len
+ && logical_indices[ordered_indices[ordered_index]].as_usize()
+ < run_end_value
+ {
+ physical_indices[ordered_indices[ordered_index]] =
physical_index;
+ ordered_index += 1;
+ }
+ }
+
+ // If there are input values >= run_ends.last_value then we'll not be
able to convert
+ // all logical indices to physical indices.
+ if ordered_index < logical_indices.len() {
+ let logical_index =
+ logical_indices[ordered_indices[ordered_index]].as_usize();
+ return Err(ArrowError::InvalidArgumentError(format!(
+ "Cannot convert all logical indices to physical indices. The
logical index cannot be converted is {logical_index}.",
+ )));
+ }
+ Ok(physical_indices)
+ }
}
impl<R: RunEndIndexType> From<ArrayData> for RunArray<R> {
@@ -348,37 +438,6 @@ impl<'a, R: RunEndIndexType, V> TypedRunArray<'a, R, V> {
pub fn values(&self) -> &'a V {
self.values
}
-
- /// Returns index to the physcial array for the given index to the logical
array.
- /// Performs a binary search on the run_ends array for the input index.
- #[inline]
- pub fn get_physical_index(&self, logical_index: usize) -> Option<usize> {
- if logical_index >= self.run_array.len() {
- return None;
- }
- let mut st: usize = 0;
- let mut en: usize = self.run_ends().len();
- while st + 1 < en {
- let mid: usize = (st + en) / 2;
- if logical_index
- < unsafe {
- // Safety:
- // The value of mid will always be between 1 and len - 1,
- // where len is length of run ends array.
- // This is based on the fact that `st` starts with 0 and
- // `en` starts with len. The condition `st + 1 < en`
ensures
- // `st` and `en` differs atleast by two. So the value of
`mid`
- // will never be either `st` or `en`
- self.run_ends().value_unchecked(mid - 1).as_usize()
- }
- {
- en = mid
- } else {
- st = mid
- }
- }
- Some(st)
- }
}
impl<'a, R: RunEndIndexType, V: Sync> Array for TypedRunArray<'a, R, V> {
@@ -417,7 +476,7 @@ where
}
unsafe fn value_unchecked(&self, logical_index: usize) -> Self::Item {
- let physical_index = self.get_physical_index(logical_index).unwrap();
+ let physical_index =
self.run_array.get_physical_index(logical_index).unwrap();
self.values().value_unchecked(physical_index)
}
}
@@ -447,13 +506,15 @@ mod tests {
use super::*;
use crate::builder::PrimitiveRunBuilder;
+ use crate::cast::as_primitive_array;
use crate::types::{Int16Type, Int32Type, Int8Type, UInt32Type};
use crate::{Array, Int16Array, Int32Array, StringArray};
- fn build_input_array(approx_size: usize) -> Vec<Option<i32>> {
+ fn build_input_array(size: usize) -> Vec<Option<i32>> {
// The input array is created by shuffling and repeating
// the seed values random number of times.
let mut seed: Vec<Option<i32>> = vec![
+ None,
None,
None,
Some(1),
@@ -462,26 +523,32 @@ mod tests {
Some(4),
Some(5),
Some(6),
+ Some(7),
+ Some(8),
+ Some(9),
];
+ let mut result: Vec<Option<i32>> = Vec::with_capacity(size);
let mut ix = 0;
- let mut result: Vec<Option<i32>> = Vec::with_capacity(approx_size);
let mut rng = thread_rng();
- while result.len() < approx_size {
+ // run length can go up to 8. Cap the max run length for smaller
arrays to size / 2.
+ let max_run_length = 8_usize.min(1_usize.max(size / 2));
+ while result.len() < size {
// shuffle the seed array if all the values are iterated.
if ix == 0 {
seed.shuffle(&mut rng);
}
- // repeat the items between 1 and 7 times.
- let num = rand::thread_rng().gen_range(1..8);
+ // repeat the items between 1 and 8 times. Cap the length for
smaller sized arrays
+ let num =
+
max_run_length.min(rand::thread_rng().gen_range(1..=max_run_length));
for _ in 0..num {
result.push(seed[ix]);
}
ix += 1;
- if ix == 8 {
+ if ix == seed.len() {
ix = 0
}
}
- println!("Size of input array: {}", result.len());
+ result.resize(size, None);
result
}
@@ -718,14 +785,62 @@ mod tests {
let run_array = builder.finish();
let typed = run_array.downcast::<PrimitiveArray<Int32Type>>().unwrap();
+ // Access every index and check if the value in the input array
matches returned value.
for (i, inp_val) in input_array.iter().enumerate() {
if let Some(val) = inp_val {
let actual = typed.value(i);
assert_eq!(*val, actual)
} else {
- let physical_ix = typed.get_physical_index(i).unwrap();
+ let physical_ix = run_array.get_physical_index(i).unwrap();
assert!(typed.values().is_null(physical_ix));
};
}
}
+
+ #[test]
+ fn test_get_physical_indices() {
+ // Test for logical lengths starting from 10 to 250 increasing by 10
+ for logical_len in (0..250).step_by(10) {
+ let input_array = build_input_array(logical_len);
+
+ // create run array using input_array
+ let mut builder = PrimitiveRunBuilder::<Int32Type,
Int32Type>::new();
+ builder.extend(input_array.clone().into_iter());
+
+ let run_array = builder.finish();
+ let physical_values_array =
+ as_primitive_array::<Int32Type>(run_array.values());
+
+ // create an array consisting of all the indices repeated twice
and shuffled.
+ let mut logical_indices: Vec<u32> = (0_u32..(logical_len as
u32)).collect();
+ // add same indices once more
+ logical_indices.append(&mut logical_indices.clone());
+ let mut rng = thread_rng();
+ logical_indices.shuffle(&mut rng);
+
+ let physical_indices =
+ run_array.get_physical_indices(&logical_indices).unwrap();
+
+ assert_eq!(logical_indices.len(), physical_indices.len());
+
+ // check value in logical index in the input_array matches
physical index in typed_run_array
+ logical_indices
+ .iter()
+ .map(|f| f.as_usize())
+ .zip(physical_indices.iter())
+ .for_each(|(logical_ix, physical_ix)| {
+ let expected = input_array[logical_ix];
+ match expected {
+ Some(val) => {
+
assert!(physical_values_array.is_valid(*physical_ix));
+ let actual =
physical_values_array.value(*physical_ix);
+ assert_eq!(val, actual);
+ }
+ None => {
+
assert!(physical_values_array.is_null(*physical_ix))
+ }
+ };
+ });
+ }
+ }
}
diff --git a/arrow-array/src/cast.rs b/arrow-array/src/cast.rs
index 02d5432c1..4bae4932c 100644
--- a/arrow-array/src/cast.rs
+++ b/arrow-array/src/cast.rs
@@ -95,6 +95,53 @@ macro_rules! downcast_integer {
};
}
+/// Given one or more expressions evaluating to an integer [`DataType`]
invokes the provided macro
+/// `m` with the corresponding integer [`RunEndIndexType`], followed by any
additional arguments
+///
+/// ```
+/// # use arrow_array::{downcast_primitive, ArrowPrimitiveType,
downcast_run_end_index};
+/// # use arrow_schema::{DataType, Field};
+///
+/// macro_rules! run_end_size_helper {
+/// ($t:ty, $o:ty) => {
+/// std::mem::size_of::<<$t as ArrowPrimitiveType>::Native>() as $o
+/// };
+/// }
+///
+/// fn run_end_index_size(t: &DataType) -> u8 {
+/// match t {
+/// DataType::RunEndEncoded(k, _) => downcast_run_end_index! {
+/// k.data_type() => (run_end_size_helper, u8),
+/// _ => unreachable!(),
+/// },
+/// _ => u8::MAX,
+/// }
+/// }
+///
+///
assert_eq!(run_end_index_size(&DataType::RunEndEncoded(Box::new(Field::new("a",
DataType::Int32, false)), Box::new(Field::new("b", DataType::Utf8, true)))), 4);
+///
assert_eq!(run_end_index_size(&DataType::RunEndEncoded(Box::new(Field::new("a",
DataType::Int64, false)), Box::new(Field::new("b", DataType::Utf8, true)))), 8);
+///
assert_eq!(run_end_index_size(&DataType::RunEndEncoded(Box::new(Field::new("a",
DataType::Int16, false)), Box::new(Field::new("b", DataType::Utf8, true)))), 2);
+/// ```
+///
+/// [`DataType`]: arrow_schema::DataType
+#[macro_export]
+macro_rules! downcast_run_end_index {
+ ($($data_type:expr),+ => ($m:path $(, $args:tt)*), $($($p:pat),+ =>
$fallback:expr $(,)*)*) => {
+ match ($($data_type),+) {
+ $crate::repeat_pat!(arrow_schema::DataType::Int16,
$($data_type),+) => {
+ $m!($crate::types::Int16Type $(, $args)*)
+ }
+ $crate::repeat_pat!(arrow_schema::DataType::Int32,
$($data_type),+) => {
+ $m!($crate::types::Int32Type $(, $args)*)
+ }
+ $crate::repeat_pat!(arrow_schema::DataType::Int64,
$($data_type),+) => {
+ $m!($crate::types::Int64Type $(, $args)*)
+ }
+ $(($($p),+) => $fallback,)*
+ }
+ };
+}
+
/// Given one or more expressions evaluating to primitive [`DataType`] invokes
the provided macro
/// `m` with the corresponding [`ArrowPrimitiveType`], followed by any
additional arguments
///
@@ -449,6 +496,85 @@ where
.expect("Unable to downcast to dictionary array")
}
+/// Force downcast of an [`Array`], such as an [`ArrayRef`] to
+/// [`RunArray<T>`], panic'ing on failure.
+///
+/// # Example
+///
+/// ```
+/// # use arrow_array::{ArrayRef, RunArray};
+/// # use arrow_array::cast::as_run_array;
+/// # use arrow_array::types::Int32Type;
+///
+/// let arr: RunArray<Int32Type> = vec![Some("foo")].into_iter().collect();
+/// let arr: ArrayRef = std::sync::Arc::new(arr);
+/// let run_array: &RunArray<Int32Type> = as_run_array::<Int32Type>(&arr);
+/// ```
+pub fn as_run_array<T>(arr: &dyn Array) -> &RunArray<T>
+where
+ T: RunEndIndexType,
+{
+ arr.as_any()
+ .downcast_ref::<RunArray<T>>()
+ .expect("Unable to downcast to run array")
+}
+
+#[macro_export]
+#[doc(hidden)]
+macro_rules! downcast_run_array_helper {
+ ($t:ty, $($values:ident),+, $e:block) => {{
+ $(let $values = $crate::cast::as_run_array::<$t>($values);)+
+ $e
+ }};
+}
+
+/// Downcast an [`Array`] to a [`RunArray`] based on its [`DataType`], accepts
+/// a number of subsequent patterns to match the data type
+///
+/// ```
+/// # use arrow_array::{Array, StringArray, downcast_run_array,
cast::as_string_array};
+/// # use arrow_schema::DataType;
+///
+/// fn print_strings(array: &dyn Array) {
+/// downcast_run_array!(
+/// array => match array.values().data_type() {
+/// DataType::Utf8 => {
+/// for v in array.downcast::<StringArray>().unwrap() {
+/// println!("{:?}", v);
+/// }
+/// }
+/// t => println!("Unsupported run array value type {}", t),
+/// },
+/// DataType::Utf8 => {
+/// for v in as_string_array(array) {
+/// println!("{:?}", v);
+/// }
+/// }
+/// t => println!("Unsupported datatype {}", t)
+/// )
+/// }
+/// ```
+///
+/// [`DataType`]: arrow_schema::DataType
+#[macro_export]
+macro_rules! downcast_run_array {
+ ($values:ident => $e:expr, $($p:pat => $fallback:expr $(,)*)*) => {
+ downcast_run_array!($values => {$e} $($p => $fallback)*)
+ };
+
+ ($values:ident => $e:block $($p:pat => $fallback:expr $(,)*)*) => {
+ match $values.data_type() {
+ arrow_schema::DataType::RunEndEncoded(k, _) => {
+ $crate::downcast_run_end_index! {
+ k.data_type() => ($crate::downcast_run_array_helper,
$values, $e),
+ k => unreachable!("unsupported run end index type: {}", k)
+ }
+ }
+ $($p => $fallback,)*
+ }
+ }
+}
+
/// Force downcast of an [`Array`], such as an [`ArrayRef`] to
/// [`GenericListArray<T>`], panic'ing on failure.
pub fn as_generic_list_array<S: OffsetSizeTrait>(
diff --git a/arrow-array/src/run_iterator.rs b/arrow-array/src/run_iterator.rs
index 6a7b785fe..8bad85a9f 100644
--- a/arrow-array/src/run_iterator.rs
+++ b/arrow-array/src/run_iterator.rs
@@ -149,7 +149,6 @@ where
// reason the next value can be accessed by decrementing physical
index once.
self.current_end_physical -= 1;
}
-
Some(if self.array.values().is_null(self.current_end_physical) {
None
} else {
@@ -180,6 +179,8 @@ where
#[cfg(test)]
mod tests {
+ use rand::{seq::SliceRandom, thread_rng, Rng};
+
use crate::{
array::{Int32Array, StringArray},
builder::PrimitiveRunBuilder,
@@ -187,6 +188,48 @@ mod tests {
Int64RunArray,
};
+ fn build_input_array(size: usize) -> Vec<Option<i32>> {
+ // The input array is created by shuffling and repeating
+ // the seed values random number of times.
+ let mut seed: Vec<Option<i32>> = vec![
+ None,
+ None,
+ None,
+ Some(1),
+ Some(2),
+ Some(3),
+ Some(4),
+ Some(5),
+ Some(6),
+ Some(7),
+ Some(8),
+ Some(9),
+ ];
+ let mut result: Vec<Option<i32>> = Vec::with_capacity(size);
+ let mut ix = 0;
+ let mut rng = thread_rng();
+ // run length can go up to 8. Cap the max run length for smaller
arrays to size / 2.
+ let max_run_length = 8_usize.min(1_usize.max(size / 2));
+ while result.len() < size {
+ // shuffle the seed array if all the values are iterated.
+ if ix == 0 {
+ seed.shuffle(&mut rng);
+ }
+ // repeat the items between 1 and 8 times. Cap the length for
smaller sized arrays
+ let num =
+
max_run_length.min(rand::thread_rng().gen_range(1..=max_run_length));
+ for _ in 0..num {
+ result.push(seed[ix]);
+ }
+ ix += 1;
+ if ix == seed.len() {
+ ix = 0
+ }
+ }
+ result.resize(size, None);
+ result
+ }
+
#[test]
fn test_primitive_array_iter_round_trip() {
let mut input_vec = vec![
@@ -239,6 +282,38 @@ mod tests {
assert_eq!(None, iter.next());
}
+ #[test]
+ fn test_run_iterator_comprehensive() {
+ // Test forward and backward iterator for different array lengths.
+ let logical_lengths = vec![1_usize, 2, 3, 4, 15, 16, 17, 63, 64, 65];
+
+ for logical_len in logical_lengths {
+ let input_array = build_input_array(logical_len);
+
+ let mut run_array_builder =
+ PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
+ run_array_builder.extend(input_array.iter().copied());
+ let run_array = run_array_builder.finish();
+ let typed_array = run_array.downcast::<Int32Array>().unwrap();
+
+ // test forward iterator
+ let mut input_iter = input_array.iter().copied();
+ let mut run_array_iter = typed_array.into_iter();
+ for _ in 0..logical_len {
+ assert_eq!(input_iter.next(), run_array_iter.next());
+ }
+ assert_eq!(None, run_array_iter.next());
+
+ // test reverse iterator
+ let mut input_iter = input_array.iter().rev().copied();
+ let mut run_array_iter = typed_array.into_iter().rev();
+ for _ in 0..logical_len {
+ assert_eq!(input_iter.next(), run_array_iter.next());
+ }
+ assert_eq!(None, run_array_iter.next());
+ }
+ }
+
#[test]
fn test_string_array_iter_round_trip() {
let input_vec = vec!["ab", "ab", "ba", "cc", "cc"];
diff --git a/arrow-select/src/take.rs b/arrow-select/src/take.rs
index d8989fa48..f8668b56e 100644
--- a/arrow-select/src/take.rs
+++ b/arrow-select/src/take.rs
@@ -19,13 +19,15 @@
use std::sync::Arc;
-use arrow_array::types::*;
use arrow_array::*;
+use arrow_array::{builder::PrimitiveRunBuilder, types::*};
use arrow_buffer::{bit_util, ArrowNativeType, Buffer, MutableBuffer};
use arrow_data::{ArrayData, ArrayDataBuilder};
use arrow_schema::{ArrowError, DataType, Field};
-use arrow_array::cast::{as_generic_binary_array, as_largestring_array,
as_string_array};
+use arrow_array::cast::{
+ as_generic_binary_array, as_largestring_array, as_primitive_array,
as_string_array,
+};
use num::{ToPrimitive, Zero};
/// Take elements by index from [Array], creating a new [Array] from those
indexes.
@@ -201,6 +203,10 @@ where
values => Ok(Arc::new(take_dict(values, indices)?)),
t => unimplemented!("Take not supported for dictionary type {:?}",
t)
}
+ DataType::RunEndEncoded(_, _) => downcast_run_array! {
+ values => Ok(Arc::new(take_run(values, indices)?)),
+ t => unimplemented!("Take not supported for run type {:?}", t)
+ }
DataType::Binary => {
Ok(Arc::new(take_bytes(as_generic_binary_array::<i32>(values),
indices)?))
}
@@ -810,6 +816,72 @@ where
Ok(DictionaryArray::<T>::from(data))
}
+macro_rules! primitive_run_take {
+ ($t:ty, $o:ty, $indices:ident, $value:ident) => {
+ take_primitive_run_values::<$o, $t>(
+ $indices,
+ as_primitive_array::<$t>($value.values()),
+ )
+ };
+}
+
+/// `take` implementation for run arrays
+///
+/// Finds physical indices for the given logical indices and builds output run
array
+/// by taking values in the input run array at the physical indices.
+/// for e.g. an input `RunArray{ run_ends = [2,4,6,8], values=[1,2,1,2] }` and
`indices=[2,7]`
+/// would be converted to `physical_indices=[1,3]` which will be used to build
+/// output `RunArray{ run_ends=[2], values=[2] }`
+
+fn take_run<T, I>(
+ run_array: &RunArray<T>,
+ logical_indices: &PrimitiveArray<I>,
+) -> Result<RunArray<T>, ArrowError>
+where
+ T: RunEndIndexType,
+ T::Native: num::Num,
+ I: ArrowPrimitiveType,
+ I::Native: ToPrimitive,
+{
+ match run_array.data_type() {
+ DataType::RunEndEncoded(_, fl) => {
+ let physical_indices =
+ run_array.get_physical_indices(logical_indices.values())?;
+
+ downcast_primitive! {
+ fl.data_type() => (primitive_run_take, T, physical_indices,
run_array),
+ dt => Err(ArrowError::NotYetImplemented(format!("take_run is
not implemented for {dt:?}")))
+ }
+ }
+ dt => Err(ArrowError::InvalidArgumentError(format!(
+ "Expected DataType::RunEndEncoded found {dt:?}"
+ ))),
+ }
+}
+
+// Builds a `RunArray` by taking values from given array for the given indices.
+fn take_primitive_run_values<R, V>(
+ physical_indices: Vec<usize>,
+ values: &PrimitiveArray<V>,
+) -> Result<RunArray<R>, ArrowError>
+where
+ R: RunEndIndexType,
+ V: ArrowPrimitiveType,
+{
+ let mut builder = PrimitiveRunBuilder::<R, V>::new();
+ let values_len = values.len();
+ for ix in physical_indices {
+ if ix >= values_len {
+ return Err(ArrowError::InvalidArgumentError("The requested index
{ix} is out of bounds for values array with length {values_len}".to_string()));
+ } else if values.is_null(ix) {
+ builder.append_null()
+ } else {
+ builder.append_value(values.value(ix))
+ }
+ }
+ Ok(builder.finish())
+}
+
/// Takes/filters a list array's inner data using the offsets of the list
array.
///
/// Where a list array has indices `[0,2,5,10]`, taking indices of `[2,0]`
returns
@@ -2086,6 +2158,28 @@ mod tests {
assert_eq!(null_buf.as_slice(), &[0b11111111]);
}
+ #[test]
+ fn test_take_runs() {
+ let logical_array: Vec<i32> = vec![1_i32, 1, 2, 2, 1, 1, 2, 2, 1, 1,
2, 2];
+
+ let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
+ builder.extend(logical_array.into_iter().map(Some));
+ let run_array = builder.finish();
+
+ let take_indices: PrimitiveArray<Int32Type> =
+ vec![2, 7, 10].into_iter().collect();
+
+ let take_out = take_run(&run_array, &take_indices).unwrap();
+
+ assert_eq!(take_out.len(), 3);
+
+ assert_eq!(take_out.run_ends().len(), 1);
+ assert_eq!(take_out.run_ends().value(0), 3);
+
+ let take_out_values =
as_primitive_array::<Int32Type>(take_out.values());
+ assert_eq!(take_out_values.value(0), 2);
+ }
+
#[test]
fn test_take_value_index_from_fixed_list() {
let list = FixedSizeListArray::from_iter_primitive::<Int32Type, _, _>(
diff --git a/arrow/Cargo.toml b/arrow/Cargo.toml
index 57e3907a2..f86ec09a9 100644
--- a/arrow/Cargo.toml
+++ b/arrow/Cargo.toml
@@ -102,6 +102,8 @@ half = { version = "2.1", default-features = false,
features = ["num-traits"] }
rand = { version = "0.8", default-features = false, features = ["std",
"std_rng"] }
tempfile = { version = "3", default-features = false }
+
+
[build-dependencies]
[[example]]
@@ -240,6 +242,21 @@ harness = false
[[bench]]
name = "string_run_builder"
harness = false
+required-features = ["test_utils"]
+
+[[bench]]
+name = "string_run_iterator"
+harness = false
+
+[[bench]]
+name = "primitive_run_accessor"
+harness = false
+required-features = ["test_utils"]
+
+[[bench]]
+name = "primitive_run_take"
+harness = false
+required-features = ["test_utils"]
[[bench]]
name = "substring_kernels"
diff --git a/arrow/benches/primitive_run_accessor.rs
b/arrow/benches/primitive_run_accessor.rs
new file mode 100644
index 000000000..868c314f9
--- /dev/null
+++ b/arrow/benches/primitive_run_accessor.rs
@@ -0,0 +1,57 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow::datatypes::Int32Type;
+use arrow::{array::PrimitiveArray,
util::bench_util::create_primitive_run_array};
+use arrow_array::ArrayAccessor;
+use criterion::{criterion_group, criterion_main, Criterion};
+
+fn criterion_benchmark(c: &mut Criterion) {
+ let mut group = c.benchmark_group("primitive_run_accessor");
+
+ let mut do_bench = |physical_array_len: usize, logical_array_len: usize| {
+ group.bench_function(
+ format!(
+ "(run_array_len:{logical_array_len},
physical_array_len:{physical_array_len})"),
+ |b| {
+ let run_array = create_primitive_run_array::<Int32Type,
Int32Type>(
+ logical_array_len,
+ physical_array_len,
+ );
+ let typed = run_array
+ .downcast::<PrimitiveArray<Int32Type>>()
+ .unwrap();
+ b.iter(|| {
+ for i in 0..logical_array_len {
+ let _ = unsafe { typed.value_unchecked(i) };
+ }
+ })
+ },
+ );
+ };
+
+ do_bench(128, 512);
+ do_bench(256, 1024);
+ do_bench(512, 2048);
+ do_bench(1024, 4096);
+ do_bench(2048, 8192);
+
+ group.finish();
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);
diff --git a/arrow/benches/primitive_run_take.rs
b/arrow/benches/primitive_run_take.rs
new file mode 100644
index 000000000..8c9a3fd04
--- /dev/null
+++ b/arrow/benches/primitive_run_take.rs
@@ -0,0 +1,79 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow::array::UInt32Builder;
+use arrow::compute::take;
+use arrow::datatypes::{Int32Type, Int64Type};
+use arrow::util::bench_util::*;
+use arrow::util::test_util::seedable_rng;
+use arrow_array::UInt32Array;
+use criterion::{criterion_group, criterion_main, Criterion};
+use rand::Rng;
+
+fn create_random_index(size: usize, null_density: f32) -> UInt32Array {
+ let mut rng = seedable_rng();
+ let mut builder = UInt32Builder::with_capacity(size);
+ for _ in 0..size {
+ if rng.gen::<f32>() < null_density {
+ builder.append_null();
+ } else {
+ let value = rng.gen_range::<u32, _>(0u32..size as u32);
+ builder.append_value(value);
+ }
+ }
+ builder.finish()
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+ let mut group = c.benchmark_group("primitive_run_take");
+
+ let mut do_bench = |physical_array_len: usize,
+ logical_array_len: usize,
+ take_len: usize| {
+ let run_array = create_primitive_run_array::<Int32Type, Int64Type>(
+ logical_array_len,
+ physical_array_len,
+ );
+ let indices = create_random_index(take_len, 0.0);
+ group.bench_function(
+ format!(
+ "(run_array_len:{logical_array_len},
physical_array_len:{physical_array_len}, take_len:{take_len})"),
+ |b| {
+ b.iter(|| {
+ criterion::black_box(take(&run_array, &indices,
None).unwrap());
+ })
+ },
+ );
+ };
+
+ do_bench(64, 512, 512);
+ do_bench(128, 512, 512);
+
+ do_bench(256, 1024, 512);
+ do_bench(256, 1024, 1024);
+
+ do_bench(512, 2048, 512);
+ do_bench(512, 2048, 1024);
+
+ do_bench(1024, 4096, 512);
+ do_bench(1024, 4096, 1024);
+
+ group.finish();
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);
diff --git a/arrow/benches/string_run_builder.rs
b/arrow/benches/string_run_builder.rs
index 2f0401bbe..dda0f35b8 100644
--- a/arrow/benches/string_run_builder.rs
+++ b/arrow/benches/string_run_builder.rs
@@ -17,26 +17,8 @@
use arrow::array::StringRunBuilder;
use arrow::datatypes::Int32Type;
+use arrow::util::bench_util::create_string_array_for_runs;
use criterion::{criterion_group, criterion_main, Criterion};
-use rand::{thread_rng, Rng};
-
-fn build_strings(
- physical_array_len: usize,
- logical_array_len: usize,
- string_len: usize,
-) -> Vec<String> {
- let mut rng = thread_rng();
- let run_len = logical_array_len / physical_array_len;
- let mut values: Vec<String> = (0..physical_array_len)
- .map(|_| (0..string_len).map(|_| rng.gen::<char>()).collect())
- .flat_map(|s| std::iter::repeat(s).take(run_len))
- .collect();
- while values.len() < logical_array_len {
- let last_val = values[values.len() - 1].clone();
- values.push(last_val);
- }
- values
-}
fn criterion_benchmark(c: &mut Criterion) {
let mut group = c.benchmark_group("string_run_builder");
@@ -50,7 +32,7 @@ fn criterion_benchmark(c: &mut Criterion) {
),
|b| {
let strings =
- build_strings(physical_array_len, logical_array_len,
string_len);
+ create_string_array_for_runs(physical_array_len,
logical_array_len, string_len);
b.iter(|| {
let mut builder =
StringRunBuilder::<Int32Type>::with_capacity(
physical_array_len,
diff --git a/arrow/benches/string_run_builder.rs
b/arrow/benches/string_run_iterator.rs
similarity index 60%
copy from arrow/benches/string_run_builder.rs
copy to arrow/benches/string_run_iterator.rs
index 2f0401bbe..cfa44e66e 100644
--- a/arrow/benches/string_run_builder.rs
+++ b/arrow/benches/string_run_iterator.rs
@@ -15,16 +15,16 @@
// specific language governing permissions and limitations
// under the License.
-use arrow::array::StringRunBuilder;
+use arrow::array::{Int32RunArray, StringArray, StringRunBuilder};
use arrow::datatypes::Int32Type;
use criterion::{criterion_group, criterion_main, Criterion};
use rand::{thread_rng, Rng};
-fn build_strings(
+fn build_strings_runs(
physical_array_len: usize,
logical_array_len: usize,
string_len: usize,
-) -> Vec<String> {
+) -> Int32RunArray {
let mut rng = thread_rng();
let run_len = logical_array_len / physical_array_len;
let mut values: Vec<String> = (0..physical_array_len)
@@ -35,43 +35,47 @@ fn build_strings(
let last_val = values[values.len() - 1].clone();
values.push(last_val);
}
- values
+ let mut builder = StringRunBuilder::<Int32Type>::with_capacity(
+ physical_array_len,
+ (string_len + 1) * physical_array_len,
+ );
+ builder.extend(values.into_iter().map(Some));
+
+ builder.finish()
}
fn criterion_benchmark(c: &mut Criterion) {
- let mut group = c.benchmark_group("string_run_builder");
+ let mut group = c.benchmark_group("string_run_iterator");
let mut do_bench = |physical_array_len: usize,
logical_array_len: usize,
string_len: usize| {
group.bench_function(
- format!(
- "(run_array_len:{logical_array_len},
physical_array_len:{physical_array_len}, string_len: {string_len})",
- ),
- |b| {
- let strings =
- build_strings(physical_array_len, logical_array_len,
string_len);
- b.iter(|| {
- let mut builder =
StringRunBuilder::<Int32Type>::with_capacity(
- physical_array_len,
- (string_len + 1) * physical_array_len,
- );
+ format!(
+ "(run_array_len:{logical_array_len},
physical_array_len:{physical_array_len}, string_len: {string_len})"),
+ |b| {
+ let run_array =
+ build_strings_runs(physical_array_len, logical_array_len,
string_len);
+ let typed = run_array.downcast::<StringArray>().unwrap();
+ b.iter(|| {
+ let iter = typed.into_iter();
+ for _ in iter {}
+ })
+ },
+ );
+ };
- for val in &strings {
- builder.append_value(val);
- }
+ do_bench(256, 1024, 5);
+ do_bench(256, 1024, 25);
+ do_bench(256, 1024, 100);
- builder.finish();
- })
- },
- );
- };
+ do_bench(512, 2048, 5);
+ do_bench(512, 2048, 25);
+ do_bench(512, 2048, 100);
- do_bench(20, 1000, 5);
- do_bench(100, 1000, 5);
- do_bench(100, 1000, 10);
- do_bench(100, 10000, 10);
- do_bench(100, 10000, 100);
+ do_bench(1024, 4096, 5);
+ do_bench(1024, 4096, 25);
+ do_bench(1024, 4096, 100);
group.finish();
}
diff --git a/arrow/benches/take_kernels.rs b/arrow/benches/take_kernels.rs
index c4677cc72..731426031 100644
--- a/arrow/benches/take_kernels.rs
+++ b/arrow/benches/take_kernels.rs
@@ -139,6 +139,13 @@ fn add_benchmark(c: &mut Criterion) {
c.bench_function("take str null values null indices 1024", |b| {
b.iter(|| bench_take(&values, &indices))
});
+
+ let values = create_primitive_run_array::<Int32Type, Int32Type>(1024, 512);
+ let indices = create_random_index(1024, 0.0);
+ c.bench_function(
+ "take primitive run logical len: 1024, physical len: 512, indices:
1024",
+ |b| b.iter(|| bench_take(&values, &indices)),
+ );
}
criterion_group!(benches, add_benchmark);
diff --git a/arrow/src/util/bench_util.rs b/arrow/src/util/bench_util.rs
index 6420b6346..33552dbe3 100644
--- a/arrow/src/util/bench_util.rs
+++ b/arrow/src/util/bench_util.rs
@@ -22,6 +22,7 @@ use crate::datatypes::*;
use crate::util::test_util::seedable_rng;
use arrow_buffer::Buffer;
use rand::distributions::uniform::SampleUniform;
+use rand::thread_rng;
use rand::Rng;
use rand::SeedableRng;
use rand::{
@@ -145,6 +146,73 @@ pub fn create_string_dict_array<K: ArrowDictionaryKeyType>(
data.iter().map(|x| x.as_deref()).collect()
}
+/// Create primitive run array for given logical and physical array lengths
+pub fn create_primitive_run_array<R: RunEndIndexType, V: ArrowPrimitiveType>(
+ logical_array_len: usize,
+ physical_array_len: usize,
+) -> RunArray<R> {
+ assert!(logical_array_len >= physical_array_len);
+ // typical length of each run
+ let run_len = logical_array_len / physical_array_len;
+
+ // Some runs should have extra length
+ let mut run_len_extra = logical_array_len % physical_array_len;
+
+ let mut values: Vec<V::Native> = (0..physical_array_len)
+ .flat_map(|s| {
+ let mut take_len = run_len;
+ if run_len_extra > 0 {
+ take_len += 1;
+ run_len_extra -= 1;
+ }
+ std::iter::repeat(V::Native::from_usize(s).unwrap()).take(take_len)
+ })
+ .collect();
+ while values.len() < logical_array_len {
+ let last_val = values[values.len() - 1];
+ values.push(last_val);
+ }
+ let mut builder = PrimitiveRunBuilder::<R,
V>::with_capacity(physical_array_len);
+ builder.extend(values.into_iter().map(Some));
+
+ builder.finish()
+}
+
+/// Create string array to be used by run array builder. The string array
+/// will result in run array with physial length of `physical_array_len`
+/// and logical length of `logical_array_len`
+pub fn create_string_array_for_runs(
+ physical_array_len: usize,
+ logical_array_len: usize,
+ string_len: usize,
+) -> Vec<String> {
+ assert!(logical_array_len >= physical_array_len);
+ let mut rng = thread_rng();
+
+ // typical length of each run
+ let run_len = logical_array_len / physical_array_len;
+
+ // Some runs should have extra length
+ let mut run_len_extra = logical_array_len % physical_array_len;
+
+ let mut values: Vec<String> = (0..physical_array_len)
+ .map(|_| (0..string_len).map(|_| rng.gen::<char>()).collect())
+ .flat_map(|s| {
+ let mut take_len = run_len;
+ if run_len_extra > 0 {
+ take_len += 1;
+ run_len_extra -= 1;
+ }
+ std::iter::repeat(s).take(take_len)
+ })
+ .collect();
+ while values.len() < logical_array_len {
+ let last_val = values[values.len() - 1].clone();
+ values.push(last_val);
+ }
+ values
+}
+
/// Creates an random (but fixed-seeded) binary array of a given size and null
density
pub fn create_binary_array<Offset: OffsetSizeTrait>(
size: usize,
