ARROW-12032: [Rust] Optimize comparison kernels

This adds a function `from_trusted_len_iter_bool` to speed up the creation of an array for booleans.

Benchmarks are a bit noisy, but seems to be ~10-20% faster for comparison kernels. This also has some positive effect on DataFusion queries, as they contain quite some (nested) comparisons in filters. For example, executing tpch query 6 in memory is ~7% faster.

```
Gnuplot not found, using plotters backend
eq Float32              time:   [54.204 us 54.284 us 54.364 us]
                        change: [-29.087% -28.838% -28.581%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 6 outliers among 100 measurements (6.00%)
  5 (5.00%) low mild
  1 (1.00%) high mild

eq scalar Float32       time:   [43.660 us 43.743 us 43.830 us]
                        change: [-30.819% -30.545% -30.269%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 5 outliers among 100 measurements (5.00%)
  4 (4.00%) high mild
  1 (1.00%) high severe

neq Float32             time:   [68.726 us 68.893 us 69.048 us]
                        change: [-14.045% -13.772% -13.490%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 1 outliers among 100 measurements (1.00%)
  1 (1.00%) high mild

neq scalar Float32      time:   [46.251 us 46.322 us 46.395 us]
                        change: [-12.204% -11.952% -11.702%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 6 outliers among 100 measurements (6.00%)
  1 (1.00%) low mild
  5 (5.00%) high mild

lt Float32              time:   [50.264 us 50.438 us 50.613 us]
                        change: [-21.300% -20.964% -20.649%] (p = 0.00 < 0.05)
                        Performance has improved.

lt scalar Float32       time:   [48.847 us 48.929 us 49.013 us]
                        change: [-10.132% -9.9180% -9.6910%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 5 outliers among 100 measurements (5.00%)
  4 (4.00%) high mild
  1 (1.00%) high severe

lt_eq Float32           time:   [46.105 us 46.198 us 46.282 us]
                        change: [-21.276% -20.966% -20.703%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 18 outliers among 100 measurements (18.00%)
  2 (2.00%) low severe
  13 (13.00%) low mild
  1 (1.00%) high mild
  2 (2.00%) high severe

lt_eq scalar Float32    time:   [47.359 us 47.456 us 47.593 us]
                        change: [+0.2766% +0.5240% +0.7821%] (p = 0.00 < 0.05)
                        Change within noise threshold.
Found 10 outliers among 100 measurements (10.00%)
  8 (8.00%) high mild
  2 (2.00%) high severe

gt Float32              time:   [57.313 us 57.363 us 57.412 us]
                        change: [-18.328% -18.177% -18.031%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 3 outliers among 100 measurements (3.00%)
  2 (2.00%) low severe
  1 (1.00%) low mild

gt scalar Float32       time:   [44.091 us 44.132 us 44.175 us]
                        change: [-9.4233% -9.2747% -9.1273%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 7 outliers among 100 measurements (7.00%)
  4 (4.00%) low mild
  3 (3.00%) high mild

gt_eq Float32           time:   [55.856 us 55.932 us 56.007 us]
                        change: [-7.4997% -7.2656% -7.0334%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 3 outliers among 100 measurements (3.00%)
  1 (1.00%) low mild
  2 (2.00%) high mild

gt_eq scalar Float32    time:   [42.365 us 42.419 us 42.482 us]
                        change: [+0.5289% +0.7174% +0.9116%] (p = 0.00 < 0.05)
                        Change within noise threshold.
Found 5 outliers among 100 measurements (5.00%)
  2 (2.00%) high mild
  3 (3.00%) high severe
```

Closes #9759 from Dandandan/optimize_comparison

Authored-by: Heres, Daniel <danielheres@gmail.com>
Signed-off-by: Andrew Lamb <andrew@nerdnetworks.org>

Author: Dandandan

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

@@ -46,17 +46,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(data))
@@ -67,17 +122,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(data))
@@ -95,7 +187,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
@@ -109,7 +201,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,