Use set_bits from transform::util

Author: tustvold

arrow/Cargo.toml

@@ -104,6 +104,10 @@ harness = false
 name = "boolean_kernels"
 harness = false
 
+[[bench]]
+name = "boolean_append_packed"
+harness = false
+
 [[bench]]
 name = "arithmetic_kernels"
 harness = false

arrow/benches/boolean_append_packed.rs

@@ -0,0 +1,38 @@
+use arrow::array::BooleanBufferBuilder;
+use criterion::{criterion_group, criterion_main, Criterion};
+use rand::{thread_rng, Rng};
+
+fn rand_bytes(len: usize) -> Vec<u8> {
+    let mut rng = thread_rng();
+    let mut buf = vec![0_u8; len];
+    rng.fill(buf.as_mut_slice());
+    buf
+}
+
+fn boolean_append_packed(c: &mut Criterion) {
+    let mut rng = thread_rng();
+    let source = rand_bytes(1024);
+    let ranges: Vec<_> = (0..100)
+        .into_iter()
+        .map(|_| {
+            let start: usize = rng.gen_range(0..1024 * 8);
+            let end: usize = rng.gen_range(start..1024 * 8);
+            start..end
+        })
+        .collect();
+
+    let total_bits: usize = ranges.iter().map(|x| x.end - x.start).sum();
+
+    c.bench_function("boolean_append_packed", |b| {
+        b.iter(|| {
+            let mut buffer = BooleanBufferBuilder::new(total_bits);
+            for range in &ranges {
+                buffer.append_packed_range(range.clone(), &source);
+            }
+            assert_eq!(buffer.len(), total_bits);
+        })
+    });
+}
+
+criterion_group!(benches, boolean_append_packed);
+criterion_main!(benches);

arrow/src/array/builder.rs

@@ -399,60 +399,6 @@ impl BooleanBufferBuilder {
         }
     }
 
-    /// Append `count` bits from `to_set`
-    ///
-    /// `to_set` is a slice of bits packed LSB-first into `[u8]`
-    ///
-    /// # Panics
-    ///
-    /// Panics if `to_set` does not contain `ceil(count / 8)` bytes
-    #[inline]
-    pub fn append_packed(&mut self, count: usize, to_set: &[u8]) {
-        assert_eq!((count + 7) >> 3, to_set.len());
-
-        let new_len = self.len + count;
-        let new_buf_len = (new_len + 7) >> 3;
-        self.buffer.reserve(new_buf_len - self.buffer.len());
-
-        let whole_bytes = count >> 3;
-        let overrun = count & 7;
-
-        let skew = self.len & 7;
-        if skew == 0 {
-            self.buffer.extend_from_slice(&to_set[..whole_bytes]);
-            if overrun > 0 {
-                let masked = to_set[whole_bytes] & ((1 << overrun) - 1);
-                self.buffer.push(masked)
-            }
-
-            self.len = new_len;
-            debug_assert_eq!(self.buffer.len(), new_buf_len);
-            return;
-        }
-
-        for to_set_byte in &to_set[..whole_bytes] {
-            let low = *to_set_byte << skew;
-            let high = *to_set_byte >> (8 - skew);
-
-            *self.buffer.last_mut().unwrap() |= low;
-            self.buffer.push(high);
-        }
-
-        if overrun > 0 {
-            let masked = to_set[whole_bytes] & ((1 << overrun) - 1);
-            let low = masked << skew;
-            *self.buffer.last_mut().unwrap() |= low;
-
-            if overrun > 8 - skew {
-                let high = masked >> (8 - skew);
-                self.buffer.push(high)
-            }
-        }
-
-        self.len = new_len;
-        debug_assert_eq!(self.buffer.len(), new_buf_len);
-    }
-
     /// Append `range` bits from `to_set`
     ///
     /// `to_set` is a slice of bits packed LSB-first into `[u8]`
