Optimize bit iteration and validity checking with SIMD popcount

AGENTS.md

@@ -151,6 +151,47 @@ Notes:
   self-explanatory code.
 - Keep public APIs small and consistent with neighboring crates.
 
+## Performance: avoid hidden-cost accessors in hot loops
+
+The most common performance trap in this codebase is calling a *per-element accessor that
+hides non-trivial work* inside an `O(n)` loop, instead of doing the work once over the whole
+chunk. The call site looks like a cheap getter, but each call re-pays a cost that is constant
+(or amortizable) across the loop, making the loop `O(n · k)`.
+
+Watch for these accessors used inside `for i in 0..n { ... }`:
+
+| Per-element accessor (in a loop) | Hidden cost | Bulk replacement |
+| --- | --- | --- |
+| `Validity::is_valid(i)` / `is_null(i)` | for array-backed validity, **allocates an `ExecutionCtx` and runs a scalar lookup per call** | `validity.execute_mask(len, ctx)?` once, then `Mask::value(i)` |
+| `array.scalar_at(i)` / `array.execute_scalar(i, ctx)` | per-element execution through the compute stack | canonicalize once (`execute::<PrimitiveArray>` / `as_slice`) then index |
+| `BitBuffer::value(i)` / `Mask::value(i)` accumulated into a count | recomputes the byte address each call; defeats popcount | `true_count()`, `BitBuffer::count_range(start, end)`, `set_indices()` |
+| `BitIterator::next()` to accumulate a running rank/prefix count | bit-at-a-time | `count_range` over each gap (SIMD popcount) |
+| re-deriving a value inside the loop (e.g. `self.validity()?` each iteration) | re-runs the derivation `n` times | hoist it above the loop |
+
+Decide per site — bulk is not always the answer:
+
+- **Sequential / contiguous access** over an accessor that hides amortizable work → hoist and
+  go bulk (materialize once, then index or iterate the chunk).
+- **Gather over arbitrary indices** → you cannot amortize a per-element *decode*, but you can
+  still materialize the backing buffer once (e.g. `execute_mask`) and then do cheap `O(1)`
+  random reads, avoiding per-call context/allocation.
+- **The accessor is already genuinely `O(1)`** (e.g. reading an already-materialized `Mask`/
+  slice, or a native bitmap) → leave it; bulk would not help.
+
+Even after materializing into a `Mask`/`BitBuffer`, **do not loop with `value(i)` per
+element** to act on each set bit — the per-element branch dominates. Iterate a `u64` word
+at a time with all-set / all-unset fast paths: use [`BitBuffer::for_each_set_index`]. It
+beats `for i in 0..len { if buf.value(i) {..} }` by 2-45x (more at low density) and beats
+collecting `set_indices()` by ~2x at mid/high density, while self-adapting from sparse to
+dense — see `vortex-mask/benches/mask_iteration.rs`. Reach for the cached `indices()` /
+`slices()` representations when you need them more than once; otherwise `for_each_set_index`
+needs no materialization.
+
+When you touch such a loop, back the change with a benchmark (see
+`vortex-array/benches/validity_is_valid.rs` for the `is_valid` case,
+`vortex-mask/benches/valid_counts.rs` for the popcount case, and
+`vortex-mask/benches/mask_iteration.rs` for the per-element-vs-word-vs-sparse comparison).
+
 ## Tests
 
 - Strongly consider `rstest` cases when parameterizing repetitive test logic.
@@ -179,6 +220,9 @@ Check new and modified lines against this list before finishing:
 - Updating expected test output to match buggy behavior without independently verifying the
   intended semantics.
 - Silently reducing the scope of an approved plan when implementation is harder than expected.
+- Calling a hidden-cost per-element accessor (`Validity::is_valid`, `scalar_at`, `BitBuffer::
+  value` accumulation) inside a hot loop instead of materializing once — see
+  "Performance: avoid hidden-cost accessors in hot loops".
 
 ## Summaries
 

encodings/runend/src/compute/filter.rs

@@ -5,6 +5,7 @@ use std::cmp::min;
 use std::ops::AddAssign;
 
 use num_traits::AsPrimitive;
+use num_traits::NumCast;
 use vortex_array::ArrayRef;
 use vortex_array::ArrayView;
 use vortex_array::ExecutionCtx;
