feat: add batched int group stats API and expand benchmark variants
Introduces a `batch_int_group_stats` API for computing presence counts, sums, minimums, and maximums across sparse and dense matrix representations. The selection layer now utilizes this batched approach to optimize aggregation semantics for boolean and numeric operations. Additionally, reorganizes the benchmarking infrastructure to support querying across presence and count index variants in both dense and sparse formats, including new packing scripts and updated statistics aggregation.
This commit is contained in:
@@ -0,0 +1,163 @@
|
||||
//! Batched, row-major-friendly group statistics for [`PersistentIntMatrix`]
|
||||
//! — the count-matrix counterpart of
|
||||
//! [`crate::bitmatrix::batch_presence_counts`]. Exists for the same two
|
||||
//! reasons: `Sparse` has no on-disk column representation (`col_view` is a
|
||||
//! deliberate `panic!` there), so a real column can only be synthesised by
|
||||
//! scanning every row — exactly what this does directly, via
|
||||
//! [`PersistentSparseCompactIntMatrix::for_each_cell_in_row`] — and even
|
||||
//! for `Columnar`/`Packed`, computing every group's stats together reads a
|
||||
//! column shared by several groups once, not once per group.
|
||||
//!
|
||||
//! Unlike the bit-matrix case, a count matrix's `sum`/`min`/`max` are
|
||||
//! genuine per-value reductions, not derivable from a single presence
|
||||
//! count — so this tracks four running quantities per group instead of
|
||||
//! one: `presence_count` (cells `>= threshold`), `sum`, `min`, `max`.
|
||||
//! `min`/`max` need one more piece of care a dense `col_view` never has to:
|
||||
//! a column absent from a sparse row is implicitly `0`, exactly like a
|
||||
//! `col_view` row that was never written to — so `max` (0 is never a new
|
||||
//! maximum once any real value has been seen) needs no correction, but
|
||||
//! `min` does: if fewer than the full group was actually present at a row,
|
||||
//! at least one member was implicitly `0`, so the true minimum is `0`
|
||||
//! regardless of what was seen among the present ones. Tracked via a fifth,
|
||||
//! internal-only accumulator (`present_count`, distinct from
|
||||
//! `presence_count`: the former is "how many group members exist at this
|
||||
//! row at all", the latter is "how many clear the threshold") and applied
|
||||
//! as a single O(n) correction pass per group once the main scan is done.
|
||||
//!
|
||||
//! `threshold == 0` is `--presence-threshold`'s CLI default, and needs its
|
||||
//! own case: every cell, present or implicitly absent, satisfies `v >= 0`,
|
||||
//! so `presence_count` is trivially `group.len()` at every row — computed
|
||||
//! directly, without paying for the scan's threshold check at all (see
|
||||
//! `PersistentIntMatrix::partial_group_presence_count`'s own `col_view`
|
||||
//! based implementation, which reaches the same result by iterating a
|
||||
//! dense view that materialises the implicit zeros explicitly).
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::io;
|
||||
|
||||
use crate::colgroup::ColGroup;
|
||||
use crate::intmatrix::PersistentIntMatrix;
|
||||
use crate::tempintvec::{TempCompactIntVec, TempCompactIntVecBuilder};
|
||||
|
||||
/// One group's batched statistics — see the module docs for exactly what
|
||||
/// each field means and how `min` differs from a plain per-cell minimum.
|
||||
pub struct IntGroupStats {
|
||||
pub presence_count: TempCompactIntVec,
|
||||
pub sum: TempCompactIntVec,
|
||||
pub min: TempCompactIntVec,
|
||||
pub max: TempCompactIntVec,
|
||||
}
|
||||
|
||||
struct GroupAccum {
|
||||
present_count: TempCompactIntVecBuilder,
|
||||
presence_count: TempCompactIntVecBuilder,
|
||||
sum: TempCompactIntVecBuilder,
|
||||
min: TempCompactIntVecBuilder,
|
||||
max: TempCompactIntVecBuilder,
|
||||
}
|
||||
|
||||
impl GroupAccum {
|
||||
fn new(n: usize) -> io::Result<Self> {
|
||||
Ok(Self {
|
||||
present_count: TempCompactIntVecBuilder::new(n)?,
|
||||
presence_count: TempCompactIntVecBuilder::new(n)?,
|
||||
sum: TempCompactIntVecBuilder::new(n)?,
|
||||
min: TempCompactIntVecBuilder::new(n)?,
|
||||
max: TempCompactIntVecBuilder::new(n)?,
|
||||
})
|
||||
}
|
||||
|
||||
/// One cell (`value` at `slot`) belonging to this group.
|
||||
#[inline]
|
||||
fn touch(&mut self, slot: usize, value: u32, threshold: u32) {
|
||||
let seen_before = self.present_count.get(slot);
|
||||
self.present_count.set(slot, seen_before + 1);
|
||||
if value >= threshold {
|
||||
self.presence_count.set(slot, self.presence_count.get(slot) + 1);
|
||||
}
|
||||
self.sum.set(slot, self.sum.get(slot) + value);
|
||||
if seen_before == 0 || value < self.min.get(slot) {
|
||||
self.min.set(slot, value);
|
||||
}
|
||||
if value > self.max.get(slot) {
|
||||
self.max.set(slot, value);
|
||||
}
|
||||
}
|
||||
|
||||
/// `min` needs a row implicitly missing a group member corrected to
|
||||
/// `0` (see module docs); `presence_count` at `threshold == 0` is
|
||||
/// trivially `group.len()` everywhere, computed directly rather than
|
||||
/// trusting the scan (which never visits a cell that was never
|
||||
/// present, so it would otherwise undercount).
|
||||
fn finish(mut self, n: usize, group_len: u32, threshold: u32) -> io::Result<IntGroupStats> {
|
||||
for slot in 0..n {
|
||||
if self.present_count.get(slot) < group_len {
|
||||
self.min.set(slot, 0);
|
||||
}
|
||||
if threshold == 0 {
|
||||
self.presence_count.set(slot, group_len);
|
||||
}
|
||||
}
|
||||
Ok(IntGroupStats {
|
||||
presence_count: self.presence_count.freeze()?,
|
||||
sum: self.sum.freeze()?,
|
||||
min: self.min.freeze()?,
|
||||
max: self.max.freeze()?,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// See the module docs. `groups` may reference overlapping columns; each
|
||||
/// column is read once regardless of how many groups need it.
|
||||
pub fn batch_int_group_stats(
|
||||
mat: &PersistentIntMatrix,
|
||||
groups: &[ColGroup],
|
||||
threshold: u32,
|
||||
) -> io::Result<Vec<IntGroupStats>> {
|
||||
let n = mat.n();
|
||||
|
||||
let mut col_to_groups: HashMap<usize, Vec<usize>> = HashMap::new();
|
||||
for (gi, g) in groups.iter().enumerate() {
|
||||
for &c in &g.indices {
|
||||
col_to_groups.entry(c).or_default().push(gi);
|
||||
}
|
||||
}
|
||||
|
||||
let mut accums: Vec<GroupAccum> = groups
|
||||
.iter()
|
||||
.map(|_| GroupAccum::new(n))
|
||||
.collect::<io::Result<_>>()?;
|
||||
|
||||
match mat {
|
||||
PersistentIntMatrix::Sparse(m) => {
|
||||
for slot in 0..n {
|
||||
m.for_each_cell_in_row(slot, |col, value| {
|
||||
if let Some(gs) = col_to_groups.get(&col) {
|
||||
for &gi in gs {
|
||||
accums[gi].touch(slot, value, threshold);
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
PersistentIntMatrix::Columnar(_) | PersistentIntMatrix::Packed(_) => {
|
||||
for (&col, gs) in &col_to_groups {
|
||||
let view = mat.col_view(col);
|
||||
for (slot, value) in view.iter().enumerate() {
|
||||
if value == 0 {
|
||||
continue; // matches Sparse's own "never visit an absent cell"
|
||||
}
|
||||
for &gi in gs {
|
||||
accums[gi].touch(slot, value, threshold);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
accums
|
||||
.into_iter()
|
||||
.zip(groups)
|
||||
.map(|(a, g)| a.finish(n, g.indices.len() as u32, threshold))
|
||||
.collect()
|
||||
}
|
||||
@@ -5,6 +5,7 @@ mod colgroup;
|
||||
mod eliasfano;
|
||||
mod fixedintvec;
|
||||
mod format;
|
||||
mod int_group_ops;
|
||||
mod intmatrix;
|
||||
mod layer_meta;
|
||||
mod matrix_builder;
|
||||
@@ -27,6 +28,7 @@ pub use bitvec::{BitIter, PersistentBitVec, PersistentBitVecBuilder};
|
||||
pub use builder::PersistentCompactIntVecBuilder;
|
||||
pub use colgroup::{ColGroup, FilterMask, MatrixGroupOps, eval_filter_mask};
|
||||
pub use eliasfano::{EliasFano, EliasFanoBuilder};
|
||||
pub use int_group_ops::{IntGroupStats, batch_int_group_stats};
|
||||
pub use fixedintvec::{PersistentFixedIntVec, PersistentFixedIntVecBuilder, bit_width_for_range};
|
||||
pub use intmatrix::{
|
||||
PersistentCompactIntMatrixBuilder, PersistentIntMatrix, pack_compact_int_matrix,
|
||||
|
||||
@@ -93,7 +93,7 @@ impl PersistentSparseCompactIntMatrix {
|
||||
/// ([`PersistentSparseBitMatrix::for_each_genome_in_row`]) with this
|
||||
/// row's private slice of whichever value stream it belongs to.
|
||||
#[inline]
|
||||
fn for_each_cell_in_row(&self, slot: usize, mut f: impl FnMut(usize, u32)) {
|
||||
pub(crate) fn for_each_cell_in_row(&self, slot: usize, mut f: impl FnMut(usize, u32)) {
|
||||
match self.support.row_rank(slot) {
|
||||
RowRank::Singleton(rank) => {
|
||||
let v = self.singleton_values.get(rank);
|
||||
|
||||
@@ -0,0 +1,115 @@
|
||||
use tempfile::tempdir;
|
||||
|
||||
use crate::{
|
||||
ColGroup, PersistentCompactIntMatrixBuilder, PersistentIntMatrix,
|
||||
PersistentSparseCompactIntMatrixBuilder, batch_int_group_stats,
|
||||
};
|
||||
|
||||
fn make_columnar(cols: &[&[u32]]) -> (tempfile::TempDir, PersistentIntMatrix) {
|
||||
let n = cols.first().map_or(0, |c| c.len());
|
||||
let dir = tempdir().unwrap();
|
||||
let mut b = PersistentCompactIntMatrixBuilder::new(n, &dir.path().join("counts")).unwrap();
|
||||
for &col in cols {
|
||||
let mut cb = b.add_col().unwrap();
|
||||
for (slot, &v) in col.iter().enumerate() {
|
||||
cb.set(slot, v);
|
||||
}
|
||||
cb.close().unwrap();
|
||||
}
|
||||
b.close().unwrap();
|
||||
let m = PersistentIntMatrix::open(dir.path()).unwrap();
|
||||
(dir, m)
|
||||
}
|
||||
|
||||
/// Row-major input (`cols[c][row]`), sparse-encoded — zero values are
|
||||
/// never pushed, matching the format's own "nonzero cells only" contract.
|
||||
fn make_sparse(cols: &[&[u32]]) -> (tempfile::TempDir, PersistentIntMatrix) {
|
||||
let n_cols = cols.len();
|
||||
let n_rows = cols.first().map_or(0, |c| c.len());
|
||||
let dir = tempdir().unwrap();
|
||||
let sparse_dir = dir.path().join("counts");
|
||||
let mut b = PersistentSparseCompactIntMatrixBuilder::new(n_rows, n_cols, &sparse_dir).unwrap();
|
||||
let mut idx = Vec::new();
|
||||
let mut vals = Vec::new();
|
||||
for row in 0..n_rows {
|
||||
idx.clear();
|
||||
vals.clear();
|
||||
for (c, col) in cols.iter().enumerate() {
|
||||
if col[row] != 0 {
|
||||
idx.push(c as u32);
|
||||
vals.push(col[row]);
|
||||
}
|
||||
}
|
||||
b.push_row(&idx, &vals);
|
||||
}
|
||||
let m = PersistentIntMatrix::Sparse(b.finish().unwrap());
|
||||
(dir, m)
|
||||
}
|
||||
|
||||
// col0=[3,0,5,2], col1=[0,4,5,0], col2=[1,4,0,2]
|
||||
// g1={0,1} (disjoint columns from g2's col2), g2={1,2} (shares col1 with g1)
|
||||
const COL0: [u32; 4] = [3, 0, 5, 2];
|
||||
const COL1: [u32; 4] = [0, 4, 5, 0];
|
||||
const COL2: [u32; 4] = [1, 4, 0, 2];
|
||||
|
||||
#[test]
|
||||
fn batch_int_group_stats_matches_hand_computed_for_columnar_and_sparse() {
|
||||
let cols: [&[u32]; 3] = [&COL0, &COL1, &COL2];
|
||||
let (_dcol, columnar) = make_columnar(&cols);
|
||||
let (_dsparse, sparse) = make_sparse(&cols);
|
||||
|
||||
let groups = [ColGroup::new("g1", vec![0, 1]), ColGroup::new("g2", vec![1, 2])];
|
||||
let threshold = 2;
|
||||
|
||||
// g1={0,1}: presence_count, sum, min, max per row.
|
||||
let g1_presence = [1u32, 1, 2, 1];
|
||||
let g1_sum = [3u64, 4, 10, 2];
|
||||
let g1_min = [0u32, 0, 5, 0]; // row2: both present (3>=thr? irrelevant to min) -> min(5,5)=5
|
||||
let g1_max = [3u32, 4, 5, 2];
|
||||
|
||||
// g2={1,2}: col1 vs col2.
|
||||
let g2_presence = [0u32, 2, 1, 1];
|
||||
let g2_sum = [1u64, 8, 5, 2];
|
||||
let g2_min = [0u32, 4, 0, 0];
|
||||
let g2_max = [1u32, 4, 5, 2];
|
||||
|
||||
for (label, mat) in [("columnar", &columnar), ("sparse", &sparse)] {
|
||||
let stats = batch_int_group_stats(mat, &groups, threshold).unwrap();
|
||||
assert_eq!(stats.len(), 2, "{label}: wrong number of groups");
|
||||
|
||||
for row in 0..4 {
|
||||
assert_eq!(stats[0].presence_count.get(row), g1_presence[row], "{label}: g1 presence row {row}");
|
||||
assert_eq!(stats[0].sum.get(row) as u64, g1_sum[row], "{label}: g1 sum row {row}");
|
||||
assert_eq!(stats[0].min.get(row), g1_min[row], "{label}: g1 min row {row}");
|
||||
assert_eq!(stats[0].max.get(row), g1_max[row], "{label}: g1 max row {row}");
|
||||
|
||||
assert_eq!(stats[1].presence_count.get(row), g2_presence[row], "{label}: g2 presence row {row}");
|
||||
assert_eq!(stats[1].sum.get(row) as u64, g2_sum[row], "{label}: g2 sum row {row}");
|
||||
assert_eq!(stats[1].min.get(row), g2_min[row], "{label}: g2 min row {row}");
|
||||
assert_eq!(stats[1].max.get(row), g2_max[row], "{label}: g2 max row {row}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn threshold_zero_makes_presence_count_the_group_size_everywhere() {
|
||||
let cols: [&[u32]; 3] = [&COL0, &COL1, &COL2];
|
||||
let (_dcol, columnar) = make_columnar(&cols);
|
||||
let (_dsparse, sparse) = make_sparse(&cols);
|
||||
let groups = [ColGroup::new("g1", vec![0, 1, 2])];
|
||||
|
||||
for (label, mat) in [("columnar", &columnar), ("sparse", &sparse)] {
|
||||
let stats = batch_int_group_stats(mat, &groups, 0).unwrap();
|
||||
for row in 0..4 {
|
||||
assert_eq!(stats[0].presence_count.get(row), 3, "{label}: row {row}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_group_list_returns_empty() {
|
||||
let cols: [&[u32]; 1] = [&COL0];
|
||||
let (_d, m) = make_columnar(&cols);
|
||||
let stats = batch_int_group_stats(&m, &[], 1).unwrap();
|
||||
assert!(stats.is_empty());
|
||||
}
|
||||
@@ -3,6 +3,7 @@ mod bitvec;
|
||||
mod colgroup;
|
||||
mod eliasfano;
|
||||
mod fixedintvec;
|
||||
mod int_group_ops;
|
||||
mod intmatrix;
|
||||
mod rankselect;
|
||||
mod sparse;
|
||||
|
||||
@@ -1,18 +1,24 @@
|
||||
//! Per-layer column projection/aggregation — the mechanics
|
||||
//! [`crate::select::Select`] runs once per partition-layer. Generic over
|
||||
//! source content (`Count`/`Presence`) via `obicompactvec::MatrixGroupOps`
|
||||
//! (object-safe: one `&dyn MatrixGroupOps` covers both matrix kinds, so
|
||||
//! this never branches on source content beyond the single `match` that
|
||||
//! opens it) and over destination content via [`DstBuilder`] — no
|
||||
//! duplicated per-op-per-content code path.
|
||||
//! [`crate::select::Select`] runs once per partition-layer. `Presence` and
|
||||
//! `Count` sources each get their own batched, row-major-friendly pass
|
||||
//! (`obicompactvec::batch_presence_counts`/`batch_int_group_stats`)
|
||||
//! computing every output spec's statistics together in one shared scan,
|
||||
//! rather than the generic `MatrixGroupOps` per-group dispatch this used to
|
||||
//! go through (still used by `obikfilter`, whose ad hoc `FilterMask`
|
||||
//! column-index lists don't fit the same "one shared batch" shape) —
|
||||
//! `sum`/`min`/`max` mean genuinely different things for the two content
|
||||
//! kinds (a bit matrix's are cheap derivations of one presence count, a
|
||||
//! count matrix's are real per-value reductions), so each gets its own
|
||||
//! `agg_result_from_*` derivation, not a shared `AggOp` dispatch.
|
||||
|
||||
use std::fs;
|
||||
use std::io;
|
||||
use std::path::Path;
|
||||
|
||||
use obicompactvec::{
|
||||
ColGroup, MatrixBuilder, MatrixGroupOps, PersistentBitMatrix, PersistentIntMatrix, TempBitVec,
|
||||
TempBitVecBuilder, TempCompactIntVec, TempCompactIntVecBuilder, batch_presence_counts,
|
||||
ColGroup, IntGroupStats, MatrixBuilder, PersistentBitMatrix, PersistentIntMatrix, TempBitVec,
|
||||
TempBitVecBuilder, TempCompactIntVec, TempCompactIntVecBuilder, batch_int_group_stats,
|
||||
batch_presence_counts,
|
||||
};
|
||||
use obikindex::layer::{KmerLayer, LayerContent};
|
||||
use obikindex::{KmerIndex, OKIError, OKIResult};
|
||||
@@ -77,22 +83,6 @@ enum AggResult {
|
||||
Int(TempCompactIntVec),
|
||||
}
|
||||
|
||||
fn compute_group(
|
||||
mat: &dyn MatrixGroupOps,
|
||||
spec: &OutputCol,
|
||||
threshold: u32,
|
||||
) -> io::Result<AggResult> {
|
||||
let g = ColGroup::new(spec.label.clone(), spec.indices.clone());
|
||||
Ok(match spec.op {
|
||||
AggOp::Any => AggResult::Bit(mat.partial_group_any(&g, threshold)?),
|
||||
AggOp::All => AggResult::Bit(mat.partial_group_all(&g, threshold)?),
|
||||
AggOp::None => AggResult::Bit(mat.partial_group_none(&g, threshold)?),
|
||||
AggOp::Sum => AggResult::Int(mat.partial_group_sum(&g)?),
|
||||
AggOp::Min => AggResult::Int(mat.partial_group_min(&g)?),
|
||||
AggOp::Max => AggResult::Int(mat.partial_group_max(&g)?),
|
||||
})
|
||||
}
|
||||
|
||||
/// A bit-matrix source's every `AggOp` is a cheap derivation of one shared
|
||||
/// presence count (see [`obicompactvec::batch_presence_counts`]'s own
|
||||
/// docs for why: `sum` = the count itself, `any`/`max` = `count ≥ 1`,
|
||||
@@ -131,6 +121,35 @@ fn agg_result_from_count(op: AggOp, group_len: usize, count: TempCompactIntVec)
|
||||
})
|
||||
}
|
||||
|
||||
/// A count-matrix source's `AggOp` from [`IntGroupStats`] — unlike the
|
||||
/// bit-matrix case, `sum`/`min`/`max` are genuine per-value reductions
|
||||
/// already computed by `batch_int_group_stats`, not further derivations;
|
||||
/// only `any`/`all`/`none` still need a threshold-count comparison here.
|
||||
fn agg_result_from_int_stats(op: AggOp, group_len: usize, stats: IntGroupStats) -> io::Result<AggResult> {
|
||||
let n = stats.presence_count.len();
|
||||
let group_len = group_len as u32;
|
||||
Ok(match op {
|
||||
AggOp::Sum => AggResult::Int(stats.sum),
|
||||
AggOp::Min => AggResult::Int(stats.min),
|
||||
AggOp::Max => AggResult::Int(stats.max),
|
||||
AggOp::Any => {
|
||||
let mut b = TempBitVecBuilder::new(n)?;
|
||||
b.or_where(stats.presence_count.view(), |v| v >= 1);
|
||||
AggResult::Bit(b.freeze()?)
|
||||
}
|
||||
AggOp::All => {
|
||||
let mut b = TempBitVecBuilder::new(n)?;
|
||||
b.or_where(stats.presence_count.view(), |v| v == group_len);
|
||||
AggResult::Bit(b.freeze()?)
|
||||
}
|
||||
AggOp::None => {
|
||||
let mut b = TempBitVecBuilder::new(n)?;
|
||||
b.or_where(stats.presence_count.view(), |v| v == 0);
|
||||
AggResult::Bit(b.freeze()?)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// ── AggResult → MatrixBuilder ─────────────────────────────────────────────────
|
||||
|
||||
/// Add one already-aggregated column to `mb` — the only piece `MatrixBuilder`
|
||||
@@ -233,8 +252,14 @@ pub(crate) fn select_partition(
|
||||
// unchanged, one `col_view`-driven pass per spec.
|
||||
LayerContent::Count => {
|
||||
let mat = PersistentIntMatrix::open(src_layer.dir()).map_err(OKIError::Io)?;
|
||||
for spec in specs {
|
||||
let r = compute_group(&mat, spec, threshold).map_err(OKIError::Io)?;
|
||||
let groups: Vec<ColGroup> = specs
|
||||
.iter()
|
||||
.map(|s| ColGroup::new(s.label.clone(), s.indices.clone()))
|
||||
.collect();
|
||||
let stats = batch_int_group_stats(&mat, &groups, threshold).map_err(OKIError::Io)?;
|
||||
for (spec, stat) in specs.iter().zip(stats) {
|
||||
let r = agg_result_from_int_stats(spec.op, spec.indices.len(), stat)
|
||||
.map_err(OKIError::Io)?;
|
||||
add_result(&mut builder, r).map_err(OKIError::Io)?;
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user