OBIKmers/obikmer
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 24 Wiki Activity

perf: optimize k-mer queries with sparse index and run-based aggregation

Browse Source 
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
2026-07-07 18:56:41 +02:00
parent 040eff140c
commit ae42a061bd
3 changed files with 338 additions and 146 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docmd/architecture/query.md
+67 -39
View File
@@ -16,37 +16,43 @@ Given a set of query sequences, determine for each sequence how many of its k-me
## Algorithm
The query follows the same superkmer-based partitioning strategy used at indexing time. Everything below happens inside `process_chunk` (`query.rs`); there is no separate per-stage function — dereplication, MPHF/matrix lookup, and the Findere window are three phases of one function, operating on a single flat allocation per chunk.
The query follows the same superkmer-based partitioning strategy used at indexing time. Everything below happens inside `process_chunk` (`query.rs`); there is no separate per-stage function, but the internal data flow is staged: k-mer-level dereplication, a two-part MPHF/column-major matrix lookup (`obikpartitionner::query_partition_with`), and a sparse Findere pass, each producing sparse intermediate structures rather than one dense allocation for the whole chunk.
```
for each chunk of sequences (parallel workers via obipipeline, one call to process_chunk):
build QueryBatch (QueryBatch::from_records):
decompose all sequences into superkmers (SuperKmerIter)
deduplicate identical superkmers into
map: HashMap<RoutableSuperKmer, Vec<SKDesc>> ← SKDesc = (seq_idx, kmer_offset)
allocate KmerResults (KmerResults::new): one flat allocation for the whole chunk —
data: Vec<u32> sized total_smers × n_genomes (row-major, zero-initialised)
in_index: Vec<bool> sized total_smers
split the batch's unique superkmers by partition, via minimiser hash
decompose all sequences into superkmers (SuperKmerIter) — construction only,
not the dedup key
deduplicate at k-mer granularity, split by partition in the same pass:
by_partition: Vec<HashMap<CanonicalKmer, Vec<KmerDesc>>> ← KmerDesc = (seq_idx, pos)
allocate SmerIndex (SmerIndex::new): in_index: Vec<bool>, sized total_smers —
NOT multiplied by n_genomes
allocate by_genome: Vec<Vec<(seq_idx, pos, value)>>, one empty Vec per genome —
stays empty (zero cost) for every genome this chunk never matches
for each partition p:
query_partition_with(p, superkmers_routed_to_p, on_hit):
→ load QueryLayer(s) for p
→ for each s-mer of each superkmer: try each layer's MphfLayer::find in turn,
stop at the first hit (a k-mer belongs to at most one layer)
→ on_hit(sk_idx, kmer_idx, row): broadcast `row` into KmerResults, once per
SKDesc referencing this superkmer (all (seq_idx, position) occurrences)
for each sequence, for each genome (inlined sliding window, no separate function):
monotone-deque scan over s-mer positions → win_min[pos][g] = min over the
z-window [pos, pos+z), "not in index" treated as 0
for each sequence, for each output position (0..n_kmers_out):
accumulate confirmed k_user-mer results into acc (kmer_count / kmer_missing /
per-genome totals) and, if --detail, into cov
query_partition_with(p, kmers_for_p, on_event):
stage 1 (MPHF-only): for each unique k-mer, try each layer's MphfLayer::find
in turn, stop at the first hit; bucket confirmed hits by (layer, slot);
emit QueryHit::Found(descs) once per hit k-mer
stage 2 (column-major fetch): for each layer with ≥1 hit, for each genome
column g in 0..layer.n_cols(): scan that layer's bucketed slots, look up
col_value(g, slot); emit QueryHit::Value(descs, g, value) on nonzero
on_event dispatches: Found → SmerIndex::mark_found for every desc;
Value → push (seq_idx, pos, value) into by_genome[g]
for each genome g with ≥1 hit (sparse_findere_for_genome):
sort by_genome[g] by (seq_idx, pos); detect maximal runs of consecutive pos
within one seq_idx; monotone-deque window-minimum scoped to each run →
confirmed_by_genome[g]: Vec<(seq_idx, pos_out, value)>
accumulate genome_totals per sequence from confirmed_by_genome (per genome, direct)
accumulate kmer_count / kmer_missing per (sequence, output position), O(1) each,
using only the confirmed-any bitmap and SmerIndex — independent of n_genomes
if --detail: densify confirmed_by_genome into per-(seq, genome) coverage arrays
emit annotated sequences (emit_batch)
```
Superkmers that appear more than once in the batch (same sequence or across sequences) are deduplicated: each unique `RoutableSuperKmer` is queried once per partition, and the result is broadcast to every `SKDesc` entry that references it. See [Future work](#throughput--parallelism--identified-potential-not-yet-implemented) for a discussion of pushing this dereplication down to k-mer granularity.
Superkmers that appear more than once in the batch (same sequence or across sequences), or different superkmers that happen to share a k-mer (read overlaps, repeats, a SNP splitting an otherwise-identical run), are deduplicated at k-mer granularity: each unique `CanonicalKmer` triggers at most one MPHF lookup and, on hit, one matrix fetch, broadcast to every `KmerDesc` occurrence referencing it.
**Findere requires full-sequence aggregation.** The sliding window is run once per `(sequence, genome)` pair on the complete s-mer result range, after all partitions have contributed to `KmerResults`. Applying it per superkmer would produce false negatives at superkmer boundaries, where the z-window spans two superkmers.
**Findere requires full-sequence aggregation.** The sliding window (now per-run, not per-sequence — see [Findere z-window filter](#findere-z-window-filter)) only ever runs after all partitions have contributed their hits to `by_genome`. Applying it per superkmer would produce false negatives at superkmer boundaries, where the z-window spans two superkmers.
Batches are processed in parallel via `obipipeline` workers; the `--threads` flag controls the number of worker threads.
@@ -56,31 +62,44 @@ Batches are processed in parallel via `obipipeline` workers; the `--threads` fla
For approximate index modes, the index physically stores s-mers of size `s = k_user − z + 1`; `idx.kmer_size()` (bound to `k` in `process_chunk`) is this physically-indexed s-mer size, so decomposing the query at `k` naturally produces s-mer results.
The z-window aggregation is inlined per `(sequence, genome)` in `process_chunk` — a monotone-deque sliding-window minimum, not a separate function operating on an intermediate `Option<Box<[u32]>>` vector:
The z-window aggregation is **sparse**, per genome, implemented in `sparse_findere_for_genome` (`query.rs`) — a run-detection pass followed by a monotone-deque sliding-window minimum scoped to each run, not a dense scan over every s-mer position of every sequence:
```
for each sequence, for each genome g independently:
dq: VecDeque<(pos, value)> // reused across all (sequence, genome) pairs
for i in 0..n_smers:
v = if results.is_in_index(seq, i) { results.val(seq, i, g) } else { 0 }
evict dq entries that left the window [i-z+1, i]
maintain dq monotone non-decreasing (pop back while back.value >= v), push (i, v)
if i + 1 >= z:
pos = i + 1 - z
win_min[pos][g] = dq.front().value // minimum over the z-window
sparse_findere_for_genome(hits, z, presence, threshold):
// hits: raw (seq_idx, pos_smer, value) triples for this genome, as delivered
// by query_partition_with's QueryHit::Value — only ever nonzero entries;
// a position with no hit for this genome simply has no entry at all.
sort hits by (seq_idx, pos_smer)
for each maximal run of consecutive pos_smer values within the same seq_idx:
dq: VecDeque<(run-relative index, value)>
for k, (_, pos, value) in enumerate(run):
maintain dq monotone non-decreasing (pop back while back.value >= value)
push (k, value)
evict dq entries with run-relative index <= k - z
if k + 1 >= z:
win_min = dq.front().value
if win_min > 0:
pos_out = pos + 1 - z
confirmed.push((seq_idx, pos_out, adjust(win_min)))
return confirmed
```
Input length `n_smers`, output length `n_kmers_out = n_smers − z + 1` (k_user-mer positions). The scan is `O(n_smers)` per genome (each position enters and leaves the deque once).
A window can only be confirmed (`win_min > 0`) when all `z` s-mers in it are present *and* nonzero for this genome — which, by construction, can only happen strictly inside one contiguous run of hits (any gap — an absent or zero-valued s-mer — forces `win_min = 0` for every window spanning it, exactly matching the old dense scan's "not in index counts as 0" rule, just never materialising the zero). The deque logic is otherwise identical to the pre-sparsification version; it's scoped to run-relative indices instead of the whole sequence.
Output position `pos` is confirmed for genome `g` iff `win_min[pos][g] > 0` — equivalent to "all z consecutive s-mer values in the window are nonzero for `g`", since a monotone-deque minimum is `0` as soon as any element in the window is `0`.
This runs once per genome that has at least one hit in the chunk (`process_chunk` iterates `by_genome`, one `Vec<(seq_idx, pos_smer, value)>` per genome, built from `QueryHit::Value` during the partition loop — genomes with zero hits in this chunk have an empty `Vec` and cost nothing beyond the iteration itself). Total work is `O(hits log hits)` per genome (the sort) rather than `O(n_smers)` per genome regardless of hit count — a genuine complexity win on top of the memory one, for the common case where most `(chunk, genome)` pairs have no or few hits.
**The value reported per confirmed position is the window minimum, not the leftmost s-mer's raw value.** For presence indexes (0/1 values) this is equivalent to a logical AND either way. For count indexes it is not: the accumulated count for genome `g` at position `pos` is `min(results[pos..pos+z][g])`, the weakest link across the window — not `results[pos][g]`.
Output position `pos_out` is confirmed for genome `g` iff its run produced a nonzero `win_min` — equivalent to "all `z` consecutive s-mer values in the window are nonzero for `g`", same semantics as before.
**`kmer_missing` bookkeeping is separate from the window value** and does look specifically at the leftmost s-mer: a position with no genome confirmed (`win_min[pos][*] == 0` for all `g`) counts as `kmer_missing` iff the leftmost s-mer of that window, `results.is_in_index(seq, pos)`, is absent from the index entirely (see [`kmer_missing` semantics](#kmer_missing-semantics)).
**The value reported per confirmed position is the window minimum, not the leftmost s-mer's raw value** — unchanged from the dense version. For presence indexes (0/1 values) this is equivalent to a logical AND either way. For count indexes it is not: the accumulated count for genome `g` at position `pos_out` is the minimum across the window, the weakest link — not the leftmost s-mer's own count. The presence/count adjustment (`u32::from(win_min >= threshold)` vs. raw `win_min`) is applied once, inside `sparse_findere_for_genome`, rather than later during accumulation.
**Short sequences**: when the s-mer count is less than `z`, `n_kmers_out` is `0` (`process_chunk` skips the sequence: `if out_n == 0 { continue; }`) — no k_user-mer is emitted for that sequence.
**`kmer_missing` bookkeeping is independent of the per-genome sparse structures**, by design (see roadmap point 9): a lightweight dense `SmerIndex` (`in_index: Vec<bool>`, sized `total_smers` — **not** multiplied by `n_genomes`) is populated from `QueryHit::Found` during the partition loop, one entry per hit k-mer regardless of which genome(s) it matched. A position with no genome confirmed counts as `kmer_missing` iff the leftmost s-mer of that window is absent from `SmerIndex` entirely (see [`kmer_missing` semantics](#kmer_missing-semantics)).
**Exact indexes**: `z = 1`, the window degenerates to a single element — the "minimum" is just that element's value, and `pos == i`, i.e. a passthrough.
**Coverage (`--detail`)** is built by re-scanning each genome's confirmed-hit list (already computed, no extra pass over raw data) and densifying into the `[u32; n_kmers_out]` arrays the JSON output format requires — but only when `--detail` is actually requested; the sparse structures cost nothing extra when it isn't.
**Short sequences**: when a sequence's s-mer count is less than `z`, its run(s) — if any hits exist at all — can never reach length `z`, so no window is ever confirmed for it; no k_user-mer is emitted, same outcome as the dense version's `n_kmers_out == 0` early-skip, reached here as a natural consequence rather than a separate check.
**Exact indexes**: `z = 1`, every single-hit "run" of length 1 immediately satisfies `k + 1 >= z`, so every hit is its own confirmed window with `win_min` equal to its own value — a passthrough, as before.
### Effective z at query time
@@ -144,7 +163,7 @@ Coverage reflects confirmed k_user-mers only. The vectors are emitted in the JSO
## `kmer_missing` semantics
`kmer_missing` counts k_user-mer positions where the leftmost s-mer of the window (`results.is_in_index(seq_idx, pos)`, `KmerResults`) is `false` — i.e. absent from the index entirely. K-mers where the z-window fails because a later s-mer is absent or zero (but the leftmost one is present) are not counted as missing — the leftmost s-mer being present is used as proxy for index membership.
`kmer_missing` counts k_user-mer positions where the leftmost s-mer of the window (`smer_index.is_in_index(seq_idx, pos)`, `SmerIndex`) is `false` — i.e. absent from the index entirely. K-mers where the z-window fails because a later s-mer is absent or zero (but the leftmost one is present) are not counted as missing — the leftmost s-mer being present is used as proxy for index membership.
---
@@ -367,6 +386,15 @@ The column-major loop in stage 2 is therefore a **plain sequential loop** for no
- **Validation — this is the correctness-critical phase**: property-test comparing old (dense, pre-phase-3) and new (sparse) implementations on the same randomized input/index fixtures, asserting identical `kmer_count`, `kmer_missing`, `kmer_strict_matches`, and (with `--detail`) `coverage` for every sequence. Keep both implementations compiled side by side (behind a debug-only flag or a temporary parallel code path) only for the duration of this validation; delete the dense path once parity is confirmed — per this project's own convention, superseded code is not kept "just in case."
- **Update `docmd/architecture/query.md` itself**: once this phase lands, the "Findere z-window filter" section (which currently — correctly — describes the dense deque-over-`0..n_smers` scan) needs another pass to describe the sparse run-detection algorithm instead, as already flagged when this phase was discussed.
**Implemented as planned, no deviations discovered this time.** What shipped:
- `KmerResults` removed entirely, replaced by `SmerIndex` (`in_index: Vec<bool>` + `offsets`, unchanged size/purpose, renamed since it's no longer "results" — just the O(1)-per-position "was this k-mer found at all" bookkeeping) and `by_genome: Vec<Vec<(seq_idx, pos, value)>>` (one empty `Vec` per genome until a hit arrives — genomes with zero hits in a chunk cost nothing beyond the outer `Vec`'s own allocation).
- New `sparse_findere_for_genome(hits, z, presence, threshold) -> (Vec<ConfirmedHit>, n_runs, total_run_len)` (`query.rs`): sorts one genome's raw hits by `(seq_idx, pos)`, detects maximal runs of consecutive `pos` within one sequence, runs the same monotone-deque window-minimum as before but scoped to each run (run-relative indices for eviction, absolute `pos` for computing `pos_out`). Presence/count adjustment (`u32::from(win_min >= threshold)` vs. raw) is applied inside this function, once per confirmed hit, rather than later during accumulation.
- `process_chunk` restructured into three passes after the partition loop: (1) run `sparse_findere_for_genome` per genome, collecting `confirmed_by_genome` and run-detection stats; (2) accumulate `genome_totals` directly from `confirmed_by_genome` and mark a `confirmed_any: Vec<bool>` (sized `total_kmers_out`, not `× n_genomes`); (3) a position-only pass (`O(total_kmers_out)`, no genome factor) computing `kmer_count`/`kmer_missing` from `confirmed_any` + `SmerIndex`. `cov` (`--detail`) is populated by re-scanning `confirmed_by_genome` — only when `--detail` is actually set, otherwise skipped entirely.
- Debug log added (`"sparse Findere"`): `n_dense_would_be` (`n_occurrences × n_genomes` — what the deleted dense path would have allocated), `n_sparse_entries` (what's actually retained), `n_runs`/`avg_run_len` (per the plan's ask, to see whether hits mostly fail to form complete windows).
- **Unit tests**: `sparse_findere_matches_dense_reference_on_random_inputs` (`obikmer/src/cmd/tests/query.rs`) — 200 randomized cases (sequence count/length, `z`, presence/count mode, threshold, hit density from sparse to fully-dense) comparing `sparse_findere_for_genome` against `dense_reference_findere`, a faithful reimplementation of the deleted dense algorithm kept only as a test-local correctness oracle (no property-testing crate added — checked first, none was a workspace dependency; a small `std`-only xorshift64 PRNG stands in for one, deterministic and dependency-free). All 200 cases pass.
- **Functional validation performed**: full workspace build + `cargo test --workspace`, zero failures. End-to-end against real indexes: baseline output (no flags) unchanged from pre-phase-5 recorded values on the same fixtures; `--count-missing` correct (`kmer_missing: 0` on a self-match); `--detail` correct — coverage array length matches `kmer_count`, and critically, re-ran the two-genome cross-contamination check from phase 4 with `--detail --count-missing`: `genomeA` reads show coverage sum `106` for `genomeA` and `0` for `genomeB` (and vice versa) — confirms the sparse-to-dense `cov` reconstruction doesn't leak across genomes either, not just the scalar `kmer_strict_matches` path.
- This phase's roadmap item ("update the Findere z-window filter section") — done, see above; the "Algorithm" section's pseudocode was also updated, since it still named `KmerResults`/`SKDesc` from before phases 3–4.
### Phase 6 — Parallel gzip decompression (independent, optional)
Tracked separately in [chunkreader.md](../implementation/chunkreader.md#future-work--parallel-gzip-decompression-in-xopen); parked pending validation of `rapidgzip-rs` on real data. Not a dependency of, or a dependency for, phases 0–5 — `xopen` is shared infrastructure (`obiread`), phase 1 benefits from it but doesn't require it (phase 1 parallelises *across* files; this phase would additionally parallelise *within* one large file).