Skip to content

Commit

Permalink
Speed up checked kernels for non-null data (~1.4-5x faster) (#2749)
Browse files Browse the repository at this point in the history
* Speed up checked kernels

* Fast path for non-null

* Move some code
  • Loading branch information
Dandandan committed Sep 17, 2022
1 parent ca00b67 commit 46fcb0c
Showing 1 changed file with 27 additions and 4 deletions.
31 changes: 27 additions & 4 deletions arrow/src/compute/kernels/arity.rs
Expand Up @@ -106,15 +106,26 @@ where
let len = array.len();
let null_count = array.null_count();

let mut buffer = BufferBuilder::<O::Native>::new(len);
buffer.append_n_zeroed(array.len());
let slice = buffer.as_slice_mut();
if null_count == 0 {
let values = array.values().iter().map(|v| op(*v));
// JUSTIFICATION
// Benefit
// ~60% speedup
// Soundness
// `values` is an iterator with a known size because arrays are sized.
let buffer = unsafe { Buffer::try_from_trusted_len_iter(values)? };
return Ok(unsafe { build_primitive_array(len, buffer, 0, None) });
}

let null_buffer = array
.data_ref()
.null_buffer()
.map(|b| b.bit_slice(array.offset(), array.len()));

let mut buffer = BufferBuilder::<O::Native>::new(len);
buffer.append_n_zeroed(array.len());
let slice = buffer.as_slice_mut();

try_for_each_valid_idx(array.len(), 0, null_count, null_buffer.as_deref(), |idx| {
unsafe { *slice.get_unchecked_mut(idx) = op(array.value_unchecked(idx))? };
Ok::<_, ArrowError>(())
Expand Down Expand Up @@ -284,9 +295,21 @@ where
if a.is_empty() {
return Ok(PrimitiveArray::from(ArrayData::new_empty(&O::DATA_TYPE)));
}

let len = a.len();

if a.null_count() == 0 && b.null_count() == 0 {
let values = a.values().iter().zip(b.values()).map(|(l, r)| op(*l, *r));
let buffer = unsafe { Buffer::try_from_trusted_len_iter(values) }?;
// JUSTIFICATION
// Benefit
// ~75% speedup
// Soundness
// `values` is an iterator with a known size from a PrimitiveArray
return Ok(unsafe { build_primitive_array(len, buffer, 0, None) });
}

let null_buffer = combine_option_bitmap(&[a.data(), b.data()], len).unwrap();

let null_count = null_buffer
.as_ref()
.map(|x| len - x.count_set_bits())
Expand Down

0 comments on commit 46fcb0c

Please sign in to comment.