perf: optimize k-mer queries with sparse index and run-based aggregation
Replaces the dense `KmerResults` matrix with a sparse `SmerIndex` (`Vec<bool>` + offsets) that tracks k-mer presence independently of per-genome counts. Introduces a new aggregation pass, `sparse_findere_for_genome`, which sorts hits, detects contiguous runs, and applies monotone-deque scans to compute sliding-window minimums. This reduces query complexity from O(n_smers) to O(hits log hits), significantly lowering memory overhead and computational cost for low-density queries. Adds a deterministic PRNG and dense reference oracle in tests to validate correctness against randomized inputs without external property-testing crates.
This commit is contained in:
Eric Coissac
+167
-107
@@ -165,21 +165,22 @@ impl QueryBatch {
|
||||
}
|
||||
}
|
||||
|
||||
// ── KmerResults — allocation-free ragged result matrix ───────────────────────
|
||||
// ── SmerIndex — sparse "was this k-mer found at all" bookkeeping ─────────────
|
||||
|
||||
/// Flat storage for per-kmer query results across all sequences in a chunk.
|
||||
///
|
||||
/// Replaces `Vec<Vec<Option<Box<[u32]>>>>` — a single allocation for the whole
|
||||
/// chunk instead of one `Box<[u32]>` per found k-mer.
|
||||
struct KmerResults {
|
||||
data: Vec<u32>, // total_kmers × n_genomes, row-major
|
||||
in_index: Vec<bool>, // total_kmers — true if the kmer was found in the index
|
||||
offsets: Vec<usize>, // offsets[i]..offsets[i+1] = kmer range for sequence i
|
||||
n_genomes: usize,
|
||||
/// Tracks, per (sequence, s-mer position), whether the k-mer was found in the
|
||||
/// index at all — independent of *which* genome(s) matched. Sized
|
||||
/// `total_smers` (one `bool` per s-mer occurrence in the chunk), **not**
|
||||
/// multiplied by `n_genomes`: this is the O(1)-per-position bookkeeping that
|
||||
/// `kmer_missing` needs (the leftmost-s-mer-of-window membership test), kept
|
||||
/// dense because it's already cheap — the `n_genomes`-scaled data lives in
|
||||
/// the sparse per-genome hit lists built alongside it (see `process_chunk`).
|
||||
struct SmerIndex {
|
||||
in_index: Vec<bool>, // total_smers
|
||||
offsets: Vec<usize>, // offsets[i]..offsets[i+1] = s-mer range for sequence i
|
||||
}
|
||||
|
||||
impl KmerResults {
|
||||
fn new(n_kmers_per_seq: &[u32], n_genomes: usize) -> Self {
|
||||
impl SmerIndex {
|
||||
fn new(n_kmers_per_seq: &[u32]) -> Self {
|
||||
let mut offsets = Vec::with_capacity(n_kmers_per_seq.len() + 1);
|
||||
let mut total = 0usize;
|
||||
offsets.push(0);
|
||||
@@ -188,43 +189,96 @@ impl KmerResults {
|
||||
offsets.push(total);
|
||||
}
|
||||
Self {
|
||||
data: vec![0u32; total * n_genomes],
|
||||
in_index: vec![false; total],
|
||||
offsets,
|
||||
n_genomes,
|
||||
}
|
||||
}
|
||||
|
||||
fn n_kmers_for(&self, seq: usize) -> usize {
|
||||
self.offsets[seq + 1] - self.offsets[seq]
|
||||
}
|
||||
|
||||
/// Mark the k-mer at (seq, kmer) as found in the index — independent of
|
||||
/// any particular genome's value. Called once per hit k-mer (stage 1 of
|
||||
/// `query_partition_with`), regardless of how the column-major fetch
|
||||
/// (stage 2) later fills in per-genome values.
|
||||
/// (stage 2) later reports per-genome values.
|
||||
fn mark_found(&mut self, seq: usize, kmer: usize) {
|
||||
let abs = self.offsets[seq] + kmer;
|
||||
self.in_index[abs] = true;
|
||||
}
|
||||
|
||||
/// Set the value for one genome at (seq, kmer). Called once per nonzero
|
||||
/// `(k-mer, genome)` pair delivered by the column-major fetch (stage 2).
|
||||
fn set_one(&mut self, seq: usize, kmer: usize, g: usize, value: u32) {
|
||||
let abs = self.offsets[seq] + kmer;
|
||||
self.data[abs * self.n_genomes + g] = value;
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn is_in_index(&self, seq: usize, kmer: usize) -> bool {
|
||||
self.in_index[self.offsets[seq] + kmer]
|
||||
}
|
||||
}
|
||||
|
||||
/// Value for genome `g` at (seq, kmer); meaningful only when `is_in_index`.
|
||||
#[inline]
|
||||
fn val(&self, seq: usize, kmer: usize, g: usize) -> u32 {
|
||||
self.data[(self.offsets[seq] + kmer) * self.n_genomes + g]
|
||||
// ── Sparse Findere: per-genome run detection + sliding-window minimum ────────
|
||||
|
||||
/// One confirmed z-window: genome `g`'s window ending at k-mer `pos` (the
|
||||
/// *leftmost* s-mer of the window, i.e. the k_user-mer's output position) is
|
||||
/// fully present and nonzero, with window-minimum `value`.
|
||||
type ConfirmedHit = (u32, u32, u32); // (seq_idx, pos_out, value)
|
||||
|
||||
/// Reduce one genome's raw sparse s-mer hits — `(seq_idx, pos_smer, raw_value)`,
|
||||
/// unsorted, exactly as delivered by `QueryHit::Value` — into confirmed
|
||||
/// z-windows, without ever visiting a position that had no hit at all.
|
||||
///
|
||||
/// A z-window is confirmed only when all z s-mers in it are present *and*
|
||||
/// nonzero for this genome (matching the dense sliding-window's semantics,
|
||||
/// where "not in index" or a zero value both contribute 0 to the window
|
||||
/// minimum) — which can only happen inside a maximal run of consecutive
|
||||
/// `pos_smer` values for the same sequence. `hits` is sorted in place by
|
||||
/// `(seq_idx, pos_smer)` to expose those runs; the monotone-deque
|
||||
/// window-minimum then runs per run, on run-relative indices, identical in
|
||||
/// spirit to the dense version's whole-sequence scan.
|
||||
///
|
||||
/// Returns the confirmed hits plus `(n_runs, total_run_len)` for logging —
|
||||
/// a low average run length relative to `z` means most hits fail to form a
|
||||
/// complete window.
|
||||
fn sparse_findere_for_genome(
|
||||
hits: &mut [(u32, u32, u32)],
|
||||
z: usize,
|
||||
presence: bool,
|
||||
threshold: u32,
|
||||
) -> (Vec<ConfirmedHit>, usize, usize) {
|
||||
hits.sort_unstable_by_key(|&(seq, pos, _)| (seq, pos));
|
||||
|
||||
let mut confirmed = Vec::new();
|
||||
let mut n_runs = 0usize;
|
||||
let mut total_run_len = 0usize;
|
||||
let mut dq: VecDeque<(usize, u32)> = VecDeque::new(); // (run-relative index, value)
|
||||
|
||||
let mut i = 0;
|
||||
while i < hits.len() {
|
||||
let seq = hits[i].0;
|
||||
let mut j = i + 1;
|
||||
while j < hits.len() && hits[j].0 == seq && hits[j].1 == hits[j - 1].1 + 1 {
|
||||
j += 1;
|
||||
}
|
||||
let run = &hits[i..j];
|
||||
n_runs += 1;
|
||||
total_run_len += run.len();
|
||||
|
||||
dq.clear();
|
||||
for (k, &(_, pos, val)) in run.iter().enumerate() {
|
||||
while dq.back().map_or(false, |&(_, v)| v >= val) {
|
||||
dq.pop_back();
|
||||
}
|
||||
dq.push_back((k, val));
|
||||
while dq.front().map_or(false, |&(fk, _)| fk + z <= k) {
|
||||
dq.pop_front();
|
||||
}
|
||||
if k + 1 >= z {
|
||||
let win_min = dq.front().unwrap().1;
|
||||
if win_min > 0 {
|
||||
let pos_out = pos + 1 - z as u32;
|
||||
let c = if presence { u32::from(win_min >= threshold) } else { win_min };
|
||||
confirmed.push((seq, pos_out, c));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
i = j;
|
||||
}
|
||||
|
||||
(confirmed, n_runs, total_run_len)
|
||||
}
|
||||
|
||||
// ── Per-sequence accumulator ──────────────────────────────────────────────────
|
||||
@@ -271,8 +325,14 @@ fn process_chunk(
|
||||
let batch = QueryBatch::from_records(records, k, 6, 0.7, n_partitions);
|
||||
let n_seqs = batch.ids.len();
|
||||
|
||||
// Flat result matrix — one allocation for the whole chunk.
|
||||
let mut results = KmerResults::new(&batch.n_kmers, n_genomes);
|
||||
// Sparse bookkeeping for the whole chunk:
|
||||
// - smer_index: O(total_smers) — is this s-mer in the index at all.
|
||||
// - by_genome[g]: raw (seq_idx, pos_smer, value) hits for genome g, only
|
||||
// ever containing nonzero entries (query_partition_with never emits a
|
||||
// QueryHit::Value for a zero value) — empty for every genome this chunk
|
||||
// never matched, which is the common case for unrelated queries.
|
||||
let mut smer_index = SmerIndex::new(&batch.n_kmers);
|
||||
let mut by_genome: Vec<Vec<(u32, u32, u32)>> = (0..n_genomes).map(|_| Vec::new()).collect();
|
||||
|
||||
// Dedup-ratio bookkeeping: occurrences (from batch.n_kmers, computed
|
||||
// before dedup) vs. unique k-mers actually queried (query_stats) — the
|
||||
@@ -296,12 +356,12 @@ fn process_chunk(
|
||||
|event| match event {
|
||||
QueryHit::Found(descs) => {
|
||||
for desc in descs {
|
||||
results.mark_found(desc.seq_idx as usize, desc.pos as usize);
|
||||
smer_index.mark_found(desc.seq_idx as usize, desc.pos as usize);
|
||||
}
|
||||
}
|
||||
QueryHit::Value(descs, g, v) => {
|
||||
for desc in descs {
|
||||
results.set_one(desc.seq_idx as usize, desc.pos as usize, g, v);
|
||||
by_genome[g].push((desc.seq_idx, desc.pos, v));
|
||||
}
|
||||
}
|
||||
},
|
||||
@@ -323,94 +383,94 @@ fn process_chunk(
|
||||
"k-mer dedup + column-major fetch"
|
||||
);
|
||||
|
||||
// Sliding window minimum — one reusable buffer and one deque per batch.
|
||||
// ── Sparse Findere: per-genome run detection + sliding-window minimum ────
|
||||
//
|
||||
// win_min[pos * n_genomes + g] = min count across the z-window [pos, pos+z)
|
||||
// for genome g, where "not in index" counts as 0.
|
||||
//
|
||||
// win_min > 0 ↔ all z consecutive kmers are in the index with count > 0
|
||||
// ↔ Findere confirmation (for z=1 this degenerates to the
|
||||
// simple case with no overhead).
|
||||
//
|
||||
// Works uniformly for count matrices and presence/absence (0/1) matrices.
|
||||
let max_n_kmers = batch.n_kmers.iter().map(|&n| n as usize).max().unwrap_or(0);
|
||||
let mut win_min = vec![0u32; max_n_kmers * n_genomes];
|
||||
|
||||
let mut accs: Vec<SeqAcc> = (0..n_seqs).map(|_| SeqAcc::new(n_genomes)).collect();
|
||||
// Confirmed z-windows, per genome, replace the dense win_min matrix:
|
||||
// total retained memory is O(actual hits), not O(total_smers × n_genomes)
|
||||
// — the whole point of this pass. See sparse_findere_for_genome's doc for
|
||||
// why run detection is equivalent to the dense scan's semantics.
|
||||
let presence = force_presence || !with_counts;
|
||||
let threshold = presence_threshold;
|
||||
let z = effective_z;
|
||||
|
||||
let n_kmers_out: Vec<usize> = batch
|
||||
.n_kmers
|
||||
.iter()
|
||||
.map(|&n| {
|
||||
let n = n as usize;
|
||||
if n >= effective_z { n - effective_z + 1 } else { 0 }
|
||||
if n >= z { n - z + 1 } else { 0 }
|
||||
})
|
||||
.collect();
|
||||
let mut out_offsets = Vec::with_capacity(n_seqs + 1);
|
||||
{
|
||||
let mut total = 0usize;
|
||||
out_offsets.push(0);
|
||||
for &n in &n_kmers_out {
|
||||
total += n;
|
||||
out_offsets.push(total);
|
||||
}
|
||||
}
|
||||
let total_out = *out_offsets.last().unwrap_or(&0);
|
||||
|
||||
let n_dense_would_be = n_occurrences as u64 * n_genomes as u64;
|
||||
let mut n_sparse_entries = 0u64;
|
||||
let mut n_runs_total = 0usize;
|
||||
let mut run_len_total = 0usize;
|
||||
|
||||
let mut confirmed_by_genome: Vec<Vec<ConfirmedHit>> = Vec::with_capacity(n_genomes);
|
||||
for hits in &mut by_genome {
|
||||
n_sparse_entries += hits.len() as u64;
|
||||
let (confirmed, n_runs, run_len) = sparse_findere_for_genome(hits, z, presence, threshold);
|
||||
n_runs_total += n_runs;
|
||||
run_len_total += run_len;
|
||||
confirmed_by_genome.push(confirmed);
|
||||
}
|
||||
|
||||
debug!(
|
||||
n_dense_would_be,
|
||||
n_sparse_entries,
|
||||
n_runs = n_runs_total,
|
||||
avg_run_len = if n_runs_total > 0 { run_len_total as f64 / n_runs_total as f64 } else { 0.0 },
|
||||
z,
|
||||
"sparse Findere"
|
||||
);
|
||||
|
||||
// ── Accumulate: genome totals (per genome, from confirmed hits) ──────────
|
||||
let mut accs: Vec<SeqAcc> = (0..n_seqs).map(|_| SeqAcc::new(n_genomes)).collect();
|
||||
let mut confirmed_any = vec![false; total_out];
|
||||
|
||||
for (g, hits) in confirmed_by_genome.iter().enumerate() {
|
||||
for &(seq_idx, pos_out, c) in hits {
|
||||
let abs_out = out_offsets[seq_idx as usize] + pos_out as usize;
|
||||
confirmed_any[abs_out] = true;
|
||||
accs[seq_idx as usize].genome_totals[g] += c;
|
||||
}
|
||||
}
|
||||
|
||||
// ── Accumulate: kmer_count / kmer_missing (per position, genome-independent) ─
|
||||
for seq_idx in 0..n_seqs {
|
||||
let out_n = n_kmers_out[seq_idx];
|
||||
let acc = &mut accs[seq_idx];
|
||||
for pos in 0..out_n {
|
||||
let abs_out = out_offsets[seq_idx] + pos;
|
||||
if confirmed_any[abs_out] {
|
||||
acc.kmer_count += 1;
|
||||
} else if !smer_index.is_in_index(seq_idx, pos) {
|
||||
acc.kmer_missing += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── Coverage (--detail): densify only when actually requested ────────────
|
||||
let mut cov: Vec<Vec<Vec<u32>>> = if detail {
|
||||
n_kmers_out.iter().map(|&n| vec![vec![0u32; n]; n_genomes]).collect()
|
||||
} else {
|
||||
Vec::new()
|
||||
};
|
||||
|
||||
let presence = force_presence || !with_counts;
|
||||
let threshold = presence_threshold;
|
||||
let z = effective_z;
|
||||
|
||||
// Deque reused across all (seq, genome) pairs.
|
||||
let mut dq: VecDeque<(usize, u32)> = VecDeque::with_capacity(z + 1);
|
||||
|
||||
for seq_idx in 0..n_seqs {
|
||||
let n = results.n_kmers_for(seq_idx);
|
||||
let out_n = n_kmers_out[seq_idx];
|
||||
if out_n == 0 { continue; }
|
||||
|
||||
let mins = &mut win_min[..out_n * n_genomes];
|
||||
mins.fill(0);
|
||||
|
||||
// ── Per-genome sliding window minimum ─────────────────────────────────
|
||||
for g in 0..n_genomes {
|
||||
dq.clear();
|
||||
for i in 0..n {
|
||||
let v_i = if results.is_in_index(seq_idx, i) {
|
||||
results.val(seq_idx, i, g)
|
||||
} else {
|
||||
0
|
||||
};
|
||||
// Evict elements that have left the window.
|
||||
while dq.front().map_or(false, |&(f, _)| f + z <= i) {
|
||||
dq.pop_front();
|
||||
}
|
||||
// Maintain monotone non-decreasing back→front for minimum at front.
|
||||
while dq.back().map_or(false, |&(_, v)| v >= v_i) {
|
||||
dq.pop_back();
|
||||
}
|
||||
dq.push_back((i, v_i));
|
||||
// Window [pos, pos+z) is complete when i = pos + z - 1.
|
||||
if i + 1 >= z {
|
||||
let pos = i + 1 - z;
|
||||
mins[pos * n_genomes + g] = dq.front().unwrap().1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── Accumulate ────────────────────────────────────────────────────────
|
||||
let acc = &mut accs[seq_idx];
|
||||
for pos in 0..out_n {
|
||||
let any = (0..n_genomes).any(|g| mins[pos * n_genomes + g] > 0);
|
||||
if !any {
|
||||
if !results.is_in_index(seq_idx, pos) {
|
||||
acc.kmer_missing += 1;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
acc.kmer_count += 1;
|
||||
for g in 0..n_genomes {
|
||||
let v = mins[pos * n_genomes + g];
|
||||
if v == 0 { continue; }
|
||||
let c = if presence { u32::from(v >= threshold) } else { v };
|
||||
acc.genome_totals[g] += c;
|
||||
if detail { cov[seq_idx][g][pos] += c; }
|
||||
if detail {
|
||||
for (g, hits) in confirmed_by_genome.iter().enumerate() {
|
||||
for &(seq_idx, pos_out, c) in hits {
|
||||
cov[seq_idx as usize][g][pos_out as usize] += c;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -122,3 +122,107 @@ fn partition_routing_is_a_pure_function_of_the_kmer() {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── sparse_findere_for_genome vs. a dense reference implementation ──────────
|
||||
//
|
||||
// No property-testing crate (proptest/quickcheck) is a workspace dependency
|
||||
// (checked before writing this — not adding one for a single test module,
|
||||
// per this project's dependency-approval rule). A tiny deterministic xorshift
|
||||
// PRNG, std-only, stands in for one.
|
||||
|
||||
/// Faithful reimplementation of the pre-phase-5 dense sliding-window scan —
|
||||
/// the algorithm `sparse_findere_for_genome` replaced — used here only as a
|
||||
/// correctness oracle, not in production code. Operates on one genome's
|
||||
/// hits across possibly many sequences, exactly like the sparse version.
|
||||
fn dense_reference_findere(
|
||||
hits: &[(u32, u32, u32)],
|
||||
seq_lens: &[usize],
|
||||
z: usize,
|
||||
presence: bool,
|
||||
threshold: u32,
|
||||
) -> Vec<(u32, u32, u32)> {
|
||||
let mut by_seq: Vec<Vec<u32>> = seq_lens.iter().map(|&n| vec![0u32; n]).collect();
|
||||
for &(seq, pos, val) in hits {
|
||||
by_seq[seq as usize][pos as usize] = val;
|
||||
}
|
||||
|
||||
let mut confirmed = Vec::new();
|
||||
for (seq_idx, values) in by_seq.iter().enumerate() {
|
||||
let n = values.len();
|
||||
let mut dq: std::collections::VecDeque<(usize, u32)> = std::collections::VecDeque::new();
|
||||
for i in 0..n {
|
||||
let v_i = values[i];
|
||||
while dq.front().map_or(false, |&(f, _)| f + z <= i) {
|
||||
dq.pop_front();
|
||||
}
|
||||
while dq.back().map_or(false, |&(_, v)| v >= v_i) {
|
||||
dq.pop_back();
|
||||
}
|
||||
dq.push_back((i, v_i));
|
||||
if i + 1 >= z {
|
||||
let win_min = dq.front().unwrap().1;
|
||||
if win_min > 0 {
|
||||
let pos_out = (i + 1 - z) as u32;
|
||||
let c = if presence { u32::from(win_min >= threshold) } else { win_min };
|
||||
confirmed.push((seq_idx as u32, pos_out, c));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
confirmed
|
||||
}
|
||||
|
||||
/// Minimal std-only xorshift64 PRNG — deterministic, seedable, no dependency.
|
||||
struct Xorshift64(u64);
|
||||
impl Xorshift64 {
|
||||
fn next(&mut self) -> u64 {
|
||||
self.0 ^= self.0 << 13;
|
||||
self.0 ^= self.0 >> 7;
|
||||
self.0 ^= self.0 << 17;
|
||||
self.0
|
||||
}
|
||||
fn range(&mut self, n: u32) -> u32 {
|
||||
(self.next() % n as u64) as u32
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sparse_findere_matches_dense_reference_on_random_inputs() {
|
||||
let mut rng = Xorshift64(0x5eed_5eed_5eed_5eedu64);
|
||||
|
||||
for case in 0..200 {
|
||||
let n_seqs = 1 + rng.range(4) as usize;
|
||||
let seq_lens: Vec<usize> = (0..n_seqs).map(|_| 1 + rng.range(30) as usize).collect();
|
||||
let z = 1 + rng.range(4) as usize;
|
||||
let presence = rng.range(2) == 0;
|
||||
let threshold = 1 + rng.range(3);
|
||||
|
||||
// Sparse density varies across cases, including edge cases (empty,
|
||||
// fully dense) — deliberately not uniform, to stress both few-hits
|
||||
// and many-overlapping-runs scenarios.
|
||||
let density = rng.range(101);
|
||||
let mut hits: Vec<(u32, u32, u32)> = Vec::new();
|
||||
for (seq_idx, &len) in seq_lens.iter().enumerate() {
|
||||
for pos in 0..len {
|
||||
if rng.range(100) < density {
|
||||
let val = 1 + rng.range(5); // never 0 — matches QueryHit::Value's invariant
|
||||
hits.push((seq_idx as u32, pos as u32, val));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let mut sparse_input = hits.clone();
|
||||
let (mut sparse_result, _, _) =
|
||||
sparse_findere_for_genome(&mut sparse_input, z, presence, threshold);
|
||||
let mut dense_result = dense_reference_findere(&hits, &seq_lens, z, presence, threshold);
|
||||
|
||||
sparse_result.sort_unstable();
|
||||
dense_result.sort_unstable();
|
||||
|
||||
assert_eq!(
|
||||
sparse_result, dense_result,
|
||||
"case {case}: n_seqs={n_seqs} seq_lens={seq_lens:?} z={z} presence={presence} \
|
||||
threshold={threshold} density={density} hits={hits:?}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user