This is an automated email from the ASF dual-hosted git repository.
dheres 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 b46fc9287 Help LLVM vectorize comparison kernel ~50-80% faster (#2646)
b46fc9287 is described below
commit b46fc9287da22534f0599651597fc13cb675ddff
Author: Raphael Taylor-Davies <[email protected]>
AuthorDate: Sun Sep 4 18:55:49 2022 +0100
Help LLVM vectorize comparison kernel ~50-80% faster (#2646)
* Help LLVM vectorize comparison kernel
* Add MutableBuffer::collect_bool
* Add SAFETY comments
---
arrow/src/buffer/mutable.rs | 68 ++++++++++++++++++---------------
arrow/src/compute/kernels/comparison.rs | 19 ++++-----
2 files changed, 45 insertions(+), 42 deletions(-)
diff --git a/arrow/src/buffer/mutable.rs b/arrow/src/buffer/mutable.rs
index 1c662ec23..96c837922 100644
--- a/arrow/src/buffer/mutable.rs
+++ b/arrow/src/buffer/mutable.rs
@@ -373,6 +373,41 @@ impl MutableBuffer {
assert!(len <= self.capacity());
self.len = len;
}
+
+ /// Invokes `f` with values `0..len` collecting the boolean results into a
new `MutableBuffer`
+ ///
+ /// This is similar to `from_trusted_len_iter_bool`, however, can be
significantly faster
+ /// as it eliminates the conditional `Iterator::next`
+ #[inline]
+ pub(crate) fn collect_bool<F: FnMut(usize) -> bool>(len: usize, mut f: F)
-> Self {
+ let mut buffer = Self::new(bit_util::ceil(len, 8));
+
+ let chunks = len / 8;
+ let remainder = len % 8;
+ for chunk in 0..chunks {
+ let mut packed = 0;
+ for bit_idx in 0..8 {
+ let i = bit_idx + chunk * 8;
+ packed |= (f(i) as u8) << bit_idx;
+ }
+
+ // SAFETY: Already allocated sufficient capacity
+ unsafe { buffer.push_unchecked(packed) }
+ }
+
+ if remainder != 0 {
+ let mut packed = 0;
+ for bit_idx in 0..remainder {
+ let i = bit_idx + chunks * 8;
+ packed |= (f(i) as u8) << bit_idx;
+ }
+
+ // SAFETY: Already allocated sufficient capacity
+ unsafe { buffer.push_unchecked(packed) }
+ }
+
+ buffer
+ }
}
/// # Safety
@@ -496,38 +531,9 @@ impl MutableBuffer {
mut iterator: I,
) -> Self {
let (_, upper) = iterator.size_hint();
- let upper = upper.expect("from_trusted_len_iter requires an upper
limit");
-
- let mut result = {
- let byte_capacity: usize = upper.saturating_add(7) / 8;
- MutableBuffer::new(byte_capacity)
- };
+ let len = upper.expect("from_trusted_len_iter requires an upper
limit");
- 'a: loop {
- let mut byte_accum: u8 = 0;
- let mut mask: u8 = 1;
-
- //collect (up to) 8 bits into a byte
- while mask != 0 {
- if let Some(value) = iterator.next() {
- byte_accum |= match value {
- true => mask,
- false => 0,
- };
- mask <<= 1;
- } else {
- if mask != 1 {
- // Add last byte
- result.push_unchecked(byte_accum);
- }
- break 'a;
- }
- }
-
- // Soundness: from_trusted_len
- result.push_unchecked(byte_accum);
- }
- result
+ Self::collect_bool(len, |_| iterator.next().unwrap())
}
/// Creates a [`MutableBuffer`] from an [`Iterator`] with a trusted
(upper) length or errors
diff --git a/arrow/src/compute/kernels/comparison.rs
b/arrow/src/compute/kernels/comparison.rs
index dd9d4fc5d..cba2d6e7d 100644
--- a/arrow/src/compute/kernels/comparison.rs
+++ b/arrow/src/compute/kernels/comparison.rs
@@ -59,12 +59,10 @@ where
let null_bit_buffer =
combine_option_bitmap(&[left.data_ref(), right.data_ref()],
left.len())?;
- // Safety:
- // `i < $left.len()` and $left.len() == $right.len()
- let comparison = (0..left.len())
- .map(|i| unsafe { op(left.value_unchecked(i),
right.value_unchecked(i)) });
- // same size as $left.len() and $right.len()
- let buffer = unsafe {
MutableBuffer::from_trusted_len_iter_bool(comparison) };
+ let buffer = MutableBuffer::collect_bool(left.len(), |i| unsafe {
+ // SAFETY: i in range 0..len
+ op(left.value_unchecked(i), right.value_unchecked(i))
+ });
let data = unsafe {
ArrayData::new_unchecked(
@@ -91,11 +89,10 @@ where
.null_buffer()
.map(|b| b.bit_slice(left.offset(), left.len()));
- // Safety:
- // `i < $left.len()`
- let comparison = (0..left.len()).map(|i| unsafe {
op(left.value_unchecked(i)) });
- // same as $left.len()
- let buffer = unsafe {
MutableBuffer::from_trusted_len_iter_bool(comparison) };
+ let buffer = MutableBuffer::collect_bool(left.len(), |i| unsafe {
+ // SAFETY: i in range 0..len
+ op(left.value_unchecked(i))
+ });
let data = unsafe {
ArrayData::new_unchecked(
