Add try_unary, binary, try_binary kernels ~90% faster #2666
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
tustvold
force-pushed
the
add-try-unary
branch
from
0815487
to
882a655
Compare
tustvold
commented
Sep 6, 2022
arrow/src/compute/kernels/arity.rs
Outdated
| use std::sync::Arc; | ||
|
|
||
| fn try_for_each_valid<F: FnMut(usize) -> Result<()>>( |
There was a problem hiding this comment.
This internal iteration style optimises much better than the external style when there are multiple variants, I may create a HybridBitIndexIterator that special cases for_each and try_fold 🤔
liukun4515
reviewed
Sep 6, 2022
| @@ -198,6 +259,80 @@ where | |||
| } | |||
| } | |||
|
|
|||
| pub fn binary<A, B, F, O>( | |||
liukun4515
reviewed
Sep 6, 2022
| #[inline] | ||
| fn into_primitive_array_data<I: ArrowPrimitiveType, O: ArrowPrimitiveType>( | ||
| array: &PrimitiveArray<I>, | ||
| unsafe fn build_primitive_array<O: ArrowPrimitiveType>( |
There was a problem hiding this comment.
the new name makes more sense than before
liukun4515
reviewed
Sep 6, 2022
| } | ||
| try_binary(left, right, |a, b| { | ||
| op(a, b).ok_or_else(|| { | ||
| ArrowError::ComputeError(format!("Overflow happened on: {:?}, {:?}", a, b)) |
There was a problem hiding this comment.
Can we point out the op which cause the overflow?
viirya
reviewed
Sep 6, 2022
| @@ -24,8 +24,51 @@ use super::{ | |||
| PrimitiveArray, | |||
| }; | |||
|
|
|||
| /// an iterator that returns Some(T) or None, that can be used on any [`ArrayAccessor`] | |||
| // Note: This implementation is based on std's [Vec]s' [IntoIter]. | |||
| /// An iterator that returns Some(T) or None, that can be used on any [`ArrayAccessor`] | |||
viirya
reviewed
Sep 6, 2022
| let null_count = array.null_count(); | ||
|
|
||
| let mut buffer = BufferBuilder::<O::Native>::new(len); | ||
| buffer.append_n_zeroed(array.len()); |
There was a problem hiding this comment.
Hmm, do we need zero-initialized values for the buffer?
There was a problem hiding this comment.
It is UB if we don't initialize all values in the buffer, even the null slots. We must therefore zero out the nulls, in the past I have found it is faster to zero initialize everything, and override the valid indexes, than to interleave appending nulls and values.
viirya
reviewed
Sep 6, 2022
arrow/src/compute/kernels/arity.rs
Outdated
Comment on lines
43
to
48
| if selectivity > 0.8 { | ||
| BitSliceIterator::new(nulls.unwrap(), 0, len) | ||
| .flat_map(|(start, end)| start..end) | ||
| .try_for_each(f) | ||
| } else { | ||
| BitIndexIterator::new(nulls.unwrap(), 0, len).try_for_each(f) |
There was a problem hiding this comment.
Am I read this incorrectly? I think higher selectivity here is more null values in the array? So I guess BitIndexIterator is more performant?
There was a problem hiding this comment.
Oops, i think you're right. I copied this from the filter kernel and forgot to reverse this
HaoYang670
reviewed
Sep 7, 2022
| @@ -198,6 +259,80 @@ where | |||
| } | |||
| } | |||
|
|
|||
| pub fn binary<A, B, F, O>( | |||
There was a problem hiding this comment.
I guess we need to add docs for binary and try_binary.
| return Ok(PrimitiveArray::from(ArrayData::new_empty(&O::DATA_TYPE))); | ||
| } | ||
|
|
||
| let null_buffer = combine_option_bitmap(&[a.data(), b.data()], len).unwrap(); |
There was a problem hiding this comment.
Suggested change
| let null_buffer = combine_option_bitmap(&[a.data(), b.data()], len).unwrap(); | |
| let null_buffer = combine_option_bitmap(&[a.data(), b.data()], len)?; |
| op: F, | ||
| ) -> Result<PrimitiveArray<T>> | ||
| ) -> Result<PrimitiveArray<LT>> |
There was a problem hiding this comment.
Why is the return type same as the left type?
There was a problem hiding this comment.
In the style, the type of left and right can be different.
Do we need to support this？
There was a problem hiding this comment.
I believe it is required to allow adding a duration to a time
tustvold
force-pushed
the
add-try-unary
branch
from
90e12b1
to
fb9ae62
Compare
Final benchmarksSo anything from 50% to 99% faster for the checked kernels 🎉 . The unchecked kernels regularly vary by +-10% but we're talking single-digit microseconds, so this isn't all that surprising. |
tustvold
force-pushed
the
add-try-unary
branch
from
fdd3040
to
f05a3e6
Compare
viirya
approved these changes
Sep 11, 2022
|
Benchmark runs are scheduled for baseline = d88ed6a and contender = 2d28010. 2d28010 is a master commit associated with this PR. Results will be available as each benchmark for each run completes. |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
Closes #.
Rationale for this change
ArrayIter is great from an ergonomics perspective and definitely something we should continue to provide to users, however, the way it interleaves null masks severely makes for pretty poor performance. For primitive arrays, this can result in orders of magnitude more time spent handling the null mask than the values themselves.
What changes are included in this PR?
Are there any user-facing changes?
FYI @viirya @liukun4515