Push smluomvxpptv #59
+218
-36
@@ -16,27 +16,37 @@ 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.
|
||||
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.
|
||||
|
||||
```
|
||||
for each chunk of sequences (parallel workers via obipipeline):
|
||||
build QueryBatch: decompose all sequences into s-mers via superkmers, deduplicate
|
||||
allocate seq_results[seq_idx][smer_pos] = None ← per-sequence s-mer result vectors
|
||||
split superkmers by partition via minimiser hash
|
||||
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
|
||||
for each partition p:
|
||||
query_partition(p, superkmers_routed_to_p)
|
||||
query_partition_with(p, superkmers_routed_to_p, on_hit):
|
||||
→ load QueryLayer(s) for p
|
||||
→ for each s-mer in each superkmer: MphfLayer::find(smer)
|
||||
fill seq_results[seq_idx][kmer_offset + j] from partition results
|
||||
for each sequence:
|
||||
apply_findere(seq_results[seq_idx], effective_z) ← per full sequence
|
||||
accumulate confirmed k-mer results into acc and cov
|
||||
emit annotated sequences
|
||||
→ 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
|
||||
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.
|
||||
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.
|
||||
|
||||
**Findere requires full-sequence aggregation.** `apply_findere` is applied once per sequence on the complete s-mer result vector, after all partitions have contributed. 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 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.
|
||||
|
||||
Batches are processed in parallel via `obipipeline` workers; the `--threads` flag controls the number of worker threads.
|
||||
|
||||
@@ -44,25 +54,33 @@ Batches are processed in parallel via `obipipeline` workers; the `--threads` fla
|
||||
|
||||
## Findere z-window filter
|
||||
|
||||
For approximate index modes, the index physically stores s-mers of size `s = k_user − z + 1`. At query time, `set_k(s)` is in effect, so queries naturally produce s-mer results. `apply_findere` then aggregates z consecutive s-mer results into one k_user-mer answer:
|
||||
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.
|
||||
|
||||
```rust
|
||||
fn apply_findere(
|
||||
results: &[Option<Box<[u32]>>], // N s-mer results
|
||||
z: usize,
|
||||
n_genomes: usize,
|
||||
) -> Vec<Option<Box<[u32]>>> // N − z + 1 k_user-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:
|
||||
|
||||
```
|
||||
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
|
||||
```
|
||||
|
||||
Input length N (s-mers), output length N − z + 1 (k_user-mers).
|
||||
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).
|
||||
|
||||
For each genome g independently, a sliding window of size z scans the input. Output position i is confirmed for genome g iff all z values `results[i..i+z][g]` are nonzero (`None` counts as zero for all genomes). The scan is O(n) per genome.
|
||||
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`.
|
||||
|
||||
Output values come from `results[i]` (leftmost s-mer of each window); genomes not confirmed are zeroed. If all genomes are zero, the position is returned as `None`.
|
||||
**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]`.
|
||||
|
||||
**Short sequences**: when the s-mer count is less than z, no complete window can form — `apply_findere` returns an empty vector. K-mers from sequences shorter than k_user are not emitted.
|
||||
**`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)).
|
||||
|
||||
**Exact indexes**: `z = 1`, `apply_findere` is a passthrough (output length = input length).
|
||||
**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.
|
||||
|
||||
**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.
|
||||
|
||||
### Effective z at query time
|
||||
|
||||
@@ -85,14 +103,17 @@ The `-z` CLI option overrides the index metadata value. A higher z increases str
|
||||
|
||||
### `QueryLayer` variant selection
|
||||
|
||||
`QueryLayer::open` in `query_layer.rs` selects the data matrix to pair with `MphfLayer`:
|
||||
`QueryLayer::open` (`obikpartitionner/src/query_layer.rs:28-45`) only ever returns two variants — `Presence` or `Count`, checked in this order:
|
||||
|
||||
| Condition | Variant | Data returned per k-mer |
|
||||
|---|---|---|
|
||||
| `with_counts=true` and `counts/` exists | `Count` | raw count per genome |
|
||||
| `presence/` exists | `Presence` | 0/1 per genome (bit matrix) |
|
||||
| only `counts/` exists | `Count` | counts used as-is |
|
||||
| neither exists | `SetOnly` | 1 for every genome |
|
||||
| Order | Condition | Variant | Data returned per k-mer |
|
||||
|---|---|---|---|
|
||||
| 1 | `with_counts=true` and `counts/` exists | `Count` | raw count per genome |
|
||||
| 2 | (else) `presence/` exists, or `counts/` doesn't exist at all | `Presence` | see below |
|
||||
| 3 | (else — `counts/` exists, `presence/` doesn't, `with_counts=false`) | `Count` | counts used as-is |
|
||||
|
||||
There is no `QueryLayer::SetOnly` variant. The "no on-disk matrix at all" case is handled one level down: `Presence` wraps `PersistentBitMatrix`, whose own `open()` (`obicompactvec/src/bitmatrix.rs:260-288`) auto-detects among **three** internal representations — `Packed` (`presence/matrix.pbmx`), `Columnar` (`presence/meta.json`), or `Implicit { n_rows, n_cols }` when neither file exists (built from `layer_meta.json`, `fill_row` returning all-`1`s without touching disk). This is where "1 for every genome" actually happens — not at the `QueryLayer` level.
|
||||
|
||||
**Worth double-checking, not confirmed as a bug**: `PersistentBitMatrix::open`'s `Implicit` branch constructs `Implicit { n_rows: meta.n, n_cols: 1 }` — `n_cols` is hardcoded to `1`, not to the layer's actual `n_genomes`. `fill_row` for `Implicit` only writes `buf[..1]`, leaving the rest of a longer `n_genomes`-sized buffer untouched (zeroed by the caller beforehand). If this path is ever reached for a layer covering more than one genome, only genome index 0 would read as present. Whether that's reachable in practice (layers might always be single-genome when they fall back to `Implicit`) wasn't verified here — flagging for follow-up, not fixing.
|
||||
|
||||
---
|
||||
|
||||
@@ -123,7 +144,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 first s-mer (`seq_results[seq_idx][pos]`) is `None` — i.e. absent from the index entirely. K-mers where the z-window fails because a later s-mer is absent or zero are not counted as missing (the first 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 (`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.
|
||||
|
||||
---
|
||||
|
||||
@@ -152,7 +173,7 @@ Genome keys follow the iteration order of `meta.genomes`.
|
||||
| Key | Type | Condition | Semantics |
|
||||
|---|---|---|---|
|
||||
| `kmer_count` | int | always | k-mers confirmed (post-Findere) with at least one genome match |
|
||||
| `kmer_missing` | int | `--count-missing` | k-mers absent from the index entirely (pre-Findere None) |
|
||||
| `kmer_missing` | int | `--count-missing` | k-mers absent from the index entirely (leftmost s-mer of the window not found) |
|
||||
| `kmer_strict_matches` | object | always | per-genome accumulated value, non-zero entries only (label → count or 0/1) |
|
||||
| `coverage` | object | `--detail` | per-genome array of per-position contributions (label → [u32]) |
|
||||
|
||||
@@ -165,7 +186,7 @@ Genome keys follow the iteration order of `meta.genomes`.
|
||||
```
|
||||
obikmer query <index> [--detail] [--mismatch] [--count-missing]
|
||||
[--force-presence] [--presence-threshold <n>]
|
||||
[-z <z>] [-T <threads>]
|
||||
[-z <z>] [-T <threads>] [--chunk-size <MiB>]
|
||||
<query.fa> [<query2.fa> ...]
|
||||
```
|
||||
|
||||
@@ -177,6 +198,7 @@ obikmer query <index> [--detail] [--mismatch] [--count-missing]
|
||||
| `--force-presence` | off | Report 0/1 per genome regardless of index counts |
|
||||
| `--presence-threshold` | 1 | Minimum count to declare genome present |
|
||||
| `-T` / `--threads` | all CPUs | Worker threads |
|
||||
| `--chunk-size` | auto (from available RAM and thread count) | I/O chunk size in MiB — see [Future work, point 3](#throughput--parallelism--identified-potential-not-yet-implemented) for why the auto-sizing formula currently under-estimates memory on indexes with many genomes |
|
||||
|
||||
`--mismatch` is accepted but currently ignored with a warning on stderr.
|
||||
|
||||
@@ -187,3 +209,163 @@ obikmer query <index> [--detail] [--mismatch] [--count-missing]
|
||||
- **`--mismatch`**: 1-mismatch approximate matching — generate `3·k` single-substitution variants per k-mer, look each up independently.
|
||||
- **Read classification** (`--classify`): assign each read to the genome with the highest match score.
|
||||
- **Whitelist / blacklist filtering**: threshold-based accept/reject on per-genome match scores.
|
||||
|
||||
### Throughput & parallelism — identified potential (not yet implemented)
|
||||
|
||||
Observed on a 192-core (8×24 NUMA) machine: `query` uses ~10 cores or fewer, and the default chunk size gets the process OOM-killed. Root causes and candidate fixes, in dependency order:
|
||||
|
||||
**1. Single-threaded I/O source (main core-utilization bottleneck).**
|
||||
`run()` builds `all_chunks` via `paths.into_iter().flat_map(read_sequence_chunks_sized(...))` and passes it directly as the `input` iterator to `pipe.apply()`. In `obipipeline::Pipe::apply` (`scheduler.rs`), `input.next()` is called exclusively from the dedicated source thread — so file opening, decompression, and FASTA/FASTQ chunk-boundary parsing for *all* input files run serially in one thread, regardless of `--threads`. Compare with `steps::scatter` (used by `index`) and `cmd/superkmer.rs`: there, file opening + streaming is itself a `Flat` pipeline stage (`||?`), executed across the `n_workers` pool, with `obipipeline::throttle(paths, max_open)` bounding concurrently-open files in the source thread. That pattern parallelises I/O across files (and NUMA nodes); `query.rs` cannot.
|
||||
Fix direction: restructure `query`'s pipe with an initial `Flat` stage analogous to `scatter`'s, opening/chunking files across workers instead of in `flat_map`.
|
||||
|
||||
**2. Gzip decompression is inherently single-threaded per file.**
|
||||
`niffler`/`flate2` (used by `xopen`) do standard DEFLATE, which has no parallel-decodable structure for an arbitrary stream. Fix (1) parallelises *across* files but not *within* one large gzip file. Parking a possible fix (`rapidgzip-rs`) is tracked in [chunkreader.md](../implementation/chunkreader.md#future-work--parallel-gzip-decompression-in-xopen).
|
||||
|
||||
**3. Chunk-size memory formula ignores `n_genomes`.**
|
||||
`chunk_bytes = available_memory_bytes() / (n_workers * 16)` (`query.rs:407-414`) assumes a fixed ~8–16× overhead per raw input byte. But `KmerResults::new` (`query.rs:165-179`) allocates `data: Vec<u32>` sized `total_kmers_in_chunk × n_genomes` — dense, **for every k-mer position in the chunk, hit or not** — plus `win_min` and (with `--detail`) `cov`, same scaling. Real per-chunk memory is `O(n_genomes)`, not constant; the formula doesn't know `n_genomes` at all. This is the direct cause of the OOM kill on indexes with many reference genomes.
|
||||
|
||||
**4. MPHF lookup and matrix-row fetch are fused, not staged.**
|
||||
`QueryLayer::find_into` (`obikpartitionner/src/query_layer.rs:48-67`) does the MPHF `find` *and* the `fill_row` matrix read in one call per k-mer, inside a single-threaded loop (`query_partition_with`). There is no separation between "is this k-mer indexed" (cheap, `O(1)`, independent of `n_genomes`) and "what are its per-genome values" (the expensive, `n_genomes`-scaling part).
|
||||
|
||||
**5. Dereplication should happen at k-mer granularity, directly — not via an intermediate superkmer-level dedup.**
|
||||
`QueryBatch::from_records` currently dereplicates at the *superkmer* level (`HashMap<RoutableSuperKmer, Vec<SKDesc>>`, `query.rs:112`). This misses redundancy between k-mers shared by *different* superkmers (read overlaps, repeats, a SNP splitting an otherwise-identical run). Superkmer *construction* (`SuperKmerIter`) stays mandatory — it is the mechanism that computes minimizers/partition routing, not an optional dedup layer — but the dedup structure built on top of it should key directly on `CanonicalKmer`, in the same pass: `HashMap<CanonicalKmer, Vec<(seq_idx, pos)>>`. This also means the MPHF `find` itself runs once per **distinct** k-mer instead of once per occurrence — a win independent of the matrix-fetch cost below.
|
||||
|
||||
**6. Stage 1 output: bucket confirmed hits by layer, keyed by MPHF slot.**
|
||||
For each unique canonical k-mer, MPHF lookup across a partition's layers stops at the first match (`query_partition_with:105-111`) — a k-mer belongs to at most one layer. So stage 1's output can be reshaped directly into:
|
||||
```
|
||||
HashMap<layer_idx, HashMap<slot, Vec<(seq_idx, pos)>>>
|
||||
```
|
||||
replacing the `CanonicalKmer` key by the resolved `slot` (compact integer, and exactly what stage 2 needs to address the matrix). K-mers matching no layer simply have no entry here (they still count toward `in_index`/`kmer_missing` bookkeeping, which stays `O(1)` per position, independent of `n_genomes`).
|
||||
|
||||
**7. Partition-level parallelism is currently absent — and a NUMA-aware mechanism for exactly this already exists, unused, in `obikindex`.**
|
||||
`process_chunk`'s partition loop (`query.rs:250-278`, `for (part_idx, part_sks) in by_part.iter().enumerate()`) processes every partition of a chunk sequentially on the single worker thread that owns that chunk. This is a parallelism axis on its own, independent of the column question below.
|
||||
More importantly: `docmd/architecture/numa_partition_runner.md` and `numa_worker_pools.md` document `PartitionRunner` (`obikindex/src/numa.rs`), **already implemented** and already used by `merge.rs`, `index.rs` (`build_layers`), `select.rs`, `reindex.rs`, `rebuild.rs` — one controller thread per NUMA node, a Rayon pool pinned to that node's CPUs (`hwlocality`, `numa` feature, default-on in `obikindex/Cargo.toml`), adaptive worker activation driven by *both* a CPU-efficiency signal and an I/O-throughput signal (`CpuSample`/`IoSample`, `/proc/self/io` on Linux). It exists precisely because a naive `into_par_iter()` on the global Rayon pool measurably degrades ×60 on this codebase's own 192-core/8-NUMA reference machine (`numa_worker_pools.md`, § Problem) once workers contend for cross-socket memory bandwidth on shared mmap'd/hashed structures — exactly the shape of the matrix-column scan in point 8 below.
|
||||
`obikmer` already depends on `obikindex` (`obikmer/Cargo.toml`, for `KmerIndex`), so `PartitionRunner` is directly reachable from `cmd/query.rs` — no new dependency. Both the partition-level loop and (see point 8) the genome-column scan should be driven through it rather than through ad-hoc `rayon::into_par_iter()`, to avoid reproducing the already-measured-and-fixed contention problem. Also relevant: the "CPU-only signal stalls on I/O-bound stages" issue documented for `pack_matrices` (mmap-heavy, page-fault-bound) applies just as much to a column-major mmap scan over persistent matrices — reuse the existing dual CPU/IO activation signal rather than re-deriving one.
|
||||
|
||||
**8. Stage 2: column-major matrix fetch, parallel across genome columns — via `PartitionRunner`, not naive `rayon`.**
|
||||
Both persistent matrix formats are column-oriented on disk: `ColumnarCompactIntMatrix`/`ColumnarBitMatrix` (`obicompactvec/src/{intmatrix,bitmatrix}.rs`) mmap one file per genome column; `PackedCompactIntMatrix`/`PackedBitMatrix` mmap one region-offset per column in a single file. `fill_row(slot, buf)` as used today (`query.rs:262-272` via `on_hit`) reads **one slot across all `n_genomes` columns** per hit — the worst possible access pattern for this layout (up to `n_genomes` scattered mmap regions touched per single k-mer).
|
||||
Better: for each layer, walk the matrix **column by column** (genome by genome): for each genome, scan the `slot` keys collected in step 6 for that layer and call `col.get(slot)`, keeping only nonzero results, and broadcast to the associated `(seq_idx, pos)` list. Total `get()` calls are unchanged (`n_hits × n_genomes` in the worst case) — the win is locality (sequential access within one mmap'd column at a time, not scattered across all columns per hit), not fewer operations.
|
||||
Columns are independent (read-only, disjoint mmap regions) → embarrassingly parallel across genomes, *but* — per point 7 — `obicompactvec`'s existing `into_par_iter()` over `0..n_cols` (`sum()`, `count_nonzero()`, pairwise distance matrices) is the **naive, unpinned** pattern the rest of the codebase is actively migrating away from, not a model to copy here. Route this through `PartitionRunner` (or the same NUMA-pool machinery) instead. Two things to settle when this is designed: how the partition axis (point 7), the column axis, and the existing chunk-level `n_workers` `obipipeline` pool compose without oversubscribing the machine (three different concurrency mechanisms — raw-thread pipe workers, `PartitionRunner`'s pinned Rayon pools, and whatever drives the column scan — need a single reconciled thread budget, not three independent ones); and the threshold below which per-column dispatch overhead outweighs the gain (small `n_genomes` or small per-layer hit counts) — to be measured, not assumed.
|
||||
|
||||
**9. Sparse per-genome representation, fed directly to Findere.**
|
||||
Stage 2's output should be `HashMap<genome_idx, Vec<(seq_idx, position, count)>>`, **sorted by `(seq_idx, position)`** once collected, instead of a dense `KmerResults`-style matrix — the key must carry `seq_idx`, not just `genome_idx`, because a chunk batches many sequences and `position` is only meaningful within one; a plain `Vec<(position, count)>` per genome would silently mix positions from different sequences and corrupt the sliding-window scan. This bounds retained memory by actual nonzero hits on both axes (position sparsity from non-matching k-mers, genome sparsity from a matched k-mer typically belonging to only a handful of genomes out of possibly many). The Findere sliding-window (`process_chunk`, the `win_min`/deque loop) would need reworking to run per `(sequence, genome)` over its sparse, sorted `(position, count)` list — detect runs of ≥`z` consecutive positions, window-min within each run — instead of today's dense `O(total_kmers × n_genomes)` scan. This is also a genuine complexity win (`O(hits log hits)` per genome vs. dense scan), not just memory.
|
||||
**Not covered by this sparsification**: `--detail`'s `cov` accumulator (`query.rs:304-308`) has the identical `n_genomes`-dense scaling problem and wasn't folded into points above. It doesn't need to be retained densely throughout processing, though — only the final JSON serialization (`emit_batch`) requires a dense `[u32]` per `(seq, genome)`, and only for the sequences actually being output with `--detail`. Densification can stay a late, output-time-only step, reconstructed from the sparse per-genome lists.
|
||||
|
||||
**Secondary patterns available from `scatter.rs`/`superkmer.rs`, not yet in `query.rs`:**
|
||||
- `throttle()` + `CommonArgs::effective_max_open()` to bound concurrently-open input files (query.rs defines its own `QueryArgs`, doesn't reuse this).
|
||||
- Progress bar with EMA throughput + live active-worker gauges (`obisys::spinner`, `flat_active`/`transform_active` counters) — diagnostic value for locating the bottleneck.
|
||||
- `obisys::Reporter`/`Stage::start`/`stop` timing per phase (used by `index`, `filter`; absent from `query`).
|
||||
|
||||
None of this is implemented yet — parked here as a coherent roadmap while the design is discussed further. Suggested dependency order: (1) I/O parallelism → (3) genome-aware chunk sizing → (4)–(9) staged/k-mer-deduped/NUMA-aware-partition-and-column-major/sparse query engine (larger refactor, biggest structural payoff — reuses `PartitionRunner` rather than inventing a new parallelism mechanism) → (2) parallel gzip (separate, orthogonal, tracked in chunkreader.md) → secondary diagnostics patterns.
|
||||
|
||||
---
|
||||
|
||||
## Implementation plan
|
||||
|
||||
Concrete, phased translation of the roadmap above. Phases 0–2 are small, independent, low-risk, and each individually testable against current `query` output — land them first, in order, and measure on the reference 192-core/8-NUMA machine before deciding whether phases 3–5 (the staged/sparse engine, the larger structural payoff) are still worth their cost. Phases 3–5 are one coordinated change spanning `obikmer`, `obikpartitionner`, and `obicompactvec` — they should not be split across releases mid-way, because the intermediate state (e.g. k-mer-level dedup feeding the old dense `KmerResults`) has no correctness or performance benefit on its own. Phase 6 is unrelated to phases 0–5 and can happen any time, independently, if `rapidgzip-rs` is validated (see [chunkreader.md](../implementation/chunkreader.md#future-work--parallel-gzip-decompression-in-xopen)).
|
||||
|
||||
Instrumentation is deliberately sequenced *before* the I/O fix (reordering the roadmap's own listed order), because every later phase's justification rests on a measurement ("to be measured, not assumed" appears throughout the roadmap above) — without it, phases 3–5 would be undertaken on faith.
|
||||
|
||||
Performance measurement on the reference 192-core/8-NUMA machine is done by the project owner, not from this development environment (macOS, 16 cores — `PartitionRunner`'s NUMA pinning is Linux-only, so even phase 4's mechanism can't be functionally exercised for its actual purpose here). Each phase below is therefore written to be *self-measuring*: the debug-level logging it adds must be enough, on its own, to judge whether that phase's algorithmic choice paid off from a cluster run's logs, without needing to attach a profiler.
|
||||
|
||||
### Conventions applied to every phase below
|
||||
|
||||
**Debug logging.** Every phase that changes an algorithmic choice (not phase 0, which *is* the logging) adds `tracing::debug!`/`trace!` at points that let a cluster run's logs answer "did this help": counts, ratios, and timings that quantify the specific claim that phase makes — e.g. phase 3 must log how many MPHF `find` calls were saved by k-mer-level dedup (the whole justification for that phase), phase 4 must log per-column scan timings, phase 5 must log actual retained-memory / sparsity ratios achieved. Prefer one structured `debug!` per chunk (fields, not prose) over free-text — the cluster logs will be the only evidence available for judging these choices, so they need to be grep/awk-able, not just readable.
|
||||
|
||||
**Unit tests.** This project's convention (`obiread`, `obikseq`, `obidebruinj`, `obicompactvec`, `obilayeredmap`, `obiskio`, `obifastwrite`) is `#[cfg(test)] #[path = "tests/<name>.rs"] mod tests;` at the bottom of the source file, with the actual test code in a sibling `src/tests/<name>.rs`. Neither `obikmer` nor `obikpartitionner` (the two crates phases 3 and 5 touch most) currently have a `src/tests/` directory at all — this needs creating, following the existing pattern exactly, not inventing a new one.
|
||||
|
||||
**Workflow (`jj`).** Work happens in a fresh `jj` commit, easy to abandon. `jj new` between phases is reasonable where it helps isolate a phase for review, but only when the working copy compiles at that point (project convention) — phase 3's internal sub-steps (batch dedup change, then `query_layer.rs` split, then the new return shape) will likely not each compile independently since they're one coupled change, so treat "commit boundary" and "plan phase boundary" as related but not forced to match 1:1; use judgement per phase rather than mechanically splitting on every bullet.
|
||||
|
||||
### Phase 0 — Instrumentation (prerequisite for measuring every later phase)
|
||||
|
||||
**Goal**: make core utilization, throughput, and per-stage timing visible on a real run, so phases 1–5 can be justified with numbers instead of assumption.
|
||||
|
||||
- `obikmer/src/cmd/query.rs`: wrap `run()`'s main loop with `obisys::Reporter`/`Stage::start("query")`/`.stop()`, printed at the end via `rep.print()` — same pattern as `index.rs`/`filter.rs`.
|
||||
- Add an `obisys::spinner("query")` progress bar around the `pipe.apply(...)` loop, with an EMA throughput readout (bases/s or k-mers/s, mirroring `steps::scatter`'s `ema_rate` computation, `scatter.rs:88-118`) and live gauges for "chunks in flight" / "workers busy" — reuse the `AtomicU32` counter pattern from `scatter.rs` (`flat_active`, `transform_active`) rather than inventing a new one.
|
||||
- Add `max_open_files: Option<usize>` to `QueryArgs` and a `effective_max_open()` method mirroring `CommonArgs::effective_max_open()` (`obikmer/src/cli.rs:90-94`) — needed by phase 1's `throttle()` call. (`QueryArgs` can't just embed `CommonArgs` — it doesn't take `kmer_size`/`minimizer_size`/`partitions`/`level_max`/`theta` from the CLI, those come from the index metadata — so this is a small standalone addition, not a flatten.)
|
||||
- Add one structured `debug!` per `process_chunk` call: chunk byte size, sequence count, s-mer count, wall time, and (once later phases exist) the fields they add — this single log line is the baseline every later phase's own logging gets compared against.
|
||||
- **Validation**: none needed beyond "the numbers appear and look sane" — this phase changes no query logic or output.
|
||||
- **Deliverable used by every phase below**: a before/after throughput and core-utilization measurement on the reference machine.
|
||||
|
||||
### Phase 1 — Parallel per-file I/O (fixes root cause of low core utilization)
|
||||
|
||||
**Goal**: file opening, decompression, and chunk-boundary parsing run across the `n_workers` pool instead of serially in the pipe's dedicated source thread.
|
||||
|
||||
- `obikmer/src/cmd/query.rs`:
|
||||
- Replace the `paths.into_iter().flat_map(read_sequence_chunks_sized(...))` construction (current `run()`, building `all_chunks`) with `obipipeline::throttle(paths.into_iter(), args.effective_max_open())`, passed as the pipe's `input`.
|
||||
- Add a new `QueryData::Path(PathBuf)` variant (alongside `Chunk`/`Output`) to carry the throttled path through the pipe's type-erasure mechanism.
|
||||
- Add a new **first** pipe stage, `Flat`/fallible (`||?`), modeled on `scatter.rs:60-86` and `superkmer.rs:54-65`: given a `Throttled<PathBuf>`, call `read_sequence_chunks_sized(path, chunk_bytes)` and yield each `Rope` chunk, keeping `pw.guard` alive until the file's iterator is exhausted (reuse or adapt `scatter.rs`'s `GuardedIter` wrapper — same lifetime problem, same fix).
|
||||
- The existing `process_chunk` transform stage becomes the pipe's **second** stage, unchanged in its own logic — it still receives one `Rope` chunk at a time, just no longer all coming from one serial source.
|
||||
- `make_pipe!` invocation grows from one stage (`Chunk => Output`) to two (`Path => Chunk => Output`).
|
||||
- Log, per file: time spent waiting on the `throttle()` slot (queueing due to `max_open`), and time spent opening/decompressing/producing the first chunk — this is what directly proves (or disproves) that I/O is now spread across workers instead of serialized.
|
||||
- **Validation**: run `query` on a small multi-file input, diff output against the pre-change version — content must be identical; **record order across files is not guaranteed to be preserved** even before this change (chunk-level dispatch across `n_workers` already reorders completions), so the diff must be order-insensitive (sort by read id, or compare as sets) if it wasn't already.
|
||||
- **Measure**: core utilization on the reference machine with several large input files, compare against phase 0's baseline.
|
||||
|
||||
### Phase 2 — Genome-aware chunk-size formula (fixes OOM)
|
||||
|
||||
**Goal**: `chunk_bytes` reflects actual per-chunk memory (`O(n_genomes)`), not a fixed multiplier.
|
||||
|
||||
- `obikmer/src/cmd/query.rs`, `run()`: `n_genomes` and `args.detail` are already computed above the `chunk_bytes` calculation (`n_genomes` at the top of `run()`, before line 407 in the current file) — reorder if needed, then replace:
|
||||
```rust
|
||||
let computed = avail / (n_workers as u64 * 16);
|
||||
```
|
||||
with a formula that scales the divisor by `n_genomes` (and roughly doubles it when `--detail` is set, since `cov` duplicates the per-genome accumulation): e.g. `per_chunk_multiplier = base_overhead + n_genomes as u64 * BYTES_PER_KMER_PER_GENOME * if detail { 2 } else { 1 }`, replacing the flat `16`. `BYTES_PER_KMER_PER_GENOME` should be derived from `KmerResults`'s actual layout (`4` bytes per `u32` entry in `data`, plus the `bool` in `in_index`, plus `win_min`'s equal-sized buffer) rather than guessed.
|
||||
- `args.chunk_size` (manual `--chunk-size` override) keeps taking priority, unchanged.
|
||||
- Log the resolved `chunk_bytes`, `n_genomes`, and the estimated peak per-chunk memory (`chunk_bytes` × the same multiplier used to derive it) once at startup — lets a cluster run confirm the estimate was actually respected, not just that the process didn't get OOM-killed (which could also happen to be true for the wrong reason).
|
||||
- **Validation**: build a test index with a large `n_genomes` (e.g. hundreds), run `query` with default chunk sizing under a memory limit (`ulimit -v` or a cgroup), confirm it no longer gets OOM-killed and that memory scales as predicted when `n_genomes` grows.
|
||||
- **Note**: this phase is superseded once phase 5 lands (sparse retained memory no longer scales with `n_genomes × total_kmers` at all) — but it's needed immediately regardless, since phases 3–5 are a bigger, riskier change and users need a working `query` in the meantime.
|
||||
|
||||
### Phase 3 — K-mer-level dereplication, staged MPHF/matrix lookup
|
||||
|
||||
**Goal**: replace superkmer-level dedup with k-mer-level dedup (roadmap point 5), and split the fused MPHF-find/matrix-fetch (point 4) so stage 1's output is bucketed by layer and MPHF slot (point 6).
|
||||
|
||||
- `obikmer/src/cmd/query.rs`:
|
||||
- Replace `QueryBatch::from_records`'s dedup map (`HashMap<RoutableSuperKmer, Vec<SKDesc>>`, current `query.rs:112`) with a per-partition `HashMap<CanonicalKmer, Vec<(seq_idx: u32, pos: u32)>>`, built in the same `SuperKmerIter` pass: superkmer construction and partition routing (`part_idx` from the superkmer's minimizer hash) are unchanged, only the granularity of what gets deduplicated changes — each `CanonicalKmer` within a superkmer is inserted individually instead of the whole superkmer being the dedup key.
|
||||
- **Verified**: `CanonicalKmer` (`obikseq/src/kmer.rs:390`, `pub type CanonicalKmer = CanonicalKmerOf<KLen>`) — the underlying `CanonicalKmerOf<L>` derives `Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash` (`kmer.rs:269`). Usable as a `HashMap`/`HashSet` key as-is, no change needed.
|
||||
- `obikpartitionner/src/query_layer.rs`:
|
||||
- Split `QueryLayer::find_into` (`query_layer.rs:48-67`) into two methods: `find_slot(&self, kmer: CanonicalKmer) -> Option<usize>` (MPHF only, no matrix touch) and keep `fill_row` as-is for phase 4 to call later.
|
||||
- Replace `query_partition_with`'s inner loop (`query_layer.rs:103-113`) with a version that, for each unique `CanonicalKmer`, calls `find_slot` across the partition's layers (stopping at first hit, same as today), and instead of immediately filling a row, records `(layer_idx, slot)`.
|
||||
- New return shape for the partition-level query, replacing today's `on_hit(sk_idx, kmer_idx, row)` callback: `HashMap<layer_idx, HashMap<slot, Vec<(seq_idx, pos)>>>` (roadmap point 6) — built directly from the k-mer dedup map's `Vec<(seq_idx,pos)>` values, keyed by the resolved slot instead of the k-mer.
|
||||
- **This phase alone has no throughput benefit yet** (matrix fetch still happens, just deferred) beyond the k-mer-level dedup itself (fewer MPHF calls when queries have overlapping/repeated k-mers) — its purpose is to produce the input phase 4 needs. Land phase 3+4 together, not phase 3 alone, per the "don't split 3–5 across releases" note above.
|
||||
- Log, per chunk: total k-mer occurrences vs. unique `CanonicalKmer` count (the dedup ratio — the entire justification for this phase) and the resulting MPHF `find` call count. If the dedup ratio is close to `1.0` on real query data (little redundancy), that's the cluster run telling us this phase wasn't worth it — the logging needs to be able to say that, not just confirm the happy path.
|
||||
- **Unit tests**: create `obikmer/src/cmd/tests/query.rs` (new `src/tests/` dir for this crate, following the project's `#[cfg(test)] #[path = "tests/query.rs"] mod tests;` convention) and `obikpartitionner/src/tests/query_layer.rs` (likewise new for this crate). Cover: the k-mer-level dedup map construction on synthetic sequences with known repeated/overlapping k-mers (assert unique-kmer count and occurrence lists); the `find_slot`/bucket-by-layer-and-slot construction against a small hand-built `QueryLayer` fixture, asserting the `(layer_idx, slot, seq_idx, pos)` tuples match what the old per-occurrence loop would have produced.
|
||||
|
||||
### Phase 4 — Column-major matrix fetch, NUMA-aware (roadmap points 7–8)
|
||||
|
||||
**Goal**: replace `fill_row`-per-hit (row-major, worst-case mmap locality) with a column-major scan, driven through `PartitionRunner` rather than ad-hoc `rayon`, and extend the same mechanism to the partition loop itself.
|
||||
|
||||
- `obicompactvec`: add public per-column accessors alongside the existing `pub(crate)` ones — e.g. `PersistentCompactIntMatrix::col_get(c: usize, slot: usize) -> u32` and equivalent for `PersistentBitMatrix`/`Packed*` variants (the underlying `col(c)`/`col_view(c)` + `.get(slot)` machinery already exists internally, per `intmatrix.rs`/`bitmatrix.rs` — this phase only needs to expose it).
|
||||
- `obikpartitionner`: new function, e.g. `fetch_layer_columns(layer: &QueryLayer, hits: &HashMap<slot, Vec<(seq_idx,pos)>>, n_genomes: usize) -> HashMap<genome_idx, Vec<(seq_idx, pos, count)>>` — for each genome column `g` in `0..n_genomes`, scan `hits`'s slot keys, call the new `col_get`, keep nonzero results, append to the per-genome accumulator.
|
||||
- `obikmer/src/cmd/query.rs` (or wherever this ends up being orchestrated): drive both the per-partition loop (today `query.rs:250-278`) and, within each partition/layer, the per-genome column loop, through `obikindex::numa::PartitionRunner::run()` (`obikindex/src/numa.rs`) instead of a plain loop or `rayon::into_par_iter()`. `obikmer` already depends on `obikindex`, so this needs no new dependency.
|
||||
- **Composition with existing parallelism — open design decision, resolve empirically before finishing this phase** (flagged as unresolved in the roadmap, point 8): three concurrency mechanisms now coexist — `obipipeline`'s raw-thread chunk-worker pool, `PartitionRunner`'s NUMA-pinned Rayon pools, and whichever axis (partition or column) is chosen as `PartitionRunner`'s unit of work. Proposed starting point to *measure*, not a final answer:
|
||||
1. Reduce `obipipeline`'s `n_workers` for the chunk pipe to something small and NUMA-node-shaped (e.g. one chunk in flight per NUMA node) once `PartitionRunner` is doing the fine-grained fan-out within `process_chunk` — avoids the two mechanisms both trying to claim all 192 cores independently.
|
||||
2. Start with `PartitionRunner` driving the **column** axis only (point 8) inside each partition/layer, keep the **partition** loop (point 7) sequential within that; add partition-level `PartitionRunner` usage only if profiling still shows idle cores with many small partitions.
|
||||
3. Reuse `PartitionRunner`'s existing dual CPU/IO activation signal as-is (`numa_partition_runner.md`'s `IoSample`/`CpuSample`) rather than tuning new thresholds up front — the column scan's mmap-page-fault-bound shape is close enough to `pack_matrices`'s documented profile to start from the same constants.
|
||||
- Log, per layer processed: number of genome columns scanned, number of `get()` calls, wall time, and (via `PartitionRunner`'s own `on_done`) active-worker count over time per NUMA node — this is the only way to see, from cluster logs alone, whether the composition question above actually got resolved well or whether one mechanism is starving the other.
|
||||
- **Unit tests**: `obikpartitionner/src/tests/query_layer.rs` (same file as phase 3, extended) — cover `fetch_layer_columns` on a small hand-built multi-column matrix fixture (mix of `Columnar`/`Packed`/`Implicit`), asserting the sparse `(seq_idx, pos, count)` output matches a naive row-major reference implementation computed in the test itself. `obicompactvec/src/tests/` (existing `tests/mod.rs` umbrella, per its established convention) gets the new `col_get` accessors' own unit tests alongside its current matrix tests.
|
||||
- **Validation**: microbenchmark `fetch_layer_columns` against today's per-hit `fill_row` loop on a large-`n_genomes` layer, measuring wall time and (on Linux) page-fault counters (`/proc/self/stat` minor/major faults) — the claimed win is locality, so it should show up there even before end-to-end `query` numbers move.
|
||||
|
||||
### Phase 5 — Sparse Findere rework (roadmap point 9)
|
||||
|
||||
**Goal**: replace the dense `KmerResults`/`win_min` sliding-window scan with one operating on phase 4's sparse per-genome output.
|
||||
|
||||
- `obikmer/src/cmd/query.rs`, `process_chunk`:
|
||||
- Remove `KmerResults` (`query.rs:157-202`) and the dense `win_min` allocation (`query.rs:290-291`, sized `max_n_kmers × n_genomes`).
|
||||
- Keep a lightweight dense `in_index: Vec<bool>` per chunk (sized `total_kmers`, independent of `n_genomes`) from phase 3's stage 1 — still needed for `kmer_missing` bookkeeping (leftmost-s-mer-of-window membership test), which phase 4's sparse structure doesn't carry (a k-mer with no genome hit has no entry there at all).
|
||||
- New per-`(seq_idx, genome)` scan: for each genome's `Vec<(seq_idx, pos, count)>` (sorted, per phase 4), group by `seq_idx` (contiguous after sort), then within each sequence's positions detect runs of `pos, pos+1, pos+2, ...` of length ≥ `z`; within each run, the existing monotone-deque window-minimum logic (`query.rs`'s current `dq` loop, conceptually unchanged) applies — but the deque now only scans real entries in the run, never zero-filled gaps.
|
||||
- Update `SeqAcc` accumulation and `emit_batch` to consume this per-genome sparse iteration instead of `results.val`/`results.is_in_index`.
|
||||
- `--detail`/`cov`: build sparsely during the same scan (only positions with a confirmed contribution get an entry), densify into the `[u32]` JSON array only in `emit_batch`, only for genomes/sequences actually being serialized (per roadmap point 9's note, `query.rs:304-308`'s current dense allocation goes away).
|
||||
- Log, per chunk: total sparse entries retained vs. what the old dense `KmerResults` would have allocated (`total_smers × n_genomes`) — the sparsity ratio is this phase's entire reason for existing, so it must be directly visible in the logs, not inferred from process RSS. Also log the run-detection stats (number of runs found, average run length) — a low average run length relative to `z` would mean most positions still fail to form a full window, worth knowing.
|
||||
- **Unit tests**: `obikmer/src/cmd/tests/query.rs` (extended from phase 3) — the property test described below is the primary deliverable here, not an afterthought; write it as an actual `#[test]` (or a small internal fuzz/property-style loop over randomized fixtures if a property-testing crate isn't already a dependency — check before adding one, per this project's dependency-approval rule) rather than a one-off manual comparison.
|
||||
- **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.
|
||||
|
||||
### 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).
|
||||
|
||||
### Cross-cutting risks
|
||||
|
||||
- **Thread-budget oversubscription** (phase 4): the single biggest unresolved design question in this whole plan — see phase 4's composition note. Should be settled with real measurements early in phase 4, not assumed from the design alone.
|
||||
- **`obicompactvec` API surface growth** (phase 4): new public per-column accessors are additive (existing `fill_row`/`row` stay for other callers — `dump`, `select`, distance computations) — no breaking change expected, but worth checking `obicompactvec`'s other callers aren't already relying on `fill_row` being the only/cheapest access path in a way that would make maintaining two access patterns (row-major and column-major) a real maintenance cost rather than a one-off addition.
|
||||
- **`PersistentBitMatrix::Implicit`'s hardcoded `n_cols: 1` — resolved, not a bug.** `LayerMeta`'s own doc comment (`obicompactvec/src/layer_meta.rs:1-9`) states it is written "alongside `mphf.bin`" and read by `PersistentBitMatrix::open` "to determine `n_rows` for **the implicit (mono-genome presence/absence) case**" — i.e. `Implicit` is a documented single-genome fast path (no presence matrix needed when there is trivially one genome), not a generic "no matrix built yet" fallback. `n_cols: 1` is correct by design for the case it's meant to handle. Phase 4's column loop is safe as planned — this was worth checking once, doesn't need further action.
|
||||
|
||||
Reference in New Issue
Block a user