Faster bitmask iteration (#1228)

* Add UnalignedBitChunks (#1227) * Clippy * Fix flaky test * Improve test legibility * Fix SlicesIterator offset direction * Format * Fix byte-aligned termination * Test edge-cases * More tests Co-authored-by: Andrew Lamb <[email protected]> * Review feedback * Make UnalignedBitChunkIterator crate local Co-authored-by: Andrew Lamb <[email protected]>
apache · Feb 2, 2022 · f055fb0 · f055fb0
1 parent 91d12ec
commit f055fb0
Show file tree

Hide file tree

Showing 4 changed files with 547 additions and 166 deletions.
diff --git a/arrow/src/buffer/immutable.rs b/arrow/src/buffer/immutable.rs
@@ -21,7 +21,7 @@ use std::ptr::NonNull;
 use std::sync::Arc;
 use std::{convert::AsRef, usize};
 
-use crate::util::bit_chunk_iterator::BitChunks;
+use crate::util::bit_chunk_iterator::{BitChunks, UnalignedBitChunk};
 use crate::{
     bytes::{Bytes, Deallocation},
     datatypes::ArrowNativeType,
@@ -205,11 +205,7 @@ impl Buffer {
     /// Returns the number of 1-bits in this buffer, starting from `offset` with `length` bits
     /// inspected. Note that both `offset` and `length` are measured in bits.
     pub fn count_set_bits_offset(&self, offset: usize, len: usize) -> usize {
-        let chunks = self.bit_chunks(offset, len);
-        let mut count = chunks.iter().map(|c| c.count_ones() as usize).sum();
-        count += chunks.remainder_bits().count_ones() as usize;
-
-        count
+        UnalignedBitChunk::new(self.as_slice(), offset, len).count_ones()
     }
 }