Optimise Bitmask iteration speed

src/kernels/bitmask/dispatch.rs

@@ -353,3 +353,29 @@ pub fn all_false_mask(mask: &Bitmask) -> bool {
         crate::kernels::bitmask::std::all_false_mask(mask)
     }
 }
+
+// --- Bit Position Iteration ---
+
+/// Iterates window-relative indices of a bitmask window where the bit
+/// equals `bit_value`, in ascending order.
+///
+/// Commonly used to enumerate the valid or null positions of a nullable
+/// array without materialising a full index vector.
+///
+/// # Parameters
+/// - `m`: Bitmask window as `(mask, offset, length)` tuple
+/// - `bit_value`: The bit value to search for, true for set bits and false for cleared bits
+///
+/// # Returns
+/// An iterator of window-relative bit indices matching `bit_value`.
+#[inline(always)]
+pub fn iter_window_bits(m: BitmaskVT<'_>, bit_value: bool) -> impl Iterator<Item = usize> + '_ {
+    #[cfg(feature = "simd")]
+    {
+        crate::kernels::bitmask::simd::iter_window_bits_simd::<W8>(m, bit_value)
+    }
+    #[cfg(not(feature = "simd"))]
+    {
+        crate::kernels::bitmask::std::iter_window_bits(m, bit_value)
+    }
+}

src/kernels/bitmask/mod.rs

@@ -58,6 +58,9 @@
 //! - **`all_true_mask`**: Test if all bits in bitmask are set to 1
 //! - **`all_false_mask`**: Test if all bits in bitmask are set to 0
 //!
+//! ### **Bit Position Iteration**
+//! - **`iter_window_bits`**: Enumerate window-relative indices of set or cleared bits
+//!
 //! ## Arrow Compatibility
 //!
 //! All operations maintain full compatibility with Apache Arrow's bitmask format:
@@ -136,6 +139,27 @@ pub fn bitmask_window_bytes_mut(mask: &mut Bitmask, offset: usize, len: usize) -
     &mut mask.bits[start..end]
 }
 
+/// Assemble one 64-bit word of the mask starting at `bit_start`, LSB
+/// first, reading past the buffer's end as zeros. An unaligned start
+/// combines the two overlapping loads with a shift, so the caller walks
+/// any bit window in whole words.
+#[inline(always)]
+pub fn load_word(mask: &Bitmask, bit_start: usize) -> u64 {
+    let byte_start = bit_start / 8;
+    let shift = bit_start % 8;
+    let mut buf = [0u8; 9];
+    let end = (byte_start + 9).min(mask.bits.len());
+    if byte_start < end {
+        buf[..end - byte_start].copy_from_slice(&mask.bits[byte_start..end]);
+    }
+    let lo = u64::from_le_bytes(buf[0..8].try_into().unwrap());
+    if shift == 0 {
+        lo
+    } else {
+        (lo >> shift) | ((buf[8] as u64) << (64 - shift))
+    }
+}
+
 /// Zero all slack bits ≥ `bm.len()`.
 #[inline(always)]
 pub fn clear_trailing_bits(bm: &mut Bitmask) {

src/kernels/bitmask/simd.rs

@@ -62,7 +62,7 @@ use crate::kernels::arithmetic::simd::{W8, W16, W32, W64};
 use crate::{Bitmask, BitmaskVT};
 
 use crate::enums::operators::{LogicalOperator, UnaryOperator};
-use crate::kernels::bitmask::{bitmask_window_bytes, bitmask_window_bytes_mut};
+use crate::kernels::bitmask::{bitmask_window_bytes, bitmask_window_bytes_mut, load_word};
 
 /// Primitive bit ops
 
@@ -928,6 +928,84 @@ where
     true
 }
 
