ARROW-12032: [Rust] Optimize comparison kernels

rust/arrow/benches/comparison_kernels.rs

@@ -121,8 +121,8 @@ fn bench_nlike_utf8_scalar(arr_a: &StringArray, value_b: &str) {
 
 fn add_benchmark(c: &mut Criterion) {
     let size = 65536;
-    let arr_a = create_primitive_array::<Float32Type>(size, 0.0);
-    let arr_b = create_primitive_array::<Float32Type>(size, 0.0);
+    let arr_a = create_primitive_array_with_seed::<Float32Type>(size, 0.0, 42);
+    let arr_b = create_primitive_array_with_seed::<Float32Type>(size, 0.0, 43);
 
     let arr_string = create_string_array(size, 0.0);
 

rust/arrow/src/buffer/mutable.rs

@@ -415,6 +415,61 @@ impl MutableBuffer {
         buffer
     }
 
+    /// Creates a [`MutableBuffer`] from a boolean [`Iterator`] with a trusted (upper) length.
+    /// # use arrow::buffer::MutableBuffer;
+    /// # Example
+    /// ```
+    /// # use arrow::buffer::MutableBuffer;
+    /// let v = vec![false, true, false];
+    /// let iter = v.iter().map(|x| *x || true);
+    /// let buffer = unsafe { MutableBuffer::from_trusted_len_iter_bool(iter) };
+    /// assert_eq!(buffer.len(), 1) // 3 booleans have 1 byte
+    /// ```
+    /// # Safety
+    /// This method assumes that the iterator's size is correct and is undefined behavior
+    /// to use it on an iterator that reports an incorrect length.
+    // This implementation is required for two reasons:
+    // 1. there is no trait `TrustedLen` in stable rust and therefore
+    //    we can't specialize `extend` for `TrustedLen` like `Vec` does.
+    // 2. `from_trusted_len_iter_bool` is faster.
+    pub unsafe fn from_trusted_len_iter_bool<I: Iterator<Item = bool>>(
+        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)
+        };
+
+        '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
+    }
+
     /// Creates a [`MutableBuffer`] from an [`Iterator`] with a trusted (upper) length or errors
     /// if any of the items of the iterator is an error.
     /// Prefer this to `collect` whenever possible, as it is faster ~60% faster.

rust/arrow/src/compute/kernels/comparison.rs

