diff --git a/src/obicompactvec/src/bitmatrix/columnar.rs b/src/obicompactvec/src/bitmatrix/columnar.rs index 273cc773..0b05fc1a 100644 --- a/src/obicompactvec/src/bitmatrix/columnar.rs +++ b/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() { - buf[c] = col.get(slot) as u32; - } + fill_row_generic(slot, buf, |c, s| self.col(c).get(s) as u32) } pub(crate) fn count_ones(&self) -> Array1 { - let counts: Vec = (0..self.n_cols()) - .into_par_iter() - .map(|c| self.col(c).count_ones()) - .collect(); - Array1::from_vec(counts) + par_col_reduce(self.n_cols(), |c| self.col(c).count_ones()) } pub(crate) fn partial_jaccard_dist_matrix(&self) -> (Array2, Array2) { diff --git a/src/obicompactvec/src/bitmatrix/mod.rs b/src/obicompactvec/src/bitmatrix/mod.rs index 6435ca2b..40bb06e9 100644 --- a/src/obicompactvec/src/bitmatrix/mod.rs +++ b/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")) diff --git a/src/obicompactvec/src/bitmatrix/packed.rs b/src/obicompactvec/src/bitmatrix/packed.rs index 840abdb5..fcae6466 100644 --- a/src/obicompactvec/src/bitmatrix/packed.rs +++ b/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() { - buf[c] = ((self.mmap[data_off + (slot >> 3)] >> (slot & 7)) & 1) as u32; - } + fill_row_generic(slot, buf, |c, s| self.col_slice(c).get(s) as u32) } #[inline] pub(crate) fn row(&self, slot: usize) -> Box<[bool]> { - (0..self.n_cols).map(|c| { - (self.mmap[self.data_offsets[c] + (slot >> 3)] >> (slot & 7)) & 1 != 0 - }).collect() + row_generic(slot, self.n_cols, |c, s| self.col_slice(c).get(s)) } #[inline] @@ -91,11 +86,7 @@ impl PackedBitMatrix { } pub(crate) fn count_ones(&self) -> Array1 { - Array1::from_vec( - (0..self.n_cols).into_par_iter() - .map(|c| self.col_slice(c).count_ones()) - .collect() - ) + par_col_reduce(self.n_cols, |c| self.col_slice(c).count_ones()) } pub(crate) fn partial_jaccard_dist_matrix(&self) -> (Array2, Array2) { diff --git a/src/obicompactvec/src/bitmatrix/pairwise.rs b/src/obicompactvec/src/bitmatrix/pairwise.rs index dba8079c..4873cd30 100644 --- a/src/obicompactvec/src/bitmatrix/pairwise.rs +++ b/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(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(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(n_cols: usize, f: impl Fn(usize) -> T + Sync + Send) -> Array1 { + 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() } diff --git a/src/obicompactvec/src/builder.rs b/src/obicompactvec/src/builder.rs index 72bdc95b..e682f9b6 100644 --- a/src/obicompactvec/src/builder.rs +++ b/src/obicompactvec/src/builder.rs @@ -192,20 +192,33 @@ impl PersistentCompactIntVecBuilder { } pub fn add(&mut self, other: IntSliceView<'_>) { - let n = self.n; - for s in 0..n { - let sb = self.primary_bytes()[s]; - let ob = other.primary_bytes()[s]; - if sb < 255 && ob < 255 { - let sum = sb as u32 + ob as u32; - if sum < 255 { self.primary_bytes_mut()[s] = sum as u8; } - else { self.set(s, sum); } + // Fast path: both operands fit in a byte — mutate in place, zipped, + // no repeated slice-indexing per element. Slots that already + // overflow (either side) or newly overflow (sum ≥ 255) are + // collected here and settled after the loop, via `get`/`set` + // (which need the overflow map, not touchable while it's borrowed + // through this loop). + let mut newly_overflow: Vec<(usize, u32)> = Vec::new(); + let mut already_overflow: Vec = 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); - let ov = other.get(s); - self.set(s, sv + ov); + already_overflow.push(s); } } + 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; - for s in 0..n { - let sb = self.primary_bytes()[s]; - let ob = other.primary_bytes()[s]; - if sb < 255 { - self.primary_bytes_mut()[s] = if ob < 255 { sb.saturating_sub(ob) } else { 0 }; + // Same shape as `add`: fast byte-level path zipped in place: when + // `self` isn't overflow, either `other` isn't either (plain + // saturating subtract) or `other`'s true value is ≥ 255 — known + // larger than any non-overflow `self` value, so the result + // saturates to 0 without needing `other.get`. Only `self`-overflow + // slots need the overflow map, settled after the loop. + let mut needs_full: Vec = 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); - let ov = if ob < 255 { ob as u32 } else { other.get(s) }; - self.set(s, sv.saturating_sub(ov)); + needs_full.push(s); } } + 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<'_>) { diff --git a/src/obicompactvec/src/intmatrix.rs b/src/obicompactvec/src/intmatrix.rs index deb2c31f..dbb36338 100644 --- a/src/obicompactvec/src/intmatrix.rs +++ b/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 { - let sums: Vec = (0..self.n_cols()) - .into_par_iter() - .map(|c| self.col(c).sum()) - .collect(); - Array1::from_vec(sums) + par_col_reduce(self.n_cols(), |c| self.col(c).sum()) } pub(crate) fn count_nonzero(&self) -> Array1 { - let counts: Vec = (0..self.n_cols()) - .into_par_iter() - .map(|c| self.col(c).count_nonzero()) - .collect(); - Array1::from_vec(counts) + par_col_reduce(self.n_cols(), |c| self.col(c).count_nonzero()) } pub(crate) fn partial_bray_dist_matrix(&self) -> Array2 { @@ -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 { - Array1::from_vec( - (0..self.n_cols).into_par_iter().map(|c| self.col_view(c).sum()).collect() - ) + par_col_reduce(self.n_cols, |c| self.col_view(c).sum()) } pub(crate) fn count_nonzero(&self) -> Array1 { - Array1::from_vec( - (0..self.n_cols).into_par_iter().map(|c| self.col_view(c).count_nonzero()).collect() - ) + par_col_reduce(self.n_cols, |c| self.col_view(c).count_nonzero()) } // Every pair formula below delegates to `IntSliceView` (`views.rs`) — diff --git a/src/obicompactvec/src/tests/mod.rs b/src/obicompactvec/src/tests/mod.rs index e5cd74cd..ee22d26e 100644 --- a/src/obicompactvec/src/tests/mod.rs +++ b/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; diff --git a/src/obicompactvec/src/views.rs b/src/obicompactvec/src/views.rs index dc80dc74..93113c4e 100644 --- a/src/obicompactvec/src/views.rs +++ b/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()) - .map(|(a, b)| { let d = a as f64 - b as f64; d * d }) - .sum() + let term = |a: u32, b: u32| { let d = a as f64 - b as f64; d * d }; + let mut total: f64 = 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 += 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()) - .fold((0u64, 0u64), |(inter, uni), (a, b)| { - let ap = a >= threshold; - let bp = b >= threshold; - (inter + (ap & bp) as u64, uni + (ap | bp) as u64) - }) + let term = |a: u32, b: u32| { + let ap = a >= threshold; + let bp = b >= threshold; + ((ap & bp) as u64, (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 ──────────────────────────────────────