@@ -461,31 +407,16 @@ impl BooleanBufferBuilder {
     ///
     /// Panics if `to_set` does not contain `ceil(range.end / 8)` bytes
     pub fn append_packed_range(&mut self, range: Range<usize>, to_set: &[u8]) {
-        let count = range.end - range.start;
-        if count == 0 {
-            return;
-        }
-
-        let start_byte = range.start >> 3;
-        let end_byte = (range.end + 7) >> 3;
-        let skew = range.start & 7;
-
-        // `append_packed` requires the provided `to_set` to be byte aligned, therefore
-        // if the range being copied is not byte aligned we must first append
-        // the leading bits to reach a byte boundary
-        if skew == 0 {
-            // No skew can simply append bytes directly
-            self.append_packed(count, &to_set[start_byte..end_byte])
-        } else if start_byte + 1 == end_byte {
-            // Append bits from single byte
-            self.append_packed(count, &[to_set[start_byte] >> skew])
-        } else {
-            // Append trailing bits from first byte to reach byte boundary, then append
-            // bits from the remaining byte-aligned mask
-            let offset = 8 - skew;
-            self.append_packed(offset, &[to_set[start_byte] >> skew]);
-            self.append_packed(count - offset, &to_set[(start_byte + 1)..end_byte]);
-        }
+        let offset_write = self.len;
+        let len = range.end - range.start;
+        self.advance(len);
+        crate::util::bit_mask::set_bits(
+            self.buffer.as_slice_mut(),
+            to_set,
+            offset_write,
+            range.start,
+            len,
+        );
     }
 
     #[inline]
@@ -2925,37 +2856,6 @@ mod tests {
         assert!(buffer.get_bit(11));
     }
 
-    #[test]
-    fn test_bit_append_packed() {
-        let mut buffer = BooleanBufferBuilder::new(0);
-
-        buffer.append_packed(8, &[0b11111111]);
-        assert_eq!(buffer.buffer.as_slice(), &[0b11111111]);
-        assert_eq!(buffer.len(), 8);
-
-        buffer.append_packed(3, &[0b01010010]);
-        assert_eq!(buffer.buffer.as_slice(), &[0b11111111, 0b00000010]);
-        assert_eq!(buffer.len(), 11);
-
-        buffer.append_packed(5, &[0b00010100]);
-        assert_eq!(buffer.buffer.as_slice(), &[0b11111111, 0b10100010]);
-        assert_eq!(buffer.len(), 16);
-
-        buffer.append_packed(2, &[0b11110010]);
-        assert_eq!(
-            buffer.buffer.as_slice(),
-            &[0b11111111, 0b10100010, 0b00000010]
-        );
-        assert_eq!(buffer.len(), 18);
-
-        buffer.append_packed(15, &[0b11011010, 0b01010101]);
-        assert_eq!(
-            buffer.buffer.as_slice(),
-            &[0b11111111, 0b10100010, 0b01101010, 0b01010111, 0b00000001]
-        );
-        assert_eq!(buffer.len(), 33);
-    }
-
     #[test]
     fn test_bool_buffer_fuzz() {
         use rand::prelude::*;
@@ -2964,19 +2864,6 @@ mod tests {
         let mut all_bools = vec![];
         let mut rng = rand::thread_rng();
 
-        for _ in 0..100 {
-            let mask_length = (rng.next_u32() % 50) as usize;
-            let bools: Vec<_> = std::iter::from_fn(|| Some(rng.next_u32() & 1 == 0))
-                .take(mask_length)
-                .collect();
-
-            let mut compacted = BooleanBufferBuilder::new(mask_length);
-            compacted.append_slice(&bools);
-
-            buffer.append_packed(mask_length, compacted.buffer.as_slice());
-            all_bools.extend_from_slice(&bools);
-        }
-
         let src_len = 32;
         let (src, compacted_src) = {
             let src: Vec<_> = std::iter::from_fn(|| Some(rng.next_u32() & 1 == 0))

arrow/src/array/transform/boolean.rs

@@ -16,11 +16,8 @@
 // under the License.
 
 use crate::array::ArrayData;
-
-use super::{
-    Extend, _MutableArrayData,
-    utils::{resize_for_bits, set_bits},
-};
+use crate::util::bit_mask::set_bits;
+use super::{Extend, _MutableArrayData, utils::resize_for_bits};
 
 pub(super) fn build_extend(array: &ArrayData) -> Extend {
     let values = array.buffers()[0].as_slice();

arrow/src/array/transform/mod.rs

@@ -97,7 +97,7 @@ fn build_extend_null_bits(array: &ArrayData, use_nulls: bool) -> ExtendNullBits
         let bytes = bitmap.bits.as_slice();
         Box::new(move |mutable, start, len| {
             utils::resize_for_bits(&mut mutable.null_buffer, mutable.len + len);
-            mutable.null_count += utils::set_bits(
+            mutable.null_count += crate::util::bit_mask::set_bits(
                 mutable.null_buffer.as_slice_mut(),
                 bytes,
                 mutable.len,