Use BitChunks in equal_bits (apache#2186)

tustvold · Jul 27, 2022 · 0e0c6e0 · 0e0c6e0
1 parent e096ec7
commit 0e0c6e0
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 8 deletions.
diff --git a/arrow/benches/equal.rs b/arrow/benches/equal.rs
@@ -51,6 +51,9 @@ fn add_benchmark(c: &mut Criterion) {
 
     let arr_a = create_boolean_array(513, 0.0, 0.5);
     c.bench_function("equal_bool_513", |b| b.iter(|| bench_equal(&arr_a)));
+
+    let arr_a = create_boolean_array(1024, 0.0, 0.5);
+    c.bench_function("equal_bool_1024", |b| b.iter(|| bench_equal(&arr_a)));
 }
 
 criterion_group!(benches, add_benchmark);

diff --git a/arrow/src/array/equal/boolean.rs b/arrow/src/array/equal/boolean.rs
@@ -85,10 +85,10 @@ pub(super) fn boolean_equal(
             let rhs_pos = rhs_start + i;
             let lhs_is_null = !get_bit(lhs_null_bytes, lhs_pos);
             let rhs_is_null = !get_bit(rhs_null_bytes, rhs_pos);
+            let lhs_is_true = get_bit(lhs_values, lhs_pos);
+            let rhs_is_true = get_bit(rhs_values, rhs_pos);
 
-            lhs_is_null
-                || (lhs_is_null == rhs_is_null)
-                    && equal_bits(lhs_values, rhs_values, lhs_pos, rhs_pos, 1)
+            lhs_is_null == rhs_is_null && (lhs_is_null || (lhs_is_true == rhs_is_true))
         })
     }
 }

diff --git a/arrow/src/array/equal/utils.rs b/arrow/src/array/equal/utils.rs
@@ -17,7 +17,7 @@
 
 use crate::array::{data::count_nulls, ArrayData};
 use crate::datatypes::DataType;
-use crate::util::bit_util;
+use crate::util::bit_chunk_iterator::BitChunks;
 
 // whether bits along the positions are equal
 // `lhs_start`, `rhs_start` and `len` are _measured in bits_.
@@ -29,10 +29,16 @@ pub(super) fn equal_bits(
     rhs_start: usize,
     len: usize,
 ) -> bool {
-    (0..len).all(|i| {
-        bit_util::get_bit(lhs_values, lhs_start + i)
-            == bit_util::get_bit(rhs_values, rhs_start + i)
-    })
+    let lhs = BitChunks::new(lhs_values, lhs_start, len);
+    let rhs = BitChunks::new(rhs_values, rhs_start, len);
+
+    for (a, b) in lhs.iter().zip(rhs.iter()) {
+        if a != b {
+            return false;
+        }
+    }
+
+    lhs.remainder_bits() == rhs.remainder_bits()
 }
 
 #[inline]