From 8967a20ed7298dd92579c416c292ab502ba64c2a Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Thu, 13 Aug 2026 10:12:00 +0200 Subject: [PATCH] refactor(cmd): restructure modules and decompose query command Convert single-file modules to directory-based layouts across the cmd crate. Decompose the monolithic query command into dedicated submodules for batching, chunk processing, sparse finding, and output formatting. Introduce a new utils module to handle index management operations including statistics reporting, label renaming, and partition analysis. --- .../src/cmd/{annotate.rs => annotate/mod.rs} | 0 src/obikmer/src/cmd/{dump.rs => dump/mod.rs} | 0 .../src/cmd/{estimate.rs => estimate/mod.rs} | 0 .../src/cmd/{filter.rs => filter/mod.rs} | 0 .../src/cmd/{index.rs => index/mod.rs} | 0 .../src/cmd/{merge.rs => merge/mod.rs} | 0 src/obikmer/src/cmd/{pack.rs => pack/mod.rs} | 0 src/obikmer/src/cmd/query.rs | 851 ------------------ src/obikmer/src/cmd/query/batch.rs | 74 ++ src/obikmer/src/cmd/query/chunk.rs | 288 ++++++ src/obikmer/src/cmd/query/findere.rs | 71 ++ src/obikmer/src/cmd/query/mod.rs | 340 +++++++ src/obikmer/src/cmd/query/output.rs | 50 + src/obikmer/src/cmd/query/smer_index.rs | 41 + .../cmd/{tests/query.rs => query/tests.rs} | 7 +- .../src/cmd/{reindex.rs => reindex/mod.rs} | 0 .../src/cmd/{select.rs => select/mod.rs} | 0 .../cmd/{superkmer.rs => superkmer/mod.rs} | 0 .../src/cmd/{unitig.rs => unitig/mod.rs} | 0 src/obikmer/src/cmd/utils/maintenance.rs | 112 +++ src/obikmer/src/cmd/utils/mod.rs | 85 ++ .../{utils.rs => utils/partition_stats.rs} | 196 +--- 22 files changed, 1069 insertions(+), 1046 deletions(-) rename src/obikmer/src/cmd/{annotate.rs => annotate/mod.rs} (100%) rename src/obikmer/src/cmd/{dump.rs => dump/mod.rs} (100%) rename src/obikmer/src/cmd/{estimate.rs => estimate/mod.rs} (100%) rename src/obikmer/src/cmd/{filter.rs => filter/mod.rs} (100%) rename src/obikmer/src/cmd/{index.rs => index/mod.rs} (100%) rename src/obikmer/src/cmd/{merge.rs => merge/mod.rs} (100%) rename src/obikmer/src/cmd/{pack.rs => pack/mod.rs} (100%) delete mode 100644 src/obikmer/src/cmd/query.rs create mode 100644 src/obikmer/src/cmd/query/batch.rs create mode 100644 src/obikmer/src/cmd/query/chunk.rs create mode 100644 src/obikmer/src/cmd/query/findere.rs create mode 100644 src/obikmer/src/cmd/query/mod.rs create mode 100644 src/obikmer/src/cmd/query/output.rs create mode 100644 src/obikmer/src/cmd/query/smer_index.rs rename src/obikmer/src/cmd/{tests/query.rs => query/tests.rs} (98%) rename src/obikmer/src/cmd/{reindex.rs => reindex/mod.rs} (100%) rename src/obikmer/src/cmd/{select.rs => select/mod.rs} (100%) rename src/obikmer/src/cmd/{superkmer.rs => superkmer/mod.rs} (100%) rename src/obikmer/src/cmd/{unitig.rs => unitig/mod.rs} (100%) create mode 100644 src/obikmer/src/cmd/utils/maintenance.rs create mode 100644 src/obikmer/src/cmd/utils/mod.rs rename src/obikmer/src/cmd/{utils.rs => utils/partition_stats.rs} (50%) diff --git a/src/obikmer/src/cmd/annotate.rs b/src/obikmer/src/cmd/annotate/mod.rs similarity index 100% rename from src/obikmer/src/cmd/annotate.rs rename to src/obikmer/src/cmd/annotate/mod.rs diff --git a/src/obikmer/src/cmd/dump.rs b/src/obikmer/src/cmd/dump/mod.rs similarity index 100% rename from src/obikmer/src/cmd/dump.rs rename to src/obikmer/src/cmd/dump/mod.rs diff --git a/src/obikmer/src/cmd/estimate.rs b/src/obikmer/src/cmd/estimate/mod.rs similarity index 100% rename from src/obikmer/src/cmd/estimate.rs rename to src/obikmer/src/cmd/estimate/mod.rs diff --git a/src/obikmer/src/cmd/filter.rs b/src/obikmer/src/cmd/filter/mod.rs similarity index 100% rename from src/obikmer/src/cmd/filter.rs rename to src/obikmer/src/cmd/filter/mod.rs diff --git a/src/obikmer/src/cmd/index.rs b/src/obikmer/src/cmd/index/mod.rs similarity index 100% rename from src/obikmer/src/cmd/index.rs rename to src/obikmer/src/cmd/index/mod.rs diff --git a/src/obikmer/src/cmd/merge.rs b/src/obikmer/src/cmd/merge/mod.rs similarity index 100% rename from src/obikmer/src/cmd/merge.rs rename to src/obikmer/src/cmd/merge/mod.rs diff --git a/src/obikmer/src/cmd/pack.rs b/src/obikmer/src/cmd/pack/mod.rs similarity index 100% rename from src/obikmer/src/cmd/pack.rs rename to src/obikmer/src/cmd/pack/mod.rs diff --git a/src/obikmer/src/cmd/query.rs b/src/obikmer/src/cmd/query.rs deleted file mode 100644 index b63f8283..00000000 --- a/src/obikmer/src/cmd/query.rs +++ /dev/null @@ -1,851 +0,0 @@ -use std::collections::{HashMap, VecDeque}; -use std::io::{self, BufWriter, Write}; -use std::path::PathBuf; -use std::sync::Arc; -use std::sync::atomic::{AtomicU32, AtomicU64, Ordering}; -use std::time::Instant; - -use clap::Args; -use obikindex::KmerIndex; -use obikpartitionner::{KmerDesc, QueryHit, QueryStats}; -use obikrope::Rope; -use obikseq::CanonicalKmer; -use obilayeredmap::IndexMode; -use obipipeline::{Throttled, ThrottleGuard, throttle}; -use obiread::chunk::read_sequence_chunks_sized; -use obiread::record::{SeqRecord, parse_chunk}; -use obiskbuilder::SuperKmerIter; -use obisys::{Reporter, Stage, available_memory_bytes, spinner}; -use tracing::{debug, info}; - -// ── Pipeline data ───────────────────────────────────────────────────────────── - -enum QueryData { - Path(Throttled), - Chunk(Rope), - Output(Vec), -} - -// SAFETY: Rope contains Cell which is !Sync, but pipeline items are owned -// exclusively through channels — no item is ever shared across threads. -unsafe impl Send for QueryData {} -unsafe impl Sync for QueryData {} - -// ── CLI ─────────────────────────────────────────────────────────────────────── - -#[derive(Args)] -pub struct QueryArgs { - /// Index directory - pub index: PathBuf, - - /// Input sequences (FASTA/FASTQ, optionally gzip-compressed) - #[arg(num_args = 1..)] - pub inputs: Vec, - - /// Report per-position coverage vectors per genome (adds "coverage" to JSON) - #[arg(long)] - pub detail: bool, - - /// Enable 1-mismatch approximate matching - #[arg(long)] - pub mismatch: bool, - - /// Count k-mers absent from the index (adds kmer_missing annotation) - #[arg(long)] - pub count_missing: bool, - - /// Report per-genome presence (0/1) instead of raw counts - #[arg(long)] - pub force_presence: bool, - - /// Minimum accumulated match count to declare a genome present (implies --force-presence) - #[arg(long, default_value_t = 1)] - pub presence_threshold: u32, - - /// Override the Findere z parameter from index metadata - #[arg(short = 'z', long)] - pub findere_z: Option, - - /// Number of worker threads - #[arg( - short = 'T', - long, - default_value_t = obisys::effective_parallelism() - )] - pub threads: usize, - - /// I/O chunk size in MiB (default: auto-sized from available RAM and thread count) - #[arg(long)] - pub chunk_size: Option, - - /// Maximum number of input files open simultaneously. - /// Defaults to threads/4 (minimum 1). Keep below the number of workers - /// to ensure CPU workers are always available for the transform stage. - #[arg(long)] - pub max_open_files: Option, -} - -impl QueryArgs { - pub fn effective_max_open(&self) -> usize { - self.max_open_files - .unwrap_or_else(|| (self.threads / 4).max(1)) - .max(1) - } -} - -// ── QueryBatch ──────────────────────────────────────────────────────────────── - -/// A batch of query sequences, with k-mers deduplicated directly (not just at -/// the superkmer level) and pre-split by partition. -/// -/// Superkmer *construction* (`SuperKmerIter`) is still required — it's the -/// mechanism that computes minimizers and partition routing — but the dedup -/// key is the canonical k-mer, not the superkmer: two different superkmers -/// that happen to share a k-mer (read overlaps, repeats, a SNP splitting an -/// otherwise-identical run) are deduplicated too, not just identical whole -/// superkmers. This also means each unique k-mer triggers at most one MPHF -/// lookup, not one per occurrence. -pub struct QueryBatch { - /// Sequence ids in batch order. - pub ids: Vec, - /// Raw sequence bytes (for output), in batch order. - pub seqs: Vec>, - /// Total kmer count per sequence (used for `--detail` coverage allocation). - pub n_kmers: Vec, - /// Deduplicated k-mer occurrences, one map per partition. - pub by_partition: Vec>>, -} - -impl QueryBatch { - /// Build a batch from a vec of parsed sequence records, deduplicating - /// k-mers and routing them to partitions in the same pass. - pub fn from_records( - records: Vec, - k: usize, - level_max: usize, - theta: f64, - n_partitions: usize, - ) -> Self { - let mut ids = Vec::with_capacity(records.len()); - let mut seqs = Vec::with_capacity(records.len()); - let mut n_kmers = Vec::with_capacity(records.len()); - let mask = (n_partitions as u64) - 1; - let mut by_partition: Vec>> = - (0..n_partitions).map(|_| HashMap::new()).collect(); - - for (seq_idx, record) in records.into_iter().enumerate() { - let mut kmer_offset = 0u32; - - for rsk in SuperKmerIter::new(&record.normalized, k, level_max, theta) { - let part_idx = (rsk.minimizer().seq_hash() & mask) as usize; - let map = &mut by_partition[part_idx]; - for (j, kmer) in rsk.superkmer().iter_canonical_kmers().enumerate() { - map.entry(kmer).or_default().push(KmerDesc { - seq_idx: seq_idx as u32, - pos: kmer_offset + j as u32, - }); - } - let n = (rsk.seql() - k + 1) as u32; - kmer_offset += n; - } - - ids.push(record.id); - seqs.push(record.sequence); - n_kmers.push(kmer_offset); - } - - Self { - ids, - seqs, - n_kmers, - by_partition, - } - } -} - -// ── SmerIndex — sparse "was this k-mer found at all" bookkeeping ───────────── - -/// 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, // total_smers - offsets: Vec, // offsets[i]..offsets[i+1] = s-mer range for sequence i -} - -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); - for &n in n_kmers_per_seq { - total += n as usize; - offsets.push(total); - } - Self { - in_index: vec![false; total], - offsets, - } - } - - /// 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 reports per-genome values. - fn mark_found(&mut self, seq: usize, kmer: usize) { - let abs = self.offsets[seq] + kmer; - self.in_index[abs] = true; - } - - #[inline] - fn is_in_index(&self, seq: usize, kmer: usize) -> bool { - self.in_index[self.offsets[seq] + kmer] - } -} - -// ── 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, 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 ────────────────────────────────────────────────── - -struct SeqAcc { - kmer_count: u32, - kmer_missing: u32, - genome_totals: Vec, -} - -impl SeqAcc { - fn new(n_genomes: usize) -> Self { - Self { - kmer_count: 0, - kmer_missing: 0, - genome_totals: vec![0u32; n_genomes], - } - } -} - -// ── process_chunk ───────────────────────────────────────────────────────────── - -fn process_chunk( - idx: &KmerIndex, - rope: Rope, - k: usize, - n_genomes: usize, - n_partitions: usize, - with_counts: bool, - effective_z: usize, - detail: bool, - count_missing: bool, - force_presence: bool, - presence_threshold: u32, -) -> Vec { - let chunk_start = Instant::now(); - let chunk_bytes = rope.len(); - - let records = parse_chunk(&rope, k); - if records.is_empty() { - return Vec::new(); - } - - let batch = QueryBatch::from_records(records, k, 6, 0.7, n_partitions); - let n_seqs = batch.ids.len(); - - // Estimate QueryBatch::by_partition's actual memory footprint: the - // k-mer-level dedup map (roadmap point 5) — one HashMap> per partition, sized by *unique* k-mers, not shrunk by - // dedup. On real workloads with a low intra-chunk duplication rate this - // can dwarf every other per-chunk structure, including the sparse - // Findere ones logged further down — unlike those, chunk_bytes's formula - // (run()) does not account for this at all today. Measured by allocated - // capacity, not logical length, to reflect real memory pressure - // (HashMap/Vec growth slack) — `by_partition` is alive for the entire - // process_chunk call (never drained, only iterated by reference), so - // this is its footprint for the whole chunk lifetime, not a transient. - let hashmap_slot_bytes = (std::mem::size_of::() - + std::mem::size_of::>() - + 1) as u64; // +1 ≈ hashbrown control byte per slot - let by_partition_map_bytes: u64 = batch - .by_partition - .iter() - .map(|m| m.capacity() as u64 * hashmap_slot_bytes) - .sum(); - let by_partition_desc_bytes: u64 = batch - .by_partition - .iter() - .flat_map(|m| m.values()) - .map(|v| v.capacity() as u64 * std::mem::size_of::() as u64) - .sum(); - let by_partition_bytes = by_partition_map_bytes + by_partition_desc_bytes; - - debug!( - n_unique_kmers_total = batch.by_partition.iter().map(|m| m.len() as u64).sum::(), - by_partition_map_bytes, - by_partition_desc_bytes, - by_partition_bytes, - chunk_bytes, - "by_partition memory retained" - ); - - // 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> = (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 - // entire justification for k-mer-level dereplication (see query.md, - // Future work point 5). If this ratio stays close to 1.0 on real data, - // dereplication isn't paying for itself and that should show up here. - let n_occurrences: u64 = batch.n_kmers.iter().map(|&n| n as u64).sum(); - let mut query_stats = QueryStats::default(); - - for (part_idx, kmers) in batch.by_partition.iter().enumerate() { - if kmers.is_empty() { - continue; - } - - let stats = idx.partition() - .query_partition_with( - part_idx, - kmers, - n_genomes, - with_counts, - |event| match event { - QueryHit::Found(descs) => { - for desc in descs { - smer_index.mark_found(desc.seq_idx as usize, desc.pos as usize); - } - } - QueryHit::Value(descs, g, v) => { - for desc in descs { - by_genome[g].push((desc.seq_idx, desc.pos, v)); - } - } - }, - ) - .unwrap_or_else(|e| { - eprintln!("query error on partition {part_idx}: {e}"); - std::process::exit(1); - }); - query_stats += stats; - } - - debug!( - n_occurrences, - n_unique_kmers = query_stats.n_unique_kmers, - n_mphf_calls = query_stats.n_mphf_calls, - n_hits = query_stats.n_hits, - n_columns_scanned = query_stats.n_columns_scanned, - n_col_get_calls = query_stats.n_col_get_calls, - "k-mer dedup + column-major fetch" - ); - - // ── Sparse Findere: per-genome run detection + sliding-window minimum ──── - // - // 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 = batch - .n_kmers - .iter() - .map(|&n| { - let n = n as usize; - 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::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" - ); - - // Actual bytes retained by the sparse hit structures (by_genome + - // confirmed_by_genome, both alive simultaneously at this point — see the - // chunk-size formula's comment in `run()`), by allocated capacity rather - // than logical length so this reflects real memory pressure including - // Vec growth slack. `empirical_multiplier` is directly comparable to - // BYTES_PER_KMER_PER_GENOME (`run()`) — the ratio a cluster run's logs - // need to judge whether that constant is over- or under-conservative for - // real data, instead of guessing. - const HIT_ENTRY_BYTES: u64 = std::mem::size_of::<(u32, u32, u32)>() as u64; - let by_genome_bytes: u64 = by_genome.iter().map(|v| v.capacity() as u64 * HIT_ENTRY_BYTES).sum(); - let confirmed_bytes: u64 = confirmed_by_genome.iter().map(|v| v.capacity() as u64 * HIT_ENTRY_BYTES).sum(); - let retained_bytes = by_genome_bytes + confirmed_bytes; - - debug!( - by_genome_bytes, - confirmed_bytes, - retained_bytes, - chunk_bytes, - empirical_multiplier = retained_bytes as f64 / chunk_bytes.max(1) as f64, - "sparse memory retained" - ); - - // ── Accumulate: genome totals (per genome, from confirmed hits) ────────── - let mut accs: Vec = (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>> = if detail { - n_kmers_out.iter().map(|&n| vec![vec![0u32; n]; n_genomes]).collect() - } else { - Vec::new() - }; - 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; - } - } - } - - // Capacity estimate: actual sequence + ID bytes, plus JSON overhead per record. - // JSON per record ≈ 50 fixed chars + ~20 per genome (label + count value) + 100 (overhead). - let seq_bytes: usize = batch.seqs.iter().map(|s| s.len()).sum(); - let id_bytes: usize = batch.ids.iter().map(|s| s.len()).sum(); - let cap = seq_bytes + id_bytes + n_seqs * (4 + 50 + n_genomes * 20) + 100; - let mut buf = Vec::with_capacity(cap); - emit_batch( - &batch, - &accs, - idx.meta(), - count_missing, - detail, - &cov, - &mut buf, - ); - - debug!( - chunk_bytes, - n_seqs, - n_smers = batch.n_kmers.iter().map(|&n| n as u64).sum::(), - wall_ms = chunk_start.elapsed().as_millis() as u64, - "process_chunk" - ); - - buf -} - -// ── GuardedChunkIter — keeps the throttle slot guard alive until the file is exhausted ── - -/// Wraps a per-file `Rope` chunk iterator together with its `ThrottleGuard`, -/// so the guard (and the throttle slot it holds) is only released once the -/// file has been fully read — never earlier, never held past that point. -struct GuardedChunkIter { - inner: Box + Send>, - _guard: ThrottleGuard, - files_open: Arc, -} - -impl Iterator for GuardedChunkIter { - type Item = Rope; - fn next(&mut self) -> Option { - self.inner.next() - } -} - -impl Drop for GuardedChunkIter { - fn drop(&mut self) { - self.files_open.fetch_sub(1, Ordering::Relaxed); - } -} - -// ── Entry point ─────────────────────────────────────────────────────────────── - -pub fn run(args: QueryArgs) { - let idx = Arc::new(KmerIndex::open(&args.index).unwrap_or_else(|e| { - eprintln!("error opening index: {e}"); - std::process::exit(1); - })); - - let k = idx.kmer_size(); - let n_genomes = idx.meta().genomes.len(); - let n_partitions = idx.n_partitions(); - let with_counts = idx.meta().config.with_counts; - let n_workers = args.threads.max(1); - - // Chunk size: each chunk stays in memory for its entire processing lifetime. - // - // Per-chunk memory is no longer a dense n_genomes-wide buffer (removed in - // the sparse Findere rework, see process_chunk) — it now scales with - // *actual hit count*, not with total_kmers_in_chunk × n_genomes - // unconditionally. BYTES_PER_KMER_PER_GENOME below is therefore a - // pathological-case bound, not a typical-case estimate: it protects - // against a fully-dense hit pattern (every k-mer of the query matching - // every genome — a degenerate case, e.g. low-complexity input theta- - // filtering should mostly reject, or an index of near-duplicate genomes), - // where by_genome and confirmed_by_genome (process_chunk) both end up - // holding one (seq_idx, pos, value) entry — 3 × u32 = 12 bytes, vs. 4 - // bytes for the old dense encoding, where position was implicit in the - // array index — per (k-mer, genome) pair, and *coexist simultaneously* - // (by_genome isn't freed before confirmed_by_genome is built), for a - // worst case of ~24 bytes/pair before Vec growth slack. `cov` remains - // fully dense when --detail is set (unaffected by the sparse rework), - // still roughly doubling the n_genomes-scaled cost. - // - // For realistic, sparse hit patterns actual memory is far below this - // bound — see the "sparse memory retained" debug log in process_chunk, - // which reports the empirical bytes-per-raw-byte multiplier actually - // observed per chunk, directly comparable to BYTES_PER_KMER_PER_GENOME - // below. Tightening this constant for typical-case throughput (at the - // cost of pathological-case safety margin) is a deliberate tuning - // decision to make from that data, not something to guess at here. - // - // BASE_OVERHEAD approximates what scales with chunk_bytes alone, - // independent of n_genomes: the Rope itself, parsed SeqRecord sequence + - // normalised bytes, the superkmer dedup map, and the JSON output buffer. - // Like the n_genomes-scaled term, this is an estimate — validate against - // actual peak RSS (Stage::stop's `rss` in the summary table) on real - // workloads rather than trusting it blindly. - // - // We target ≤ 50 % of available RAM across all concurrent workers - // (SAFETY_FACTOR). - const BASE_OVERHEAD: u64 = 4; - const BYTES_PER_KMER_PER_GENOME: u64 = 8; // pathological-case bound — see comment above - const SAFETY_FACTOR: u64 = 2; - - let detail_factor: u64 = if args.detail { 2 } else { 1 }; - let overhead_multiplier = - BASE_OVERHEAD + n_genomes as u64 * BYTES_PER_KMER_PER_GENOME * detail_factor; - - let chunk_bytes = args - .chunk_size - .map(|mb| mb * 1024 * 1024) - .unwrap_or_else(|| { - let avail = available_memory_bytes(); - let computed = avail / (n_workers as u64 * overhead_multiplier * SAFETY_FACTOR); - computed.clamp(4 * 1024 * 1024, 256 * 1024 * 1024) as usize - }); - - debug!( - chunk_bytes, - n_genomes, - detail = args.detail, - overhead_multiplier, - estimated_peak_chunk_bytes = chunk_bytes as u64 * overhead_multiplier, - "chunk-size formula resolved" - ); - - let effective_z: usize = args - .findere_z - .unwrap_or_else(|| match idx.meta().config.evidence { - IndexMode::Approx { z, .. } | IndexMode::Hybrid { z, .. } => z as usize, - IndexMode::Exact => 1, - }); - - info!( - "query: k={k}, {} genome(s), with_counts={with_counts}, z={effective_z}, \ - mismatch={}, detail={}", - n_genomes, args.mismatch, args.detail - ); - - if args.mismatch { - eprintln!("warning: --mismatch not yet implemented, ignored"); - } - - let detail = args.detail; - let count_missing = args.count_missing; - let force_presence = args.force_presence; - let presence_threshold = args.presence_threshold; - - // Throttled iterator over input file paths: at most `effective_max_open()` - // files are open at once. Opening + decompressing + chunking each file is - // now a Flat pipeline stage, executed across the `n_workers` pool — not - // serialised in the pipe's dedicated source thread (see steps::scatter / - // cmd::superkmer for the same pattern applied to indexing). - info!("query: chunk_size={}MiB, max_open_files={}", chunk_bytes / (1024 * 1024), args.effective_max_open()); - - let paths: Vec = args.inputs.iter().map(PathBuf::from).collect(); - let path_source = throttle(paths.into_iter(), args.effective_max_open()); - - // Instrumentation: total bytes processed (for the EMA throughput readout), - // number of files currently open/being chunked, and number of chunks - // currently being processed by a worker — all read from the spinner loop - // below, updated from inside the pipe closures. - let total_bytes = Arc::new(AtomicU64::new(0)); - let files_open = Arc::new(AtomicU32::new(0)); - let chunks_active = Arc::new(AtomicU32::new(0)); - - let pipe = obipipeline::make_pipe! { - QueryData : Throttled => Vec, - || { - let files_open = Arc::clone(&files_open); - move |pw: Throttled| -> GuardedChunkIter { - let path = pw.item; - let guard = pw.guard; - let path_str = path.to_str().unwrap_or("").to_owned(); - files_open.fetch_add(1, Ordering::Relaxed); - let open_start = Instant::now(); - // Hard-exit on file-open failure (mirrors the previous behaviour): - // propagating this as a pipeline Err would hit a known scheduler - // hang on early stage errors (obipipeline::scheduler::WorkerPool::run - // breaks its main loop without unblocking the still-running source - // thread, so the final `h.join()` never returns) — worth fixing in - // obipipeline itself, but out of scope here; sidestepping it like the - // original code already did is the safe choice for this change. - let iter = read_sequence_chunks_sized(&path_str, chunk_bytes).unwrap_or_else(|e| { - eprintln!("error opening {path_str}: {e}"); - std::process::exit(1); - }); - debug!( - path = %path_str, - open_ms = open_start.elapsed().as_millis() as u64, - "opened query input file" - ); - let err_path = path_str.clone(); - GuardedChunkIter { - inner: Box::new(iter.filter_map(move |r| match r { - Ok(rope) => Some(rope), - Err(e) => { - eprintln!("read error: {err_path}: {e}"); - None - } - })), - _guard: guard, - files_open: Arc::clone(&files_open), - } - } - } : Path => Chunk, - | { - let idx = Arc::clone(&idx); - let total_bytes = Arc::clone(&total_bytes); - let chunks_active = Arc::clone(&chunks_active); - move |rope: Rope| { - chunks_active.fetch_add(1, Ordering::Relaxed); - let bytes = rope.len() as u64; - let out = process_chunk( - &idx, rope, k, n_genomes, n_partitions, with_counts, - effective_z, detail, count_missing, force_presence, presence_threshold, - ); - total_bytes.fetch_add(bytes, Ordering::Relaxed); - chunks_active.fetch_sub(1, Ordering::Relaxed); - out - } - } : Chunk => Output, - }; - - let t = Stage::start("query"); - let pb = spinner("query"); - - let mut ema_rate: f64 = 0.0; - let mut last_t = Instant::now(); - let mut last_bytes: u64 = 0; - const ALPHA: f64 = 0.15; - - let mut out = BufWriter::new(io::stdout()); - for block in pipe.apply(path_source, n_workers, 2) { - if !block.is_empty() { - out.write_all(&block).expect("write error"); - } - - let now = Instant::now(); - let dt = now.duration_since(last_t).as_secs_f64(); - if dt > 0.1 { - let total = total_bytes.load(Ordering::Relaxed); - let instant = (total - last_bytes) as f64 / dt; - ema_rate = ALPHA * instant + (1.0 - ALPHA) * ema_rate; - last_t = now; - last_bytes = total; - let bp = total as f64; - let (count_str, rate_str) = if bp >= 1e9 { - (format!("{:.2} GB", bp / 1e9), format!("{:.0} MB/s", ema_rate / 1e6)) - } else { - (format!("{:.0} MB", bp / 1e6), format!("{:.0} MB/s", ema_rate / 1e6)) - }; - let active = chunks_active.load(Ordering::Relaxed); - let open = files_open.load(Ordering::Relaxed); - pb.set_message(format!("{count_str} {rate_str} [files open: {open}, chunks in flight: {active}]")); - } - } - out.flush().expect("flush error"); - - pb.finish_and_clear(); - - let mut rep = Reporter::new(); - rep.push(t.stop()); - rep.print(); -} - -// ── Output ──────────────────────────────────────────────────────────────────── - -fn emit_batch( - batch: &QueryBatch, - accs: &[SeqAcc], - meta: &obikindex::meta::IndexMeta, - count_missing: bool, - detail: bool, - cov: &[Vec>], - out: &mut impl Write, -) { - for (seq_idx, (id, seq)) in batch.ids.iter().zip(batch.seqs.iter()).enumerate() { - let acc = &accs[seq_idx]; - let mut ann = serde_json::Map::new(); - - ann.insert("kmer_count".into(), acc.kmer_count.into()); - if count_missing { - ann.insert("kmer_missing".into(), acc.kmer_missing.into()); - } - - let mut match_map = serde_json::Map::new(); - for (g, genome) in meta.genomes.iter().enumerate() { - if acc.genome_totals[g] != 0 { - match_map.insert(genome.label.clone(), acc.genome_totals[g].into()); - } - } - ann.insert("kmer_strict_matches".into(), match_map.into()); - - if detail && !cov.is_empty() { - let mut cov_map = serde_json::Map::new(); - for (g, genome) in meta.genomes.iter().enumerate() { - let v: Vec = cov[seq_idx][g].iter().map(|&x| x.into()).collect(); - cov_map.insert(genome.label.clone(), v.into()); - } - ann.insert("coverage".into(), cov_map.into()); - } - - // OBITools4 FASTA format: >id {"key":value,...} - let _ = out.write_all(b">"); - let _ = out.write_all(id.as_bytes()); - let _ = out.write_all(b" "); - let _ = serde_json::to_writer(&mut *out, &ann); - let _ = out.write_all(b"\n"); - let _ = out.write_all(seq); - let _ = out.write_all(b"\n"); - } -} - -#[cfg(test)] -#[path = "tests/query.rs"] -mod tests; diff --git a/src/obikmer/src/cmd/query/batch.rs b/src/obikmer/src/cmd/query/batch.rs new file mode 100644 index 00000000..55cd9452 --- /dev/null +++ b/src/obikmer/src/cmd/query/batch.rs @@ -0,0 +1,74 @@ +use std::collections::HashMap; + +use obikpartitionner::KmerDesc; +use obikseq::CanonicalKmer; +use obiread::record::SeqRecord; +use obiskbuilder::SuperKmerIter; + +/// A batch of query sequences, with k-mers deduplicated directly (not just at +/// the superkmer level) and pre-split by partition. +/// +/// Superkmer *construction* (`SuperKmerIter`) is still required — it's the +/// mechanism that computes minimizers and partition routing — but the dedup +/// key is the canonical k-mer, not the superkmer: two different superkmers +/// that happen to share a k-mer (read overlaps, repeats, a SNP splitting an +/// otherwise-identical run) are deduplicated too, not just identical whole +/// superkmers. This also means each unique k-mer triggers at most one MPHF +/// lookup, not one per occurrence. +pub struct QueryBatch { + /// Sequence ids in batch order. + pub ids: Vec, + /// Raw sequence bytes (for output), in batch order. + pub seqs: Vec>, + /// Total kmer count per sequence (used for `--detail` coverage allocation). + pub n_kmers: Vec, + /// Deduplicated k-mer occurrences, one map per partition. + pub by_partition: Vec>>, +} + +impl QueryBatch { + /// Build a batch from a vec of parsed sequence records, deduplicating + /// k-mers and routing them to partitions in the same pass. + pub fn from_records( + records: Vec, + k: usize, + level_max: usize, + theta: f64, + n_partitions: usize, + ) -> Self { + let mut ids = Vec::with_capacity(records.len()); + let mut seqs = Vec::with_capacity(records.len()); + let mut n_kmers = Vec::with_capacity(records.len()); + let mask = (n_partitions as u64) - 1; + let mut by_partition: Vec>> = + (0..n_partitions).map(|_| HashMap::new()).collect(); + + for (seq_idx, record) in records.into_iter().enumerate() { + let mut kmer_offset = 0u32; + + for rsk in SuperKmerIter::new(&record.normalized, k, level_max, theta) { + let part_idx = (rsk.minimizer().seq_hash() & mask) as usize; + let map = &mut by_partition[part_idx]; + for (j, kmer) in rsk.superkmer().iter_canonical_kmers().enumerate() { + map.entry(kmer).or_default().push(KmerDesc { + seq_idx: seq_idx as u32, + pos: kmer_offset + j as u32, + }); + } + let n = (rsk.seql() - k + 1) as u32; + kmer_offset += n; + } + + ids.push(record.id); + seqs.push(record.sequence); + n_kmers.push(kmer_offset); + } + + Self { + ids, + seqs, + n_kmers, + by_partition, + } + } +} diff --git a/src/obikmer/src/cmd/query/chunk.rs b/src/obikmer/src/cmd/query/chunk.rs new file mode 100644 index 00000000..509d84b1 --- /dev/null +++ b/src/obikmer/src/cmd/query/chunk.rs @@ -0,0 +1,288 @@ +use std::time::Instant; + +use obikindex::KmerIndex; +use obikpartitionner::{KmerDesc, QueryHit, QueryStats}; +use obikrope::Rope; +use obikseq::CanonicalKmer; +use obiread::record::parse_chunk; +use tracing::debug; + +use super::batch::QueryBatch; +use super::findere::{ConfirmedHit, sparse_findere_for_genome}; +use super::output::emit_batch; +use super::smer_index::SmerIndex; + +pub(super) struct SeqAcc { + pub(super) kmer_count: u32, + pub(super) kmer_missing: u32, + pub(super) genome_totals: Vec, +} + +impl SeqAcc { + fn new(n_genomes: usize) -> Self { + Self { + kmer_count: 0, + kmer_missing: 0, + genome_totals: vec![0u32; n_genomes], + } + } +} + +pub(super) fn process_chunk( + idx: &KmerIndex, + rope: Rope, + k: usize, + n_genomes: usize, + n_partitions: usize, + with_counts: bool, + effective_z: usize, + detail: bool, + count_missing: bool, + force_presence: bool, + presence_threshold: u32, +) -> Vec { + let chunk_start = Instant::now(); + let chunk_bytes = rope.len(); + + let records = parse_chunk(&rope, k); + if records.is_empty() { + return Vec::new(); + } + + let batch = QueryBatch::from_records(records, k, 6, 0.7, n_partitions); + let n_seqs = batch.ids.len(); + + // Estimate QueryBatch::by_partition's actual memory footprint: the + // k-mer-level dedup map (roadmap point 5) — one HashMap> per partition, sized by *unique* k-mers, not shrunk by + // dedup. On real workloads with a low intra-chunk duplication rate this + // can dwarf every other per-chunk structure, including the sparse + // Findere ones logged further down — unlike those, chunk_bytes's formula + // (run()) does not account for this at all today. Measured by allocated + // capacity, not logical length, to reflect real memory pressure + // (HashMap/Vec growth slack) — `by_partition` is alive for the entire + // process_chunk call (never drained, only iterated by reference), so + // this is its footprint for the whole chunk lifetime, not a transient. + let hashmap_slot_bytes = (std::mem::size_of::() + + std::mem::size_of::>() + + 1) as u64; // +1 ≈ hashbrown control byte per slot + let by_partition_map_bytes: u64 = batch + .by_partition + .iter() + .map(|m| m.capacity() as u64 * hashmap_slot_bytes) + .sum(); + let by_partition_desc_bytes: u64 = batch + .by_partition + .iter() + .flat_map(|m| m.values()) + .map(|v| v.capacity() as u64 * std::mem::size_of::() as u64) + .sum(); + let by_partition_bytes = by_partition_map_bytes + by_partition_desc_bytes; + + debug!( + n_unique_kmers_total = batch.by_partition.iter().map(|m| m.len() as u64).sum::(), + by_partition_map_bytes, + by_partition_desc_bytes, + by_partition_bytes, + chunk_bytes, + "by_partition memory retained" + ); + + // 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> = (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 + // entire justification for k-mer-level dereplication (see query.md, + // Future work point 5). If this ratio stays close to 1.0 on real data, + // dereplication isn't paying for itself and that should show up here. + let n_occurrences: u64 = batch.n_kmers.iter().map(|&n| n as u64).sum(); + let mut query_stats = QueryStats::default(); + + for (part_idx, kmers) in batch.by_partition.iter().enumerate() { + if kmers.is_empty() { + continue; + } + + let stats = idx.partition() + .query_partition_with( + part_idx, + kmers, + n_genomes, + with_counts, + |event| match event { + QueryHit::Found(descs) => { + for desc in descs { + smer_index.mark_found(desc.seq_idx as usize, desc.pos as usize); + } + } + QueryHit::Value(descs, g, v) => { + for desc in descs { + by_genome[g].push((desc.seq_idx, desc.pos, v)); + } + } + }, + ) + .unwrap_or_else(|e| { + eprintln!("query error on partition {part_idx}: {e}"); + std::process::exit(1); + }); + query_stats += stats; + } + + debug!( + n_occurrences, + n_unique_kmers = query_stats.n_unique_kmers, + n_mphf_calls = query_stats.n_mphf_calls, + n_hits = query_stats.n_hits, + n_columns_scanned = query_stats.n_columns_scanned, + n_col_get_calls = query_stats.n_col_get_calls, + "k-mer dedup + column-major fetch" + ); + + // ── Sparse Findere: per-genome run detection + sliding-window minimum ──── + // + // 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 = batch + .n_kmers + .iter() + .map(|&n| { + let n = n as usize; + 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::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" + ); + + // Actual bytes retained by the sparse hit structures (by_genome + + // confirmed_by_genome, both alive simultaneously at this point — see the + // chunk-size formula's comment in `run()`), by allocated capacity rather + // than logical length so this reflects real memory pressure including + // Vec growth slack. `empirical_multiplier` is directly comparable to + // BYTES_PER_KMER_PER_GENOME (`run()`) — the ratio a cluster run's logs + // need to judge whether that constant is over- or under-conservative for + // real data, instead of guessing. + const HIT_ENTRY_BYTES: u64 = std::mem::size_of::<(u32, u32, u32)>() as u64; + let by_genome_bytes: u64 = by_genome.iter().map(|v| v.capacity() as u64 * HIT_ENTRY_BYTES).sum(); + let confirmed_bytes: u64 = confirmed_by_genome.iter().map(|v| v.capacity() as u64 * HIT_ENTRY_BYTES).sum(); + let retained_bytes = by_genome_bytes + confirmed_bytes; + + debug!( + by_genome_bytes, + confirmed_bytes, + retained_bytes, + chunk_bytes, + empirical_multiplier = retained_bytes as f64 / chunk_bytes.max(1) as f64, + "sparse memory retained" + ); + + // ── Accumulate: genome totals (per genome, from confirmed hits) ────────── + let mut accs: Vec = (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>> = if detail { + n_kmers_out.iter().map(|&n| vec![vec![0u32; n]; n_genomes]).collect() + } else { + Vec::new() + }; + 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; + } + } + } + + // Capacity estimate: actual sequence + ID bytes, plus JSON overhead per record. + // JSON per record ≈ 50 fixed chars + ~20 per genome (label + count value) + 100 (overhead). + let seq_bytes: usize = batch.seqs.iter().map(|s| s.len()).sum(); + let id_bytes: usize = batch.ids.iter().map(|s| s.len()).sum(); + let cap = seq_bytes + id_bytes + n_seqs * (4 + 50 + n_genomes * 20) + 100; + let mut buf = Vec::with_capacity(cap); + emit_batch( + &batch, + &accs, + idx.meta(), + count_missing, + detail, + &cov, + &mut buf, + ); + + debug!( + chunk_bytes, + n_seqs, + n_smers = batch.n_kmers.iter().map(|&n| n as u64).sum::(), + wall_ms = chunk_start.elapsed().as_millis() as u64, + "process_chunk" + ); + + buf +} diff --git a/src/obikmer/src/cmd/query/findere.rs b/src/obikmer/src/cmd/query/findere.rs new file mode 100644 index 00000000..4564c6f8 --- /dev/null +++ b/src/obikmer/src/cmd/query/findere.rs @@ -0,0 +1,71 @@ +use std::collections::VecDeque; + +/// 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`. +pub(super) 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. +pub(super) fn sparse_findere_for_genome( + hits: &mut [(u32, u32, u32)], + z: usize, + presence: bool, + threshold: u32, +) -> (Vec, 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) +} diff --git a/src/obikmer/src/cmd/query/mod.rs b/src/obikmer/src/cmd/query/mod.rs new file mode 100644 index 00000000..425342b8 --- /dev/null +++ b/src/obikmer/src/cmd/query/mod.rs @@ -0,0 +1,340 @@ +mod batch; +mod chunk; +mod findere; +mod output; +mod smer_index; + +use std::io::{self, BufWriter, Write}; +use std::path::PathBuf; +use std::sync::Arc; +use std::sync::atomic::{AtomicU32, AtomicU64, Ordering}; +use std::time::Instant; + +use clap::Args; +use obikindex::KmerIndex; +use obikrope::Rope; +use obilayeredmap::IndexMode; +use obipipeline::{Throttled, ThrottleGuard, throttle}; +use obiread::chunk::read_sequence_chunks_sized; +use obisys::{Reporter, Stage, available_memory_bytes, spinner}; +use tracing::{debug, info}; + +use chunk::process_chunk; + +// ── Pipeline data ───────────────────────────────────────────────────────────── + +enum QueryData { + Path(Throttled), + Chunk(Rope), + Output(Vec), +} + +// SAFETY: Rope contains Cell which is !Sync, but pipeline items are owned +// exclusively through channels — no item is ever shared across threads. +unsafe impl Send for QueryData {} +unsafe impl Sync for QueryData {} + +// ── CLI ─────────────────────────────────────────────────────────────────────── + +#[derive(Args)] +pub struct QueryArgs { + /// Index directory + pub index: PathBuf, + + /// Input sequences (FASTA/FASTQ, optionally gzip-compressed) + #[arg(num_args = 1..)] + pub inputs: Vec, + + /// Report per-position coverage vectors per genome (adds "coverage" to JSON) + #[arg(long)] + pub detail: bool, + + /// Enable 1-mismatch approximate matching + #[arg(long)] + pub mismatch: bool, + + /// Count k-mers absent from the index (adds kmer_missing annotation) + #[arg(long)] + pub count_missing: bool, + + /// Report per-genome presence (0/1) instead of raw counts + #[arg(long)] + pub force_presence: bool, + + /// Minimum accumulated match count to declare a genome present (implies --force-presence) + #[arg(long, default_value_t = 1)] + pub presence_threshold: u32, + + /// Override the Findere z parameter from index metadata + #[arg(short = 'z', long)] + pub findere_z: Option, + + /// Number of worker threads + #[arg( + short = 'T', + long, + default_value_t = obisys::effective_parallelism() + )] + pub threads: usize, + + /// I/O chunk size in MiB (default: auto-sized from available RAM and thread count) + #[arg(long)] + pub chunk_size: Option, + + /// Maximum number of input files open simultaneously. + /// Defaults to threads/4 (minimum 1). Keep below the number of workers + /// to ensure CPU workers are always available for the transform stage. + #[arg(long)] + pub max_open_files: Option, +} + +impl QueryArgs { + pub fn effective_max_open(&self) -> usize { + self.max_open_files + .unwrap_or_else(|| (self.threads / 4).max(1)) + .max(1) + } +} + +// ── GuardedChunkIter — keeps the throttle slot guard alive until the file is exhausted ── + +/// Wraps a per-file `Rope` chunk iterator together with its `ThrottleGuard`, +/// so the guard (and the throttle slot it holds) is only released once the +/// file has been fully read — never earlier, never held past that point. +struct GuardedChunkIter { + inner: Box + Send>, + _guard: ThrottleGuard, + files_open: Arc, +} + +impl Iterator for GuardedChunkIter { + type Item = Rope; + fn next(&mut self) -> Option { + self.inner.next() + } +} + +impl Drop for GuardedChunkIter { + fn drop(&mut self) { + self.files_open.fetch_sub(1, Ordering::Relaxed); + } +} + +// ── Entry point ─────────────────────────────────────────────────────────────── + +pub fn run(args: QueryArgs) { + let idx = Arc::new(KmerIndex::open(&args.index).unwrap_or_else(|e| { + eprintln!("error opening index: {e}"); + std::process::exit(1); + })); + + let k = idx.kmer_size(); + let n_genomes = idx.meta().genomes.len(); + let n_partitions = idx.n_partitions(); + let with_counts = idx.meta().config.with_counts; + let n_workers = args.threads.max(1); + + // Chunk size: each chunk stays in memory for its entire processing lifetime. + // + // Per-chunk memory is no longer a dense n_genomes-wide buffer (removed in + // the sparse Findere rework, see process_chunk) — it now scales with + // *actual hit count*, not with total_kmers_in_chunk × n_genomes + // unconditionally. BYTES_PER_KMER_PER_GENOME below is therefore a + // pathological-case bound, not a typical-case estimate: it protects + // against a fully-dense hit pattern (every k-mer of the query matching + // every genome — a degenerate case, e.g. low-complexity input theta- + // filtering should mostly reject, or an index of near-duplicate genomes), + // where by_genome and confirmed_by_genome (process_chunk) both end up + // holding one (seq_idx, pos, value) entry — 3 × u32 = 12 bytes, vs. 4 + // bytes for the old dense encoding, where position was implicit in the + // array index — per (k-mer, genome) pair, and *coexist simultaneously* + // (by_genome isn't freed before confirmed_by_genome is built), for a + // worst case of ~24 bytes/pair before Vec growth slack. `cov` remains + // fully dense when --detail is set (unaffected by the sparse rework), + // still roughly doubling the n_genomes-scaled cost. + // + // For realistic, sparse hit patterns actual memory is far below this + // bound — see the "sparse memory retained" debug log in process_chunk, + // which reports the empirical bytes-per-raw-byte multiplier actually + // observed per chunk, directly comparable to BYTES_PER_KMER_PER_GENOME + // below. Tightening this constant for typical-case throughput (at the + // cost of pathological-case safety margin) is a deliberate tuning + // decision to make from that data, not something to guess at here. + // + // BASE_OVERHEAD approximates what scales with chunk_bytes alone, + // independent of n_genomes: the Rope itself, parsed SeqRecord sequence + + // normalised bytes, the superkmer dedup map, and the JSON output buffer. + // Like the n_genomes-scaled term, this is an estimate — validate against + // actual peak RSS (Stage::stop's `rss` in the summary table) on real + // workloads rather than trusting it blindly. + // + // We target ≤ 50 % of available RAM across all concurrent workers + // (SAFETY_FACTOR). + const BASE_OVERHEAD: u64 = 4; + const BYTES_PER_KMER_PER_GENOME: u64 = 8; // pathological-case bound — see comment above + const SAFETY_FACTOR: u64 = 2; + + let detail_factor: u64 = if args.detail { 2 } else { 1 }; + let overhead_multiplier = + BASE_OVERHEAD + n_genomes as u64 * BYTES_PER_KMER_PER_GENOME * detail_factor; + + let chunk_bytes = args + .chunk_size + .map(|mb| mb * 1024 * 1024) + .unwrap_or_else(|| { + let avail = available_memory_bytes(); + let computed = avail / (n_workers as u64 * overhead_multiplier * SAFETY_FACTOR); + computed.clamp(4 * 1024 * 1024, 256 * 1024 * 1024) as usize + }); + + debug!( + chunk_bytes, + n_genomes, + detail = args.detail, + overhead_multiplier, + estimated_peak_chunk_bytes = chunk_bytes as u64 * overhead_multiplier, + "chunk-size formula resolved" + ); + + let effective_z: usize = args + .findere_z + .unwrap_or_else(|| match idx.meta().config.evidence { + IndexMode::Approx { z, .. } | IndexMode::Hybrid { z, .. } => z as usize, + IndexMode::Exact => 1, + }); + + info!( + "query: k={k}, {} genome(s), with_counts={with_counts}, z={effective_z}, \ + mismatch={}, detail={}", + n_genomes, args.mismatch, args.detail + ); + + if args.mismatch { + eprintln!("warning: --mismatch not yet implemented, ignored"); + } + + let detail = args.detail; + let count_missing = args.count_missing; + let force_presence = args.force_presence; + let presence_threshold = args.presence_threshold; + + // Throttled iterator over input file paths: at most `effective_max_open()` + // files are open at once. Opening + decompressing + chunking each file is + // now a Flat pipeline stage, executed across the `n_workers` pool — not + // serialised in the pipe's dedicated source thread (see steps::scatter / + // cmd::superkmer for the same pattern applied to indexing). + info!("query: chunk_size={}MiB, max_open_files={}", chunk_bytes / (1024 * 1024), args.effective_max_open()); + + let paths: Vec = args.inputs.iter().map(PathBuf::from).collect(); + let path_source = throttle(paths.into_iter(), args.effective_max_open()); + + // Instrumentation: total bytes processed (for the EMA throughput readout), + // number of files currently open/being chunked, and number of chunks + // currently being processed by a worker — all read from the spinner loop + // below, updated from inside the pipe closures. + let total_bytes = Arc::new(AtomicU64::new(0)); + let files_open = Arc::new(AtomicU32::new(0)); + let chunks_active = Arc::new(AtomicU32::new(0)); + + let pipe = obipipeline::make_pipe! { + QueryData : Throttled => Vec, + || { + let files_open = Arc::clone(&files_open); + move |pw: Throttled| -> GuardedChunkIter { + let path = pw.item; + let guard = pw.guard; + let path_str = path.to_str().unwrap_or("").to_owned(); + files_open.fetch_add(1, Ordering::Relaxed); + let open_start = Instant::now(); + // Hard-exit on file-open failure (mirrors the previous behaviour): + // propagating this as a pipeline Err would hit a known scheduler + // hang on early stage errors (obipipeline::scheduler::WorkerPool::run + // breaks its main loop without unblocking the still-running source + // thread, so the final `h.join()` never returns) — worth fixing in + // obipipeline itself, but out of scope here; sidestepping it like the + // original code already did is the safe choice for this change. + let iter = read_sequence_chunks_sized(&path_str, chunk_bytes).unwrap_or_else(|e| { + eprintln!("error opening {path_str}: {e}"); + std::process::exit(1); + }); + debug!( + path = %path_str, + open_ms = open_start.elapsed().as_millis() as u64, + "opened query input file" + ); + let err_path = path_str.clone(); + GuardedChunkIter { + inner: Box::new(iter.filter_map(move |r| match r { + Ok(rope) => Some(rope), + Err(e) => { + eprintln!("read error: {err_path}: {e}"); + None + } + })), + _guard: guard, + files_open: Arc::clone(&files_open), + } + } + } : Path => Chunk, + | { + let idx = Arc::clone(&idx); + let total_bytes = Arc::clone(&total_bytes); + let chunks_active = Arc::clone(&chunks_active); + move |rope: Rope| { + chunks_active.fetch_add(1, Ordering::Relaxed); + let bytes = rope.len() as u64; + let out = process_chunk( + &idx, rope, k, n_genomes, n_partitions, with_counts, + effective_z, detail, count_missing, force_presence, presence_threshold, + ); + total_bytes.fetch_add(bytes, Ordering::Relaxed); + chunks_active.fetch_sub(1, Ordering::Relaxed); + out + } + } : Chunk => Output, + }; + + let t = Stage::start("query"); + let pb = spinner("query"); + + let mut ema_rate: f64 = 0.0; + let mut last_t = Instant::now(); + let mut last_bytes: u64 = 0; + const ALPHA: f64 = 0.15; + + let mut out = BufWriter::new(io::stdout()); + for block in pipe.apply(path_source, n_workers, 2) { + if !block.is_empty() { + out.write_all(&block).expect("write error"); + } + + let now = Instant::now(); + let dt = now.duration_since(last_t).as_secs_f64(); + if dt > 0.1 { + let total = total_bytes.load(Ordering::Relaxed); + let instant = (total - last_bytes) as f64 / dt; + ema_rate = ALPHA * instant + (1.0 - ALPHA) * ema_rate; + last_t = now; + last_bytes = total; + let bp = total as f64; + let (count_str, rate_str) = if bp >= 1e9 { + (format!("{:.2} GB", bp / 1e9), format!("{:.0} MB/s", ema_rate / 1e6)) + } else { + (format!("{:.0} MB", bp / 1e6), format!("{:.0} MB/s", ema_rate / 1e6)) + }; + let active = chunks_active.load(Ordering::Relaxed); + let open = files_open.load(Ordering::Relaxed); + pb.set_message(format!("{count_str} {rate_str} [files open: {open}, chunks in flight: {active}]")); + } + } + out.flush().expect("flush error"); + + pb.finish_and_clear(); + + let mut rep = Reporter::new(); + rep.push(t.stop()); + rep.print(); +} + +#[cfg(test)] +mod tests; diff --git a/src/obikmer/src/cmd/query/output.rs b/src/obikmer/src/cmd/query/output.rs new file mode 100644 index 00000000..2584c0ea --- /dev/null +++ b/src/obikmer/src/cmd/query/output.rs @@ -0,0 +1,50 @@ +use std::io::Write; + +use super::batch::QueryBatch; +use super::chunk::SeqAcc; + +pub(super) fn emit_batch( + batch: &QueryBatch, + accs: &[SeqAcc], + meta: &obikindex::meta::IndexMeta, + count_missing: bool, + detail: bool, + cov: &[Vec>], + out: &mut impl Write, +) { + for (seq_idx, (id, seq)) in batch.ids.iter().zip(batch.seqs.iter()).enumerate() { + let acc = &accs[seq_idx]; + let mut ann = serde_json::Map::new(); + + ann.insert("kmer_count".into(), acc.kmer_count.into()); + if count_missing { + ann.insert("kmer_missing".into(), acc.kmer_missing.into()); + } + + let mut match_map = serde_json::Map::new(); + for (g, genome) in meta.genomes.iter().enumerate() { + if acc.genome_totals[g] != 0 { + match_map.insert(genome.label.clone(), acc.genome_totals[g].into()); + } + } + ann.insert("kmer_strict_matches".into(), match_map.into()); + + if detail && !cov.is_empty() { + let mut cov_map = serde_json::Map::new(); + for (g, genome) in meta.genomes.iter().enumerate() { + let v: Vec = cov[seq_idx][g].iter().map(|&x| x.into()).collect(); + cov_map.insert(genome.label.clone(), v.into()); + } + ann.insert("coverage".into(), cov_map.into()); + } + + // OBITools4 FASTA format: >id {"key":value,...} + let _ = out.write_all(b">"); + let _ = out.write_all(id.as_bytes()); + let _ = out.write_all(b" "); + let _ = serde_json::to_writer(&mut *out, &ann); + let _ = out.write_all(b"\n"); + let _ = out.write_all(seq); + let _ = out.write_all(b"\n"); + } +} diff --git a/src/obikmer/src/cmd/query/smer_index.rs b/src/obikmer/src/cmd/query/smer_index.rs new file mode 100644 index 00000000..2c57e366 --- /dev/null +++ b/src/obikmer/src/cmd/query/smer_index.rs @@ -0,0 +1,41 @@ +/// 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`). +pub(super) struct SmerIndex { + in_index: Vec, // total_smers + offsets: Vec, // offsets[i]..offsets[i+1] = s-mer range for sequence i +} + +impl SmerIndex { + pub(super) 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); + for &n in n_kmers_per_seq { + total += n as usize; + offsets.push(total); + } + Self { + in_index: vec![false; total], + offsets, + } + } + + /// 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 reports per-genome values. + pub(super) fn mark_found(&mut self, seq: usize, kmer: usize) { + let abs = self.offsets[seq] + kmer; + self.in_index[abs] = true; + } + + #[inline] + pub(super) fn is_in_index(&self, seq: usize, kmer: usize) -> bool { + self.in_index[self.offsets[seq] + kmer] + } +} diff --git a/src/obikmer/src/cmd/tests/query.rs b/src/obikmer/src/cmd/query/tests.rs similarity index 98% rename from src/obikmer/src/cmd/tests/query.rs rename to src/obikmer/src/cmd/query/tests.rs index 6f5976e6..129069dd 100644 --- a/src/obikmer/src/cmd/tests/query.rs +++ b/src/obikmer/src/cmd/query/tests.rs @@ -1,4 +1,9 @@ -use super::*; +use obikrope::Rope; +use obikseq::CanonicalKmer; +use obiread::record::parse_chunk; + +use super::batch::QueryBatch; +use super::findere::sparse_findere_for_genome; const K: usize = 11; const M: usize = 5; diff --git a/src/obikmer/src/cmd/reindex.rs b/src/obikmer/src/cmd/reindex/mod.rs similarity index 100% rename from src/obikmer/src/cmd/reindex.rs rename to src/obikmer/src/cmd/reindex/mod.rs diff --git a/src/obikmer/src/cmd/select.rs b/src/obikmer/src/cmd/select/mod.rs similarity index 100% rename from src/obikmer/src/cmd/select.rs rename to src/obikmer/src/cmd/select/mod.rs diff --git a/src/obikmer/src/cmd/superkmer.rs b/src/obikmer/src/cmd/superkmer/mod.rs similarity index 100% rename from src/obikmer/src/cmd/superkmer.rs rename to src/obikmer/src/cmd/superkmer/mod.rs diff --git a/src/obikmer/src/cmd/unitig.rs b/src/obikmer/src/cmd/unitig/mod.rs similarity index 100% rename from src/obikmer/src/cmd/unitig.rs rename to src/obikmer/src/cmd/unitig/mod.rs diff --git a/src/obikmer/src/cmd/utils/maintenance.rs b/src/obikmer/src/cmd/utils/maintenance.rs new file mode 100644 index 00000000..c85f0087 --- /dev/null +++ b/src/obikmer/src/cmd/utils/maintenance.rs @@ -0,0 +1,112 @@ +use std::path::PathBuf; + +use obikindex::{validate_label, IndexBitsPerKmer, KmerIndex}; +use tracing::info; + +pub(super) fn run_stats(index_path: &PathBuf) { + let idx = KmerIndex::open(index_path).unwrap_or_else(|e| { + eprintln!("error opening index: {e}"); + std::process::exit(1); + }); + let (total, per_genome) = idx.genome_kmer_counts().unwrap_or_else(|e| { + eprintln!("error computing stats: {e}"); + std::process::exit(1); + }); + println!("genome,n_kmers"); + for (g, &n) in idx.meta().genomes.iter().zip(per_genome.iter()) { + println!("{},{}", g.label, n); + } + println!("total,{total}"); +} + +pub(super) fn run_bits_per_kmer(index_path: &PathBuf) { + let idx = KmerIndex::open(index_path).unwrap_or_else(|e| { + eprintln!("error opening index: {e}"); + std::process::exit(1); + }); + let stats: IndexBitsPerKmer = idx.bits_per_kmer().unwrap_or_else(|e| { + eprintln!("error computing bits/kmer: {e}"); + std::process::exit(1); + }); + println!("k-mers : {}", stats.n_kmers); + println!("genomes : {}", stats.n_genomes); + println!("mphf : {:6.2} bits/kmer", stats.mphf); + println!("evidence : {:6.2} bits/kmer", stats.evidence); + println!( + "matrix : {:6.2} bits/kmer ({:.2} bits/kmer/genome)", + stats.matrix, stats.matrix_per_genome + ); + println!("total : {:6.2} bits/kmer", stats.total); +} + +pub(super) fn run_upgrade_index(index_path: &PathBuf) { + let idx = KmerIndex::open(index_path).unwrap_or_else(|e| { + eprintln!("error opening index: {e}"); + std::process::exit(1); + }); + idx.upgrade_layer_meta().unwrap_or_else(|e| { + eprintln!("upgrade error: {e}"); + std::process::exit(1); + }); + info!("upgrade-index: layer_meta.json written to all layers that were missing it"); +} + +pub(super) fn run_rename(index_path: &PathBuf, spec: &str) { + let (old_label, new_label) = parse_rename_spec(spec); + + let mut idx = KmerIndex::open(index_path).unwrap_or_else(|e| { + eprintln!("error opening index: {e}"); + std::process::exit(1); + }); + + let pos = idx + .meta() + .genomes + .iter() + .position(|g| g.label == old_label) + .unwrap_or_else(|| { + eprintln!("error: genome '{old_label}' not found in index"); + std::process::exit(1); + }); + + validate_label(&new_label).unwrap_or_else(|e| { + eprintln!("error: --new-label: {e}"); + std::process::exit(1); + }); + + if idx.meta().genomes.iter().any(|g| g.label == new_label) { + eprintln!("error: label '{new_label}' already exists in index"); + std::process::exit(1); + } + + idx.meta_mut().genomes[pos].label = new_label.clone(); + idx.meta_mut().write(index_path).unwrap_or_else(|e| { + eprintln!("error writing index metadata: {e}"); + std::process::exit(1); + }); + + let spectrums_dir = index_path.join("spectrums"); + let old_spectrum = spectrums_dir.join(format!("{old_label}.json")); + let new_spectrum = spectrums_dir.join(format!("{new_label}.json")); + if old_spectrum.exists() { + std::fs::rename(&old_spectrum, &new_spectrum).unwrap_or_else(|e| { + eprintln!("warning: could not rename spectrum file: {e}"); + }); + } + + info!("renamed genome '{old_label}' → '{new_label}'"); +} + +fn parse_rename_spec(spec: &str) -> (String, String) { + let eq = spec.find('=').unwrap_or_else(|| { + eprintln!("error: --new-label expects NEW_LABEL=OLD_LABEL, got '{spec}'"); + std::process::exit(1); + }); + let new = spec[..eq].trim().to_string(); + let old = spec[eq + 1..].trim().to_string(); + if old.is_empty() || new.is_empty() { + eprintln!("error: --new-label: both old and new labels must be non-empty"); + std::process::exit(1); + } + (old, new) +} diff --git a/src/obikmer/src/cmd/utils/mod.rs b/src/obikmer/src/cmd/utils/mod.rs new file mode 100644 index 00000000..53efb3c4 --- /dev/null +++ b/src/obikmer/src/cmd/utils/mod.rs @@ -0,0 +1,85 @@ +mod maintenance; +mod partition_stats; + +use std::path::PathBuf; + +use clap::Args; + +use maintenance::{run_bits_per_kmer, run_stats, run_upgrade_index, run_rename}; +use partition_stats::run_partition_stats; + +#[derive(Args)] +pub struct UtilsArgs { + /// Index directories to operate on (one or more) + #[arg(required = true, num_args = 1..)] + pub indexes: Vec, + + /// Set a new genome label: NEW_LABEL=OLD_LABEL (single-index only) + #[arg(long, value_name = "NEW=OLD")] + pub new_label: Option, + + /// Add missing layer_meta.json files to each layer (single-index only) + #[arg(long)] + pub upgrade_index: bool, + + /// Print bits-per-kmer statistics (single-index only) + #[arg(long)] + pub bits_per_kmer: bool, + + /// Print per-genome k-mer counts as CSV (single-index only) + #[arg(long)] + pub stats: bool, + + /// Print partition size distribution report (accepts multiple indexes) + #[arg(long)] + pub partition_stats: bool, + + /// Write per-(partition, source) raw data as CSV to FILE (used with --partition-stats) + #[arg(long, value_name = "FILE")] + pub csv: Option, +} + +pub fn run(args: UtilsArgs) { + let mut any = false; + + if let Some(spec) = &args.new_label { + any = true; + run_rename(single_index(&args), spec); + } + + if args.upgrade_index { + any = true; + run_upgrade_index(single_index(&args)); + } + + if args.bits_per_kmer { + any = true; + run_bits_per_kmer(single_index(&args)); + } + + if args.stats { + any = true; + run_stats(single_index(&args)); + } + + if args.partition_stats { + any = true; + run_partition_stats(&args.indexes, args.csv.as_deref()); + } + + if !any { + eprintln!( + "utils: no operation specified. \ + Available: --new-label, --upgrade-index, --bits-per-kmer, --stats, --partition-stats" + ); + std::process::exit(1); + } +} + +fn single_index(args: &UtilsArgs) -> &PathBuf { + if args.indexes.len() > 1 { + eprintln!("utils: this option requires exactly one index (got {})", args.indexes.len()); + std::process::exit(1); + } + &args.indexes[0] +} diff --git a/src/obikmer/src/cmd/utils.rs b/src/obikmer/src/cmd/utils/partition_stats.rs similarity index 50% rename from src/obikmer/src/cmd/utils.rs rename to src/obikmer/src/cmd/utils/partition_stats.rs index f73c3085..1b53c322 100644 --- a/src/obikmer/src/cmd/utils.rs +++ b/src/obikmer/src/cmd/utils/partition_stats.rs @@ -1,89 +1,7 @@ use std::io::{self, Write}; use std::path::PathBuf; -use clap::Args; -use obikindex::{validate_label, IndexBitsPerKmer, KmerIndex}; -use tracing::info; - -#[derive(Args)] -pub struct UtilsArgs { - /// Index directories to operate on (one or more) - #[arg(required = true, num_args = 1..)] - pub indexes: Vec, - - /// Set a new genome label: NEW_LABEL=OLD_LABEL (single-index only) - #[arg(long, value_name = "NEW=OLD")] - pub new_label: Option, - - /// Add missing layer_meta.json files to each layer (single-index only) - #[arg(long)] - pub upgrade_index: bool, - - /// Print bits-per-kmer statistics (single-index only) - #[arg(long)] - pub bits_per_kmer: bool, - - /// Print per-genome k-mer counts as CSV (single-index only) - #[arg(long)] - pub stats: bool, - - /// Print partition size distribution report (accepts multiple indexes) - #[arg(long)] - pub partition_stats: bool, - - /// Write per-(partition, source) raw data as CSV to FILE (used with --partition-stats) - #[arg(long, value_name = "FILE")] - pub csv: Option, -} - -pub fn run(args: UtilsArgs) { - let mut any = false; - - if let Some(spec) = &args.new_label { - any = true; - run_rename(single_index(&args), spec); - } - - if args.upgrade_index { - any = true; - run_upgrade_index(single_index(&args)); - } - - if args.bits_per_kmer { - any = true; - run_bits_per_kmer(single_index(&args)); - } - - if args.stats { - any = true; - run_stats(single_index(&args)); - } - - if args.partition_stats { - any = true; - run_partition_stats(&args.indexes, args.csv.as_deref()); - } - - if !any { - eprintln!( - "utils: no operation specified. \ - Available: --new-label, --upgrade-index, --bits-per-kmer, --stats, --partition-stats" - ); - std::process::exit(1); - } -} - -// ── helpers ─────────────────────────────────────────────────────────────────── - -fn single_index(args: &UtilsArgs) -> &PathBuf { - if args.indexes.len() > 1 { - eprintln!("utils: this option requires exactly one index (got {})", args.indexes.len()); - std::process::exit(1); - } - &args.indexes[0] -} - -// ── --partition-stats ───────────────────────────────────────────────────────── +use obikindex::KmerIndex; /// Per-partition, per-source byte count of all unitigs.bin files summed across layers. struct PartRow { @@ -198,7 +116,7 @@ fn ascii_histogram(totals: &[u64], n_buckets: usize, bar_width: usize) -> String out } -fn run_partition_stats(indexes: &[PathBuf], csv_path: Option<&std::path::Path>) { +pub(super) fn run_partition_stats(indexes: &[PathBuf], csv_path: Option<&std::path::Path>) { let rows = collect_rows(indexes); if rows.is_empty() { eprintln!("partition-stats: no data found"); @@ -282,113 +200,3 @@ fn run_partition_stats(indexes: &[PathBuf], csv_path: Option<&std::path::Path>) eprintln!("CSV written to {}", csv_out.display()); } } - -// ── existing single-index operations ───────────────────────────────────────── - -fn run_stats(index_path: &PathBuf) { - let idx = KmerIndex::open(index_path).unwrap_or_else(|e| { - eprintln!("error opening index: {e}"); - std::process::exit(1); - }); - let (total, per_genome) = idx.genome_kmer_counts().unwrap_or_else(|e| { - eprintln!("error computing stats: {e}"); - std::process::exit(1); - }); - println!("genome,n_kmers"); - for (g, &n) in idx.meta().genomes.iter().zip(per_genome.iter()) { - println!("{},{}", g.label, n); - } - println!("total,{total}"); -} - -fn run_bits_per_kmer(index_path: &PathBuf) { - let idx = KmerIndex::open(index_path).unwrap_or_else(|e| { - eprintln!("error opening index: {e}"); - std::process::exit(1); - }); - let stats: IndexBitsPerKmer = idx.bits_per_kmer().unwrap_or_else(|e| { - eprintln!("error computing bits/kmer: {e}"); - std::process::exit(1); - }); - println!("k-mers : {}", stats.n_kmers); - println!("genomes : {}", stats.n_genomes); - println!("mphf : {:6.2} bits/kmer", stats.mphf); - println!("evidence : {:6.2} bits/kmer", stats.evidence); - println!( - "matrix : {:6.2} bits/kmer ({:.2} bits/kmer/genome)", - stats.matrix, stats.matrix_per_genome - ); - println!("total : {:6.2} bits/kmer", stats.total); -} - -fn run_upgrade_index(index_path: &PathBuf) { - let idx = KmerIndex::open(index_path).unwrap_or_else(|e| { - eprintln!("error opening index: {e}"); - std::process::exit(1); - }); - idx.upgrade_layer_meta().unwrap_or_else(|e| { - eprintln!("upgrade error: {e}"); - std::process::exit(1); - }); - info!("upgrade-index: layer_meta.json written to all layers that were missing it"); -} - -fn run_rename(index_path: &PathBuf, spec: &str) { - let (old_label, new_label) = parse_rename_spec(spec); - - let mut idx = KmerIndex::open(index_path).unwrap_or_else(|e| { - eprintln!("error opening index: {e}"); - std::process::exit(1); - }); - - let pos = idx - .meta() - .genomes - .iter() - .position(|g| g.label == old_label) - .unwrap_or_else(|| { - eprintln!("error: genome '{old_label}' not found in index"); - std::process::exit(1); - }); - - validate_label(&new_label).unwrap_or_else(|e| { - eprintln!("error: --new-label: {e}"); - std::process::exit(1); - }); - - if idx.meta().genomes.iter().any(|g| g.label == new_label) { - eprintln!("error: label '{new_label}' already exists in index"); - std::process::exit(1); - } - - idx.meta_mut().genomes[pos].label = new_label.clone(); - idx.meta_mut().write(index_path).unwrap_or_else(|e| { - eprintln!("error writing index metadata: {e}"); - std::process::exit(1); - }); - - let spectrums_dir = index_path.join("spectrums"); - let old_spectrum = spectrums_dir.join(format!("{old_label}.json")); - let new_spectrum = spectrums_dir.join(format!("{new_label}.json")); - if old_spectrum.exists() { - std::fs::rename(&old_spectrum, &new_spectrum).unwrap_or_else(|e| { - eprintln!("warning: could not rename spectrum file: {e}"); - }); - } - - info!("renamed genome '{old_label}' → '{new_label}'"); -} - -fn parse_rename_spec(spec: &str) -> (String, String) { - let eq = spec.find('=').unwrap_or_else(|| { - eprintln!("error: --new-label expects NEW_LABEL=OLD_LABEL, got '{spec}'"); - std::process::exit(1); - }); - let new = spec[..eq].trim().to_string(); - let old = spec[eq + 1..].trim().to_string(); - if old.is_empty() || new.is_empty() { - eprintln!("error: --new-label: both old and new labels must be non-empty"); - std::process::exit(1); - } - (old, new) -}