+/// Scans the bits of the window `m` in batches of `LANES` 64-bit words,
+/// returning window-relative indices where the bit equals `bit_value`.
+///
+/// Each batch is compared against zero with a single SIMD operation, so a
+/// batch with no matches is rejected in one comparison rather than `LANES`
+/// separate word checks. A batch flagged as containing a match falls back
+/// to a scalar trailing-zeros scan restricted to its flagged words, so the
+/// exact positions are worked out only where they exist. The words left
+/// over at the end of the window, too few to fill a batch, scan
+/// individually.
+#[inline]
+pub fn iter_window_bits_simd<const LANES: usize>(
+    m: BitmaskVT<'_>,
+    bit_value: bool,
+) -> impl Iterator<Item = usize> + '_
+where
+{
+    use std::simd::cmp::SimdPartialEq;
+
+    let (mask, offset, len) = m;
+    let n_words = len.div_ceil(64);
+    let n_batches = n_words / LANES;
+
+    let word_at = move |word_index: usize| {
+        let bit_base = word_index * 64;
+        let take = (len - bit_base).min(64);
+        let mut w = load_word(mask, offset + bit_base);
+        if !bit_value {
+            w = !w;
+        }
+        if take < 64 {
+            w &= u64::MAX >> (64 - take);
+        }
+        w
+    };
+
+    let batched = (0..n_batches).flat_map(move |bi| {
+        let batch_base = bi * LANES;
+        let mut words = [0u64; LANES];
+        for (j, word) in words.iter_mut().enumerate() {
+            *word = word_at(batch_base + j);
+        }
+        let hits = Simd::<u64, LANES>::from_array(words)
+            .simd_ne(Simd::<u64, LANES>::splat(0))
+            .to_bitmask() as u64;
+        let mut pending = hits;
+        std::iter::from_fn(move || {
+            if pending == 0 {
+                return None;
+            }
+            let lane = pending.trailing_zeros() as usize;
+            let w = &mut words[lane];
+            let tz = w.trailing_zeros() as usize;
+            *w &= *w - 1;
+            if *w == 0 {
+                pending &= pending - 1;
+            }
+            Some((batch_base + lane) * 64 + tz)
+        })
+    });
+
+    let tail = (n_batches * LANES..n_words).flat_map(move |wi| {
+        let mut w = word_at(wi);
+        let base = wi * 64;
+        std::iter::from_fn(move || {
+            if w == 0 {
+                None
+            } else {
+                let tz = w.trailing_zeros() as usize;
+                w &= w - 1;
+                Some(base + tz)
+            }
+        })
+    });
+
+    batched.chain(tail)
+}
+
 /// Generates a SIMD equality mask function for a given element type and lane count.
 /// Processes LANES elements per iteration, with a scalar tail for the remainder.
 macro_rules! impl_simd_eq_mask {
@@ -1134,6 +1212,34 @@ mod tests {
                     let all_false = Bitmask::new_set_all(64 * LANES, false);
                     assert!(all_false_mask_simd::<LANES>(&all_false));
                 }
+
+                #[test]
+                fn test_iter_window_bits_simd() {
+                    // Spans several full LANES batches plus a tail of words too few
+                    // to fill a batch, so both scan paths run.
+                    let n = 64 * LANES * 2 + 37;
+                    let bits: Vec<bool> = (0..n).map(|i| i % 3 == 0).collect();
+                    let mask = bm(&bits);
+
+                    let set: Vec<usize> = iter_window_bits_simd::<LANES>(slice(&mask), true).collect();
+                    let expected_set: Vec<usize> = (0..n).filter(|&i| bits[i]).collect();
+                    assert_eq!(set, expected_set);
+
+                    let cleared: Vec<usize> =
+                        iter_window_bits_simd::<LANES>(slice(&mask), false).collect();
+                    let expected_cleared: Vec<usize> = (0..n).filter(|&i| !bits[i]).collect();
+                    assert_eq!(cleared, expected_cleared);
+
+                    // A bit-unaligned window offset crosses word boundaries within
+                    // every batch, exercising load_word's shift-combine path.
+                    let offset = 13;
+                    let len = n - offset - 5;
+                    let windowed: Vec<usize> =
+                        iter_window_bits_simd::<LANES>((&mask, offset, len), true).collect();
+                    let expected_windowed: Vec<usize> =
+                        (0..len).filter(|&i| bits[offset + i]).collect();
+                    assert_eq!(windowed, expected_windowed);
+                }
             }
         };
     }

src/kernels/bitmask/std.rs

@@ -51,7 +51,7 @@ use crate::{Bitmask, BitmaskVT};
 
 use crate::{
     enums::operators::{LogicalOperator, UnaryOperator},
-    kernels::bitmask::{bitmask_window_bytes, bitmask_window_bytes_mut},
+    kernels::bitmask::{bitmask_window_bytes, bitmask_window_bytes_mut, load_word},
 };
 
 /// Performs bitwise binary operations (AND/OR/XOR) over two bitmask slices using word-level processing.
@@ -518,6 +518,37 @@ pub fn all_false_mask(mask: &Bitmask) -> bool {
     true
 }
 
