Help LLVM vectorize comparison kernel ~50-80% faster

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

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(