Refactor bit matrix operations and optimize distance computations

Consolidate row extraction and parallel column reduction into reusable generic utilities within the pairwise module. Replace manual iteration in builder and view methods with iterator-based zipped loops and deferred overflow processing to improve memory access patterns. Optimize distance functions by switching to bulk byte passes for branch-free SIMD vectorization, adding a merge-like helper to correct masked overflow values without secondary allocations. Add a comprehensive test validating the optimized distance paths against naive references across multiple thresholds and overflow scenarios.

src/obicompactvec/src/bitmatrix/columnar.rs

@@ -2,13 +2,12 @@ use std::io;
 use std::path::Path;
 
 use ndarray::{Array1, Array2};
-use rayon::prelude::*;
 
 use crate::bitvec::{PersistentBitVec, PersistentBitVecBuilder};
 use crate::meta::MatrixMeta;
 
 use super::col_path;
-use super::pairwise::{pairwise_matrix, pairwise2_matrix};
+use super::pairwise::{fill_row_generic, par_col_reduce, pairwise2_matrix, pairwise_matrix, row_generic};
 
 // ── ColumnarBitMatrix ─────────────────────────────────────────────────────────
 
@@ -36,22 +35,16 @@ impl ColumnarBitMatrix {
 
     #[inline]
     pub(crate) fn row(&self, slot: usize) -> Box<[bool]> {
-        self.cols.iter().map(|c| c.get(slot)).collect()
+        row_generic(slot, self.n_cols(), |c, s| self.col(c).get(s))
     }
 
     #[inline]
     pub(crate) fn fill_row(&self, slot: usize, buf: &mut [u32]) {
-        for (c, col) in self.cols.iter().enumerate() {
+        fill_row_generic(slot, buf, |c, s| self.col(c).get(s) as u32)
-            buf[c] = col.get(slot) as u32;
-        }
     }
 
     pub(crate) fn count_ones(&self) -> Array1<u64> {
-        let counts: Vec<u64> = (0..self.n_cols())
+        par_col_reduce(self.n_cols(), |c| self.col(c).count_ones())
-            .into_par_iter()
-            .map(|c| self.col(c).count_ones())
-            .collect();
-        Array1::from_vec(counts)
     }
 
     pub(crate) fn partial_jaccard_dist_matrix(&self) -> (Array2<u64>, Array2<u64>) {

src/obicompactvec/src/bitmatrix/mod.rs

@@ -24,7 +24,7 @@ pub use packed::pack_bit_matrix;
 pub use persistent::PersistentBitMatrix;
 pub use sparse::{PersistentSparseBitMatrix, PersistentSparseBitMatrixBuilder, pack_sparse_bit_matrix};
 
-pub(crate) use pairwise::{pairwise_matrix, pairwise2_matrix};
+pub(crate) use pairwise::{fill_row_generic, par_col_reduce, pairwise2_matrix, pairwise_matrix, row_generic};
 
 fn col_path(dir: &Path, col: usize) -> PathBuf {
     dir.join(format!("col_{col:06}.pbiv"))

src/obicompactvec/src/bitmatrix/packed.rs

@@ -4,14 +4,13 @@ use std::path::Path;
 
 use memmap2::Mmap;
 use ndarray::{Array1, Array2};
-use rayon::prelude::*;
 
 use crate::bitvec::PersistentBitVecBuilder;
 use crate::meta::MatrixMeta;
 use crate::views::BitSliceView;
 
 use super::col_path;
-use super::pairwise::{pairwise_matrix, pairwise2_matrix};
+use super::pairwise::{fill_row_generic, par_col_reduce, pairwise2_matrix, pairwise_matrix, row_generic};
 
 // ── PackedBitMatrix ───────────────────────────────────────────────────────────
 
@@ -53,16 +52,12 @@ impl PackedBitMatrix {
 
     #[inline]
     pub(crate) fn fill_row(&self, slot: usize, buf: &mut [u32]) {
-        for (c, &data_off) in self.data_offsets.iter().enumerate() {
+        fill_row_generic(slot, buf, |c, s| self.col_slice(c).get(s) as u32)
-            buf[c] = ((self.mmap[data_off + (slot >> 3)] >> (slot & 7)) & 1) as u32;
-        }
     }
 
     #[inline]
     pub(crate) fn row(&self, slot: usize) -> Box<[bool]> {
-        (0..self.n_cols).map(|c| {
+        row_generic(slot, self.n_cols, |c, s| self.col_slice(c).get(s))
-            (self.mmap[self.data_offsets[c] + (slot >> 3)] >> (slot & 7)) & 1 != 0
-        }).collect()
     }
 
     #[inline]
@@ -91,11 +86,7 @@ impl PackedBitMatrix {
     }
 
     pub(crate) fn count_ones(&self) -> Array1<u64> {
-        Array1::from_vec(
+        par_col_reduce(self.n_cols, |c| self.col_slice(c).count_ones())
-            (0..self.n_cols).into_par_iter()
-                .map(|c| self.col_slice(c).count_ones())
-                .collect()
-        )
     }
 
     pub(crate) fn partial_jaccard_dist_matrix(&self) -> (Array2<u64>, Array2<u64>) {

src/obicompactvec/src/bitmatrix/pairwise.rs

@@ -1,8 +1,30 @@
-use ndarray::Array2;
+use ndarray::{Array1, Array2};
 use rayon::prelude::*;
 
 // ── Shared matrix helpers (also used by intmatrix.rs) ─────────────────────────
 
+/// Shared `row`/`fill_row` bodies for matrix backends (`Columnar`/`Packed`,
+/// bit or int) whose only real difference is how to fetch one column's
+/// value at a slot — reused so the same `(0..n_cols).map(|c| get(c,
+/// slot))` loop isn't hand-written once per format, mirroring
+/// `views::nonzero_triples`'s reuse across the same backends.
+pub(crate) fn row_generic<T: Copy>(slot: usize, n_cols: usize, get: impl Fn(usize, usize) -> T) -> Box<[T]> {
+    (0..n_cols).map(|c| get(c, slot)).collect()
+}
+
+/// Like [`row_generic`], filling a caller-provided buffer instead of allocating.
+pub(crate) fn fill_row_generic<T: Copy>(slot: usize, buf: &mut [T], get: impl Fn(usize, usize) -> T) {
+    for (c, out) in buf.iter_mut().enumerate() { *out = get(c, slot); }
+}
+
+/// Shared per-column parallel reduction: `count_ones`/`sum`/`count_nonzero`
+/// all have this exact shape (`(0..n_cols).into_par_iter().map(f).collect()`)
+/// across every `Columnar`/`Packed` backend — written once here instead of
+/// once per format.
+pub(crate) fn par_col_reduce<T: Send>(n_cols: usize, f: impl Fn(usize) -> T + Sync + Send) -> Array1<T> {
+    Array1::from_vec((0..n_cols).into_par_iter().map(f).collect())
+}
+
 fn upper_pairs(n: usize) -> Vec<(usize, usize)> {
     (0..n).flat_map(|i| (i + 1..n).map(move |j| (i, j))).collect()
 }

src/obicompactvec/src/builder.rs

@@ -192,20 +192,33 @@ impl PersistentCompactIntVecBuilder {
     }
 
     pub fn add(&mut self, other: IntSliceView<'_>) {
-        let n = self.n;
+        // Fast path: both operands fit in a byte — mutate in place, zipped,
-        for s in 0..n {
+        // no repeated slice-indexing per element. Slots that already
-            let sb = self.primary_bytes()[s];
+        // overflow (either side) or newly overflow (sum ≥ 255) are
-            let ob = other.primary_bytes()[s];
+        // collected here and settled after the loop, via `get`/`set`
-            if sb < 255 && ob < 255 {
+        // (which need the overflow map, not touchable while it's borrowed
-                let sum = sb as u32 + ob as u32;
+        // through this loop).
-                if sum < 255 { self.primary_bytes_mut()[s] = sum as u8; }
+        let mut newly_overflow: Vec<(usize, u32)> = Vec::new();
-                else         { self.set(s, sum); }
+        let mut already_overflow: Vec<usize> = Vec::new();
+        for (s, (a, &b)) in self.primary_bytes_mut().iter_mut().zip(other.primary_bytes()).enumerate() {
+            if *a < 255 && b < 255 {
+                let sum = *a as u32 + b as u32;
+                if sum < 255 {
+                    *a = sum as u8;
+                } else {
+                    newly_overflow.push((s, sum));
+                    *a = 255;
+                }
             } else {
-                let sv = self.get(s);
+                already_overflow.push(s);
-                let ov = other.get(s);
-                self.set(s, sv + ov);
             }
         }
+        for (s, sum) in newly_overflow { self.overflow.insert(s, sum); }
+        for s in already_overflow {
+            let sv = self.get(s);
+            let ov = other.get(s);
+            self.set(s, sv + ov);
+        }
     }
 
     pub fn min(&mut self, other: IntSliceView<'_>) {
@@ -233,18 +246,26 @@ impl PersistentCompactIntVecBuilder {
     }
 
     pub fn diff(&mut self, other: IntSliceView<'_>) {
-        let n = self.n;
+        // Same shape as `add`: fast byte-level path zipped in place: when
-        for s in 0..n {
+        // `self` isn't overflow, either `other` isn't either (plain
-            let sb = self.primary_bytes()[s];
+        // saturating subtract) or `other`'s true value is ≥ 255 — known
-            let ob = other.primary_bytes()[s];
+        // larger than any non-overflow `self` value, so the result
-            if sb < 255 {
+        // saturates to 0 without needing `other.get`. Only `self`-overflow
-                self.primary_bytes_mut()[s] = if ob < 255 { sb.saturating_sub(ob) } else { 0 };
+        // slots need the overflow map, settled after the loop.
+        let mut needs_full: Vec<usize> = Vec::new();
+        for (s, (a, &b)) in self.primary_bytes_mut().iter_mut().zip(other.primary_bytes()).enumerate() {
+            if *a < 255 {
+                *a = if b < 255 { a.saturating_sub(b) } else { 0 };
             } else {
-                let sv = self.get(s);
+                needs_full.push(s);
-                let ov = if ob < 255 { ob as u32 } else { other.get(s) };
-                self.set(s, sv.saturating_sub(ov));
             }
         }
+        for s in needs_full {
+            let sv = self.get(s);
+            let ob = other.primary_bytes()[s];
+            let ov = if ob < 255 { ob as u32 } else { other.get(s) };
+            self.set(s, sv.saturating_sub(ov));
+        }
     }
 
     pub fn mask_with(&mut self, mask: BitSliceView<'_>) {

src/obicompactvec/src/intmatrix.rs

@@ -4,9 +4,8 @@ use std::path::{Path, PathBuf};
 
 use memmap2::Mmap;
 use ndarray::{Array1, Array2};
-use rayon::prelude::*;
 
-use crate::bitmatrix::{pairwise_matrix, pairwise2_matrix};
+use crate::bitmatrix::{fill_row_generic, par_col_reduce, pairwise2_matrix, pairwise_matrix, row_generic};
 use crate::builder::PersistentCompactIntVecBuilder;
 use crate::colgroup::{chunked_presence_count, ColGroup, MatrixGroupOps};
 use crate::format::{HEADER_SIZE, OVERFLOW_ENTRY_SIZE};
@@ -45,28 +44,20 @@ impl ColumnarCompactIntMatrix {
 
     #[inline]
     pub(crate) fn row(&self, slot: usize) -> Box<[u32]> {
-        self.cols.iter().map(|c| c.get(slot)).collect()
+        row_generic(slot, self.n_cols(), |c, s| self.col(c).get(s))
     }
 
     #[inline]
     pub(crate) fn fill_row(&self, slot: usize, buf: &mut [u32]) {
-        for (c, col) in self.cols.iter().enumerate() { buf[c] = col.get(slot); }
+        fill_row_generic(slot, buf, |c, s| self.col(c).get(s))
     }
 
     pub(crate) fn sum(&self) -> Array1<u64> {
-        let sums: Vec<u64> = (0..self.n_cols())
+        par_col_reduce(self.n_cols(), |c| self.col(c).sum())
-            .into_par_iter()
-            .map(|c| self.col(c).sum())
-            .collect();
-        Array1::from_vec(sums)
     }
 
     pub(crate) fn count_nonzero(&self) -> Array1<u64> {
-        let counts: Vec<u64> = (0..self.n_cols())
+        par_col_reduce(self.n_cols(), |c| self.col(c).count_nonzero())
-            .into_par_iter()
-            .map(|c| self.col(c).count_nonzero())
-            .collect();
-        Array1::from_vec(counts)
     }
 
     pub(crate) fn partial_bray_dist_matrix(&self) -> Array2<u64> {
@@ -166,24 +157,20 @@ impl PackedCompactIntMatrix {
 
     #[inline]
     pub(crate) fn fill_row(&self, slot: usize, buf: &mut [u32]) {
-        for c in 0..self.n_cols { buf[c] = self.get(c, slot); }
+        fill_row_generic(slot, buf, |c, s| self.get(c, s))
     }
 
     #[inline]
     pub(crate) fn row(&self, slot: usize) -> Box<[u32]> {
-        (0..self.n_cols).map(|c| self.get(c, slot)).collect()
+        row_generic(slot, self.n_cols, |c, s| self.get(c, s))
     }
 
     pub(crate) fn sum(&self) -> Array1<u64> {
-        Array1::from_vec(
+        par_col_reduce(self.n_cols, |c| self.col_view(c).sum())
-            (0..self.n_cols).into_par_iter().map(|c| self.col_view(c).sum()).collect()
-        )
     }
 
     pub(crate) fn count_nonzero(&self) -> Array1<u64> {
-        Array1::from_vec(
+        par_col_reduce(self.n_cols, |c| self.col_view(c).count_nonzero())
-            (0..self.n_cols).into_par_iter().map(|c| self.col_view(c).count_nonzero()).collect()
-        )
     }
 
     // Every pair formula below delegates to `IntSliceView` (`views.rs`) —

src/obicompactvec/src/tests/mod.rs

@@ -331,6 +331,41 @@ fn threshold_jaccard_dist_with_threshold() {
     assert!((d - 2.0 / 3.0).abs() < 1e-12, "got {d}");
 }
 
+/// Exercises `partial_bray_dist`/`partial_euclidean_dist`/
+/// `partial_threshold_jaccard_dist`'s bulk-over-primary-bytes-then-correct
+/// path (see `views::IntSliceView::for_each_overflow_pair`) against a naive
+/// per-slot reference computed via `get`. Overflow slots are scattered so
+/// every union pattern the merge walk must handle is covered: overflow only
+/// in `a`, only in `b`, in both at the same slot, and in both at different
+/// slots — plus plain non-overflow slots in between.
+#[test]
+fn partial_dists_bulk_path_matches_naive_reference_with_overflow() {
+    let a = [10u32, 300, 20, 1000, 5, 500, 0, 254];
+    let b = [400u32, 7, 20, 1000, 600, 500, 255, 254];
+    let (_da, ra) = make_pciv(&a);
+    let (_db, rb) = make_pciv(&b);
+
+    let naive_bray: u64 = a.iter().zip(&b).map(|(&x, &y)| x.min(y) as u64).sum();
+    assert_eq!(ra.partial_bray_dist(&rb), naive_bray);
+
+    let naive_euclidean: f64 = a.iter().zip(&b)
+        .map(|(&x, &y)| { let d = x as f64 - y as f64; d * d })
+        .sum();
+    assert!((ra.partial_euclidean_dist(&rb) - naive_euclidean).abs() < 1e-9);
+
+    for threshold in [1u32, 6, 255, 600] {
+        let (naive_inter, naive_union) = a.iter().zip(&b)
+            .fold((0u64, 0u64), |(i, u), (&x, &y)| {
+                let xp = x >= threshold;
+                let yp = y >= threshold;
+                (i + (xp & yp) as u64, u + (xp | yp) as u64)
+            });
+        let (inter, union) = ra.partial_threshold_jaccard_dist(&rb, threshold);
+        assert_eq!(inter, naive_inter, "threshold {threshold}");
+        assert_eq!(union, naive_union, "threshold {threshold}");
+    }
+}
+
 #[test]
 fn mixed_large_dataset() {
     let n    = 1000usize;

src/obicompactvec/src/views.rs

@@ -249,10 +249,35 @@ impl<'a> IntSliceView<'a> {
     }
 
     // ── Distance methods ──────────────────────────────────────────────────────
+    //
+    // `partial_bray_dist`, `partial_euclidean_dist` and
+    // `partial_threshold_jaccard_dist` are the hot path — called
+    // `O(n_cols²)` times per distance matrix over `n`-long columns. Each
+    // computes its formula in one branch-free pass over the raw primary
+    // bytes (vectorises far better than the element-wise `iter().zip()`
+    // used elsewhere, which re-checks the overflow sentinel on every
+    // element), then patches the result for the — normally rare — slots
+    // where that was wrong: any slot that's an overflow entry in `self` or
+    // `other`. Mirrors `format::byte_sum`'s existing "bulk over primary
+    // bytes, then correct for overflow" shape, generalised from a linear
+    // reduction (sum) to arbitrary pairwise formulas via
+    // [`for_each_overflow_pair`](Self::for_each_overflow_pair).
+    // `partial_relfreq_*`/`partial_hellinger_*` are left on the
+    // element-wise path below: they already need a full pass to normalise
+    // by `sum_a`/`sum_b`, so the bulk-primary-bytes trick buys much less there.
 
     pub fn partial_bray_dist(self, other: IntSliceView<'_>) -> u64 {
         assert_eq!(self.n, other.n, "length mismatch");
-        self.iter().zip(other.iter()).map(|(a, b)| a.min(b) as u64).sum()
+        let term = |a: u32, b: u32| (a as u64).min(b as u64);
+        let mut total: u64 = self.primary.iter().zip(other.primary)
+            .map(|(&a, &b)| term(a as u32, b as u32))
+            .sum();
+        self.for_each_overflow_pair(other, |slot, a_true, b_true| {
+            let wrong = term(self.primary[slot] as u32, other.primary[slot] as u32);
+            let right = term(a_true, b_true);
+            total = total + right - wrong;
+        });
+        total
     }
 
     pub fn bray_dist(self, other: IntSliceView<'_>) -> f64 {
@@ -281,9 +306,16 @@ impl<'a> IntSliceView<'a> {
 
     pub fn partial_euclidean_dist(self, other: IntSliceView<'_>) -> f64 {
         assert_eq!(self.n, other.n, "length mismatch");
-        self.iter().zip(other.iter())
+        let term = |a: u32, b: u32| { let d = a as f64 - b as f64; d * d };
-            .map(|(a, b)| { let d = a as f64 - b as f64; d * d })
+        let mut total: f64 = self.primary.iter().zip(other.primary)
-            .sum()
+            .map(|(&a, &b)| term(a as u32, b as u32))
+            .sum();
+        self.for_each_overflow_pair(other, |slot, a_true, b_true| {
+            let wrong = term(self.primary[slot] as u32, other.primary[slot] as u32);
+            let right = term(a_true, b_true);
+            total += right - wrong;
+        });
+        total
     }
 
     pub fn euclidean_dist(self, other: IntSliceView<'_>) -> f64 {
@@ -334,12 +366,23 @@ impl<'a> IntSliceView<'a> {
 
     pub fn partial_threshold_jaccard_dist(self, other: IntSliceView<'_>, threshold: u32) -> (u64, u64) {
         assert_eq!(self.n, other.n, "length mismatch");
-        self.iter().zip(other.iter())
+        let term = |a: u32, b: u32| {
-            .fold((0u64, 0u64), |(inter, uni), (a, b)| {
+            let ap = a >= threshold;
-                let ap = a >= threshold;
+            let bp = b >= threshold;
-                let bp = b >= threshold;
+            ((ap & bp) as u64, (ap | bp) as u64)
-                (inter + (ap & bp) as u64, uni + (ap | bp) as u64)
+        };
-            })
+        let (mut inter, mut uni) = self.primary.iter().zip(other.primary)
+            .fold((0u64, 0u64), |(inter, uni), (&a, &b)| {
+                let (ti, tu) = term(a as u32, b as u32);
+                (inter + ti, uni + tu)
+            });
+        self.for_each_overflow_pair(other, |slot, a_true, b_true| {
+            let (wi, wu) = term(self.primary[slot] as u32, other.primary[slot] as u32);
+            let (ri, ru) = term(a_true, b_true);
+            inter = inter + ri - wi;
+            uni = uni + ru - wu;
+        });
+        (inter, uni)
     }
 
     pub fn threshold_jaccard_dist(self, other: IntSliceView<'_>, threshold: u32) -> f64 {
@@ -350,6 +393,38 @@ impl<'a> IntSliceView<'a> {
     pub fn jaccard_dist(self, other: IntSliceView<'_>) -> f64 {
         self.threshold_jaccard_dist(other, 1)
     }
+
+    /// Walks the union of `self`'s and `other`'s overflow slots — the
+    /// slots where a bulk pass over raw primary bytes computed the wrong
+    /// term because a `255` sentinel isn't the true value — and calls
+    /// `correct(slot, self_true, other_true)` once per slot so the caller
+    /// can undo its (wrong) sentinel-based contribution and add the true
+    /// one. Overflow entries are sorted by slot in both views (on-disk
+    /// invariant, see `format::finalize_pciv`), so their union is walked in
+    /// one linear merge — no allocation, no full second pass over `n`.
+    fn for_each_overflow_pair(self, other: IntSliceView<'_>, mut correct: impl FnMut(usize, u32, u32)) {
+        let mut a_it = self.overflow_entries().peekable();
+        let mut b_it = other.overflow_entries().peekable();
+        loop {
+            let slot = match (a_it.peek(), b_it.peek()) {
+                (None, None) => break,
+                (Some(&(s, _)), None) => s,
+                (None, Some(&(s, _))) => s,
+                (Some(&(sa, _)), Some(&(sb, _))) => sa.min(sb),
+            };
+            let a_true = if a_it.peek().map(|&(s, _)| s) == Some(slot) {
+                a_it.next().unwrap().1
+            } else {
+                self.primary[slot] as u32
+            };
+            let b_true = if b_it.peek().map(|&(s, _)| s) == Some(slot) {
+                b_it.next().unwrap().1
+            } else {
+                other.primary[slot] as u32
+            };
+            correct(slot, a_true, b_true);
+        }
+    }
 }
 
 // ── NonzeroSlotsView / nonzero_triples ──────────────────────────────────────