From 974258f19bcb5785536ec7ac20650452b5c309ca Mon Sep 17 00:00:00 2001 From: Peter Bower <37089506+pbower@users.noreply.github.com> Date: Mon, 10 Aug 2026 23:41:47 +0100 Subject: [PATCH] Add Bitmask iteration kernels --- src/kernels/bitmask/dispatch.rs | 26 +++++++ src/kernels/bitmask/mod.rs | 24 +++++++ src/kernels/bitmask/simd.rs | 108 +++++++++++++++++++++++++++++- src/kernels/bitmask/std.rs | 33 ++++++++- src/structs/bitmask.rs | 45 ++++--------- src/structs/views/bitmask_view.rs | 84 ++++++++++++++++++++++- 6 files changed, 284 insertions(+), 36 deletions(-) diff --git a/src/kernels/bitmask/dispatch.rs b/src/kernels/bitmask/dispatch.rs index 89a72bb..db58d98 100644 --- a/src/kernels/bitmask/dispatch.rs +++ b/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 + '_ { + #[cfg(feature = "simd")] + { + crate::kernels::bitmask::simd::iter_window_bits_simd::(m, bit_value) + } + #[cfg(not(feature = "simd"))] + { + crate::kernels::bitmask::std::iter_window_bits(m, bit_value) + } +} diff --git a/src/kernels/bitmask/mod.rs b/src/kernels/bitmask/mod.rs index 4281fa1..fd51f4b 100644 --- a/src/kernels/bitmask/mod.rs +++ b/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) { diff --git a/src/kernels/bitmask/simd.rs b/src/kernels/bitmask/simd.rs index 317b1e2..bbfc6dc 100644 --- a/src/kernels/bitmask/simd.rs +++ b/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( + m: BitmaskVT<'_>, + bit_value: bool, +) -> impl Iterator + '_ +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::::from_array(words) + .simd_ne(Simd::::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::(&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 = (0..n).map(|i| i % 3 == 0).collect(); + let mask = bm(&bits); + + let set: Vec = iter_window_bits_simd::(slice(&mask), true).collect(); + let expected_set: Vec = (0..n).filter(|&i| bits[i]).collect(); + assert_eq!(set, expected_set); + + let cleared: Vec = + iter_window_bits_simd::(slice(&mask), false).collect(); + let expected_cleared: Vec = (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 = + iter_window_bits_simd::((&mask, offset, len), true).collect(); + let expected_windowed: Vec = + (0..len).filter(|&i| bits[offset + i]).collect(); + assert_eq!(windowed, expected_windowed); + } } }; } diff --git a/src/kernels/bitmask/std.rs b/src/kernels/bitmask/std.rs index ed0072a..7907450 100644 --- a/src/kernels/bitmask/std.rs +++ b/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 + '_ { + 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::*; diff --git a/src/structs/bitmask.rs b/src/structs/bitmask.rs index 4aa9f8d..9b848d1 100644 --- a/src/structs/bitmask.rs +++ b/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 + '_ { + + // 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 + '_>; } - 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 + '_>; @@ -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 + '_ { + + // 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 + '_>; } - 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 + '_>; diff --git a/src/structs/views/bitmask_view.rs b/src/structs/views/bitmask_view.rs index 0b60079..a1086a7 100644 --- a/src/structs/views/bitmask_view.rs +++ b/src/structs/views/bitmask_view.rs @@ -50,6 +50,7 @@ use std::ops::Index; #[cfg(feature = "views")] use crate::ArrayV; use crate::enums::shape_dim::ShapeDim; +use crate::kernels::bitmask::dispatch::iter_window_bits; use crate::traits::print::MAX_PREVIEW; use crate::traits::shape::Shape; use crate::{Array, Bitmask, BitmaskVT, BooleanArray}; @@ -158,13 +159,48 @@ impl<'a> BitmaskV<'a> { } /// Returns an iterator over all set bits (indices relative to the window). + /// + /// The scan walks the backing mask one 64-bit word at a time from the + /// window's bit offset, 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 + '_ { - (0..self.len).filter(move |&i| self.get(i)) + + // Exception case verifies per bit ensuring atomic consistency + #[cfg(feature = "lbuffer")] + if self.bitmask.is_lbuffer_backed() { + return Box::new((0..self.len).filter(move |&i| self.get(i))) + as Box + '_>; + } + + let scan = iter_window_bits((self.bitmask, self.offset, self.len), true); + + #[cfg(feature = "lbuffer")] + return Box::new(scan) as Box + '_>; + #[cfg(not(feature = "lbuffer"))] + return scan; } /// Returns an iterator over all cleared bits (indices relative to the window). + /// + /// The scan walks the backing mask one 64-bit word at a time from the + /// window's bit offset, 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 + '_ { - (0..self.len).filter(move |&i| !self.get(i)) + + // Exception case verifies per bit ensuring atomic consistency + #[cfg(feature = "lbuffer")] + if self.bitmask.is_lbuffer_backed() { + return Box::new((0..self.len).filter(move |&i| !self.get(i))) + as Box + '_>; + } + + let scan = iter_window_bits((self.bitmask, self.offset, self.len), false); + + #[cfg(feature = "lbuffer")] + return Box::new(scan) as Box + '_>; + #[cfg(not(feature = "lbuffer"))] + return scan; } /// Counts number of set bits in the view. @@ -492,4 +528,48 @@ mod tests { assert!(!out.get(1)); assert!(out.get(2)); } + + #[test] + fn test_iterators_match_per_bit_definition_across_offsets() { + let n = 200; + let bits: Vec = (0..n).map(|i| (i * 7 + i / 13) % 3 == 0).collect(); + let mask = Bitmask::from_bools(&bits); + + for offset in [0, 1, 5, 7, 8, 63, 64, 65, 127, 128, 130] { + for len in [0, 1, 3, 63, 64, 65, n - offset] { + if offset + len > n { + continue; + } + let view = BitmaskV::new(&mask, offset, len); + let set: Vec = view.iter_set().collect(); + let expected_set: Vec = + (0..len).filter(|&i| bits[offset + i]).collect(); + assert_eq!(set, expected_set, "iter_set at offset {offset} len {len}"); + + let cleared: Vec = view.iter_cleared().collect(); + let expected_cleared: Vec = + (0..len).filter(|&i| !bits[offset + i]).collect(); + assert_eq!( + cleared, expected_cleared, + "iter_cleared at offset {offset} len {len}" + ); + } + } + } + + #[test] + fn test_iterators_uniform_windows() { + let all_set = Bitmask::from_bools(&[true; 130]); + let view = BitmaskV::new(&all_set, 3, 120); + assert_eq!(view.iter_cleared().count(), 0); + assert_eq!(view.iter_set().count(), 120); + + let all_clear = Bitmask::from_bools(&[false; 130]); + let view = BitmaskV::new(&all_clear, 3, 120); + assert_eq!(view.iter_set().count(), 0); + assert_eq!( + view.iter_cleared().collect::>(), + (0..120).collect::>() + ); + } }