@@ -47,17 +47,72 @@ macro_rules! compare_op {
         let null_bit_buffer =
             combine_option_bitmap($left.data_ref(), $right.data_ref(), $left.len())?;
 
-        let buffer = (0..$left.len())
-            .map(|i| $op($left.value(i), $right.value(i)))
-            .collect();
+        let comparison = (0..$left.len()).map(|i| $op($left.value(i), $right.value(i)));
+        // same size as $left.len() and $right.len()
+        let buffer = unsafe { MutableBuffer::from_trusted_len_iter_bool(comparison) };
 
         let data = ArrayData::new(
             DataType::Boolean,
             $left.len(),
             None,
             null_bit_buffer,
             0,
-            vec![buffer],
+            vec![Buffer::from(buffer)],
+            vec![],
+        );
+        Ok(BooleanArray::from(Arc::new(data)))
+    }};
+}
+
+macro_rules! compare_op_primitive {
+    ($left: expr, $right:expr, $op:expr) => {{
+        if $left.len() != $right.len() {
+            return Err(ArrowError::ComputeError(
+                "Cannot perform comparison operation on arrays of different length"
+                    .to_string(),
+            ));
+        }
+
+        let null_bit_buffer =
+            combine_option_bitmap($left.data_ref(), $right.data_ref(), $left.len())?;
+
+        let mut values = MutableBuffer::from_len_zeroed(($left.len() + 7) / 8);
+        let lhs_chunks_iter = $left.values().chunks_exact(8);
+        let lhs_remainder = lhs_chunks_iter.remainder();
+        let rhs_chunks_iter = $right.values().chunks_exact(8);
+        let rhs_remainder = rhs_chunks_iter.remainder();
+        let chunks = $left.len() / 8;
+
+        values[..chunks]
+            .iter_mut()
+            .zip(lhs_chunks_iter)
+            .zip(rhs_chunks_iter)
+            .for_each(|((byte, lhs), rhs)| {
+                lhs.iter()
+                    .zip(rhs.iter())
+                    .enumerate()
+                    .for_each(|(i, (&lhs, &rhs))| {
+                        *byte |= if $op(lhs, rhs) { 1 << i } else { 0 };
+                    });
+            });
+
+        if !lhs_remainder.is_empty() {
+            let last = &mut values[chunks];
+            lhs_remainder
+                .iter()
+                .zip(rhs_remainder.iter())
+                .enumerate()
+                .for_each(|(i, (&lhs, &rhs))| {
+                    *last |= if $op(lhs, rhs) { 1 << i } else { 0 };
+                });
+        };
+        let data = ArrayData::new(
+            DataType::Boolean,
+            $left.len(),
+            None,
+            null_bit_buffer,
+            0,
+            vec![Buffer::from(values)],
             vec![],
         );
         Ok(BooleanArray::from(Arc::new(data)))
@@ -68,17 +123,54 @@ macro_rules! compare_op_scalar {
     ($left: expr, $right:expr, $op:expr) => {{
         let null_bit_buffer = $left.data().null_buffer().cloned();
 
-        let buffer = (0..$left.len())
-            .map(|i| $op($left.value(i), $right))
-            .collect();
+        let comparison = (0..$left.len()).map(|i| $op($left.value(i), $right));
+        // same as $left.len()
+        let buffer = unsafe { MutableBuffer::from_trusted_len_iter_bool(comparison) };
+
+        let data = ArrayData::new(
+            DataType::Boolean,
+            $left.len(),
+            None,
+            null_bit_buffer,
+            0,
+            vec![Buffer::from(buffer)],
+            vec![],
+        );
+        Ok(BooleanArray::from(Arc::new(data)))
+    }};
+}
+
+macro_rules! compare_op_scalar_primitive {
+    ($left: expr, $right:expr, $op:expr) => {{
+        let null_bit_buffer = $left.data().null_buffer().cloned();
+
+        let mut values = MutableBuffer::from_len_zeroed(($left.len() + 7) / 8);
+        let lhs_chunks_iter = $left.values().chunks_exact(8);
+        let lhs_remainder = lhs_chunks_iter.remainder();
+        let chunks = $left.len() / 8;
+
+        values[..chunks]
+            .iter_mut()
+            .zip(lhs_chunks_iter)
+            .for_each(|(byte, chunk)| {
+                chunk.iter().enumerate().for_each(|(i, &c_i)| {
+                    *byte |= if $op(c_i, $right) { 1 << i } else { 0 };
+                });
+            });
+        if !lhs_remainder.is_empty() {
+            let last = &mut values[chunks];
+            lhs_remainder.iter().enumerate().for_each(|(i, &lhs)| {
+                *last |= if $op(lhs, $right) { 1 << i } else { 0 };
+            });
+        };
 
         let data = ArrayData::new(
             DataType::Boolean,
             $left.len(),
             None,
             null_bit_buffer,
             0,
-            vec![buffer],
+            vec![Buffer::from(values)],
             vec![],
         );
         Ok(BooleanArray::from(Arc::new(data)))
@@ -96,7 +188,7 @@ where
     T: ArrowNumericType,
     F: Fn(T::Native, T::Native) -> bool,
 {
-    compare_op!(left, right, op)
+    compare_op_primitive!(left, right, op)
 }
 
 /// Evaluate `op(left, right)` for [`PrimitiveArray`] and scalar using
@@ -110,7 +202,7 @@ where
     T: ArrowNumericType,
     F: Fn(T::Native, T::Native) -> bool,
 {
-    compare_op_scalar!(left, right, op)
+    compare_op_scalar_primitive!(left, right, op)
 }
 
 /// Perform SQL `left LIKE right` operation on [`StringArray`] / [`LargeStringArray`].

rust/arrow/src/util/bench_util.rs

@@ -17,12 +17,15 @@
 
 //! Utils to make benchmarking easier
 
-use rand::distributions::{Alphanumeric, Distribution, Standard};
-use rand::Rng;
-
 use crate::array::*;
 use crate::datatypes::*;
 use crate::util::test_util::seedable_rng;
+use rand::Rng;
+use rand::SeedableRng;
+use rand::{
+    distributions::{Alphanumeric, Distribution, Standard},
+    prelude::StdRng,
+};
 
 /// Creates an random (but fixed-seeded) array of a given size and null density
 pub fn create_primitive_array<T>(size: usize, null_density: f32) -> PrimitiveArray<T>
@@ -43,6 +46,28 @@ where
         .collect()
 }
 
+pub fn create_primitive_array_with_seed<T>(
+    size: usize,
+    null_density: f32,
+    seed: u64,
+) -> PrimitiveArray<T>
+where
+    T: ArrowPrimitiveType,
+    Standard: Distribution<T::Native>,
+{
+    let mut rng = StdRng::seed_from_u64(seed);
+
+    (0..size)
+        .map(|_| {
+            if rng.gen::<f32>() < null_density {
+                None
+            } else {
+                Some(rng.gen())
+            }
+        })
+        .collect()
+}
+
 /// Creates an random (but fixed-seeded) array of a given size and null density
 pub fn create_boolean_array(
     size: usize,