@@ -87,16 +88,21 @@ pub fn filter_run_end_primitive<R: NativePType + AddAssign + From<bool> + AsPrim
     let mut count = R::zero();
 
     let new_mask: Mask = BitBuffer::collect_bool(run_ends.len(), |i| {
-        let mut keep = false;
         let end = min(run_ends[i].as_() - offset, length);
 
-        // Safety: predicate must be the same length as the array the ends have been taken from
-        for pred in (start..end).map(|i| unsafe {
-            mask.value_unchecked(i.try_into().vortex_expect("index must fit in usize"))
-        }) {
-            count += <R as From<bool>>::from(pred);
-            keep |= pred
-        }
+        // SIMD popcount of the predicate bits in this run. The range matches the
+        // bit-by-bit `value_unchecked` read it replaces, so `end <= mask.len()`.
+        let start_usize = start
+            .try_into()
+            .vortex_expect("run start index must fit in usize");
+        let end_usize = end
+            .try_into()
+            .vortex_expect("run end index must fit in usize");
+        let run_trues = mask.count_range(start_usize, end_usize);
+        count += <R as NumCast>::from(run_trues)
+            .vortex_expect("run popcount must fit in run-end native type");
+        let keep = run_trues > 0;
+
         // this is to avoid branching
         new_run_ends[j] = count;
         j += keep as usize;

vortex-buffer/src/bit/buf.rs

@@ -379,6 +379,20 @@ impl BitBuffer {
         count_ones(self.buffer.as_slice(), self.offset, self.len)
     }
 
+    /// Get the number of set bits in the bit range `[start, end)`.
+    ///
+    /// Unlike `self.slice(start..end).true_count()`, this counts directly over the
+    /// existing backing buffer without allocating or cloning a new [`BitBuffer`],
+    /// making it cheap to call repeatedly over many small ranges.
+    ///
+    /// Panics if `start > end` or `end > len`.
+    #[inline]
+    pub fn count_range(&self, start: usize, end: usize) -> usize {
+        assert!(start <= end, "start {start} exceeds end {end}");
+        assert!(end <= self.len, "end {end} exceeds len {}", self.len);
+        count_ones(self.buffer.as_slice(), self.offset + start, end - start)
+    }
+
     /// Returns the position of the `nth` set bit (0-indexed).
     ///
     /// This is the "select" operation on a bitmap: given a rank `nth`, find
@@ -410,6 +424,33 @@ impl BitBuffer {
         BitSliceIterator::new(self.buffer.as_slice(), self.offset, self.len)
     }
 
+    /// Invoke `f(index)` for every set bit, in ascending order, processing a `u64`
+    /// word at a time.
+    ///
+    /// This is the fast way to "do something for each set bit": it skips all-zero
+    /// words, fast-paths all-one words, and walks the remaining bits with
+    /// `trailing_zeros`. Prefer it over `for i in 0..len { if buf.value(i) { f(i) } }`
+    /// (which pays a branch per element) and over collecting [`Self::set_indices`]
+    /// (whose per-`next` iterator state does not inline as well).
+    #[inline]
+    pub fn for_each_set_index<F: FnMut(usize)>(&self, mut f: F) {
+        let mut base = 0usize;
+        for word in self.chunks().iter_padded() {
+            if word == u64::MAX {
+                for k in 0..64 {
+                    f(base + k);
+                }
+            } else {
+                let mut w = word;
+                while w != 0 {
+                    f(base + w.trailing_zeros() as usize);
+                    w &= w - 1;
+                }
+            }
+            base += 64;
+        }
+    }
+
     /// Created a new BitBuffer with offset reset to 0
     pub fn sliced(&self) -> Self {
         if self.offset.is_multiple_of(8) {
@@ -857,6 +898,75 @@ mod tests {
         }
     }
 
+    #[rstest]
+    #[case(0, 0)]
+    #[case(0, 64)]
+    #[case(5, 70)]
+    #[case(64, 130)]
+    #[case(0, 200)]
+    fn test_count_range(#[case] start: usize, #[case] end: usize) {
+        let len = 200;
+        let buf = BitBuffer::collect_bool(len, |i| i % 3 == 0);
+        let expected = (start..end).filter(|i| i % 3 == 0).count();
+        assert_eq!(buf.count_range(start, end), expected);
+        // Must agree with slicing then counting.
+        assert_eq!(
+            buf.count_range(start, end),
+            buf.slice(start..end).true_count()
+        );
+    }
+
+    #[rstest]
+    #[case(3)]
+    #[case(7)]
+    fn test_count_range_with_offset(#[case] offset: usize) {
+        let len = 150;
+        let buf = BitBuffer::collect_bool(offset + len, |i| i % 2 == 0);
+        let view = BitBuffer::new_with_offset(buf.inner().clone(), len, offset);
+        for (start, end) in [(0, len), (10, 100), (1, 2), (63, 129)] {
+            let expected = (offset + start..offset + end)
+                .filter(|i| i % 2 == 0)
+                .count();
+            assert_eq!(view.count_range(start, end), expected, "[{start}, {end})");
+        }
+    }
+
+    #[rstest]
+    #[case(0)]
+    #[case(1)]
+    #[case(63)]
+    #[case(64)]
+    #[case(65)]
+    #[case(200)]
+    #[case(1000)]
+    fn test_for_each_set_index_matches_set_indices(#[case] len: usize) {
+        let buf = BitBuffer::collect_bool(len, |i| i % 5 == 0 || i % 7 == 0);
+        let expected: Vec<usize> = buf.set_indices().collect();
+        let mut got = Vec::new();
+        buf.for_each_set_index(|i| got.push(i));
+        assert_eq!(got, expected);
+    }
+
+    #[rstest]
+    #[case(3, 200)]
+    #[case(7, 130)]
+    fn test_for_each_set_index_with_offset(#[case] offset: usize, #[case] len: usize) {
+        let base = BitBuffer::collect_bool(offset + len, |i| i % 3 == 0);
+        let view = BitBuffer::new_with_offset(base.inner().clone(), len, offset);
+        let expected: Vec<usize> = view.set_indices().collect();
+        let mut got = Vec::new();
+        view.for_each_set_index(|i| got.push(i));
+        assert_eq!(got, expected);
+    }
+
+    #[test]
+    fn test_for_each_set_index_all_set() {
+        let buf = BitBuffer::new_set(130);
+        let mut got = Vec::new();
+        buf.for_each_set_index(|i| got.push(i));
+        assert_eq!(got, (0..130).collect::<Vec<_>>());
+    }
+
     #[test]
     fn test_map_cmp_conditional() {
         // map_cmp with conditional logic based on index and bit value

vortex-duckdb/src/exporter/bool.rs

@@ -46,13 +46,16 @@ impl ColumnExporter for BoolExporter {
     ) -> VortexResult<()> {
         // DuckDB uses byte bools, not bit bools.
         // maybe we can convert into these from a compressed array sometimes?.
-        unsafe { vector.as_slice_mut(len) }.copy_from_slice(
-            &self
-                .bit_buffer
-                .slice(offset..(offset + len))
-                .iter()
-                .collect::<Vec<bool>>(),
-        );
+        // Unpack the bits directly into the destination byte slice, avoiding the
+        // throwaway `Vec<bool>` allocation and the extra copy pass. The sliced
+        // iterator already accounts for the buffer's bit offset.
+        let dst = unsafe { vector.as_slice_mut(len) };
+        for (slot, bit) in dst
+            .iter_mut()
+            .zip(self.bit_buffer.slice(offset..(offset + len)).iter())
+        {
+            *slot = bit;
+        }
 
         Ok(())
     }

vortex-mask/benches/mask_iteration.rs

@@ -110,6 +110,18 @@ fn word_trailing_zeros(bencher: Bencher, (len, density): (usize, f64)) {
         });
 }
 
+#[divan::bench(args = ARGS)]
+fn for_each_set_index(bencher: Bencher, (len, density): (usize, f64)) {
+    let (buf, values) = make(len, density);
+    bencher
+        .with_inputs(|| (&buf, &values))
+        .bench_refs(|(buf, values)| {
+            let mut acc = 0u64;
+            buf.for_each_set_index(|i| acc = acc.wrapping_add(values[i]));
+            acc
+        });
+}
+
 #[divan::bench(args = ARGS)]
 fn set_slices(bencher: Bencher, (len, density): (usize, f64)) {
     let (buf, values) = make(len, density);

vortex-mask/src/lib.rs

@@ -628,23 +628,24 @@ impl Mask {
     /// Given monotonically increasing `indices` in [0, n_rows], returns the
     /// count of valid elements up to each index.
     ///
-    /// This is O(n_rows).
+    /// This is O(n_rows), but the per-gap counts are computed with a SIMD
+    /// popcount over the underlying bit buffer rather than walking bit-by-bit.
     pub fn valid_counts_for_indices(&self, indices: &[usize]) -> Vec<usize> {
         match self {
             Self::AllTrue(_) => indices.to_vec(),
             Self::AllFalse(_) => vec![0; indices.len()],
             Self::Values(values) => {
-                let mut bool_iter = values.bit_buffer().iter();
+                let buffer = values.bit_buffer();
                 let mut valid_counts = Vec::with_capacity(indices.len());
                 let mut valid_count = 0;
-                let mut idx = 0;
+                let mut prev = 0;
                 for &next_idx in indices {
-                    while idx < next_idx {
-                        idx += 1;
-                        valid_count += bool_iter
-                            .next()
-                            .unwrap_or_else(|| vortex_panic!("Row indices exceed array length"))
-                            as usize;
+                    assert!(next_idx <= buffer.len(), "Row indices exceed array length");
+                    // `indices` is monotonically increasing, so each gap is counted once;
+                    // the total work across all gaps scans the prefix `[0, last_idx)` once.
+                    if next_idx > prev {
+                        valid_count += buffer.count_range(prev, next_idx);
+                        prev = next_idx;
                     }
                     valid_counts.push(valid_count);
                 }
@@ -779,7 +780,10 @@ impl MaskValues {
             }
 
             let mut indices = Vec::with_capacity(self.true_count);
-            indices.extend(self.buffer.set_indices());
+            // Word-at-a-time set-bit walk; faster than collecting `set_indices()`,
+            // whose per-`next` iterator state inlines less well (see
+            // `vortex-mask/benches/mask_iteration.rs`).
+            self.buffer.for_each_set_index(|i| indices.push(i));
             debug_assert!(indices.is_sorted());
             assert_eq!(indices.len(), self.true_count);
             indices