+/// Scans the bits of the window `m` one 64-bit word at a time, returning
+/// window-relative indices where the bit equals `bit_value`. A word
+/// containing no matching bits is skipped with a single comparison, and
+/// each match is located through a trailing-zeros count, so the traversal
+/// cost scales with the number of matches rather than the number of bits
+/// scanned.
+pub fn iter_window_bits(m: BitmaskVT<'_>, bit_value: bool) -> impl Iterator<Item = usize> + '_ {
+    let (mask, offset, len) = m;
+    let n_words = len.div_ceil(64);
+    (0..n_words).flat_map(move |wi| {
+        let base = wi * 64;
+        let take = (len - base).min(64);
+        let mut w = load_word(mask, offset + base);
+        if !bit_value {
+            w = !w;
+        }
+        if take < 64 {
+            w &= u64::MAX >> (64 - take);
+        }
+        std::iter::from_fn(move || {
+            if w == 0 {
+                None
+            } else {
+                let tz = w.trailing_zeros() as usize;
+                w &= w - 1;
+                Some(base + tz)
+            }
+        })
+    })
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

src/structs/bitmask.rs

@@ -32,6 +32,7 @@ use std::fmt::{Debug, Display, Formatter, Result as FmtResult};
 use std::ops::{BitAnd, BitOr, Deref, DerefMut, Index, Not};
 
 use crate::enums::shape_dim::ShapeDim;
+use crate::kernels::bitmask::dispatch::iter_window_bits;
 #[cfg(feature = "lbuffer")]
 use crate::structs::lbuffer::LBufferV;
 use crate::traits::concatenate::Concatenate;
@@ -983,29 +984,19 @@ impl Bitmask {
 
     /// Iterator over all indices with set bits (valid).
     ///
-    /// Under the `lbuffer` feature an LBuffer-backed mask reads each bit
-    /// through [`get`](Self::get) for a consistent view of the producer's
-    /// published bits. Any other storage takes the plain byte scan, so a
-    /// non-LBuffer-backed mask pays nothing for the feature being enabled.
+    /// The scan walks the mask one 64-bit word at a time, so a word with no
+    /// set bits costs a single comparison and each set bit is located
+    /// through a trailing-zeros count.
     pub fn iter_set(&self) -> impl Iterator<Item = usize> + '_ {
+
+        // Exception case ensuring atomic consistency
         #[cfg(feature = "lbuffer")]
         if self.is_lbuffer_backed() {
             return Box::new((0..self.len()).filter(move |&i| self.get(i)))
                 as Box<dyn Iterator<Item = usize> + '_>;
         }
 
-        let n = self.len();
-        let scan = self.bits.iter().enumerate().flat_map(move |(byte_i, &b)| {
-            let base = byte_i * 8;
-            (0..8).filter_map(move |bit| {
-                let idx = base + bit;
-                if idx < n && ((b >> bit) & 1) != 0 {
-                    Some(idx)
-                } else {
-                    None
-                }
-            })
-        });
+        let scan = iter_window_bits((self, 0, self.len()), true);
 
         #[cfg(feature = "lbuffer")]
         return Box::new(scan) as Box<dyn Iterator<Item = usize> + '_>;
@@ -1015,29 +1006,19 @@ impl Bitmask {
 
     /// Iterator over all indices with cleared bits (nulls).
     ///
-    /// Under the `lbuffer` feature an LBuffer-backed mask reads each bit
-    /// through [`get`](Self::get) for a consistent view of the producer's
-    /// published bits. Any other storage takes the plain byte scan, so a
-    /// non-LBuffer-backed mask pays nothing for the feature being enabled.
+    /// The scan walks the mask one 64-bit word at a time, so a fully valid
+    /// word costs a single comparison and each null is located through a
+    /// trailing-zeros count.
     pub fn iter_cleared(&self) -> impl Iterator<Item = usize> + '_ {
+
+        // Exception case ensuring atomic consistency
         #[cfg(feature = "lbuffer")]
         if self.is_lbuffer_backed() {
             return Box::new((0..self.len()).filter(move |&i| !self.get(i)))
                 as Box<dyn Iterator<Item = usize> + '_>;
         }
 
-        let n = self.len();
-        let scan = self.bits.iter().enumerate().flat_map(move |(byte_i, &b)| {
-            let base = byte_i * 8;
-            (0..8).filter_map(move |bit| {
-                let idx = base + bit;
-                if idx < n && ((b >> bit) & 1) == 0 {
-                    Some(idx)
-                } else {
-                    None
-                }
-            })
-        });
+        let scan = iter_window_bits((self, 0, self.len()), false);
 
         #[cfg(feature = "lbuffer")]
         return Box::new(scan) as Box<dyn Iterator<Item = usize> + '_>;