From bd7729b0953028b4423b2cd105e3ba3cbde9f30c Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Wed, 26 Aug 2026 19:26:29 +0200 Subject: [PATCH] refactor: Shift KmerIndex ownership to Arc for thread-safe sharing Replaces lifetime-bound references with runtime reference counting across multiple crates. This enables safe concurrent access across parallel workers without explicit cloning or manual lifetime management. Introduces the `query` and `utils` CLI commands in obikmer2, along with supporting modules for batch processing, sparse indexing, sliding-window findere logic, and output formatting. Updates dependency manifests and aligns test suites with the new ownership model. --- src/Cargo.lock | 12 + src/obikdump/src/dump.rs | 7 +- src/obikdump/src/unitig.rs | 5 +- src/obikfilter/src/filter_algo.rs | 7 +- src/obikfilter/src/filter_partition.rs | 5 +- src/obikfilter/src/partition_iter.rs | 2 +- src/obikidxcache/src/index_cache.rs | 32 +- src/obikmer2/Cargo.toml | 4 + src/obikmer2/src/cmd/dump/mod.rs | 5 +- src/obikmer2/src/cmd/filter/mod.rs | 7 +- src/obikmer2/src/cmd/mod.rs | 2 + src/obikmer2/src/cmd/pack/mod.rs | 5 +- src/obikmer2/src/cmd/query/mod.rs | 343 ++++++++++++++++++ src/obikmer2/src/cmd/unitig/mod.rs | 5 +- src/obikmer2/src/cmd/utils/maintenance.rs | 113 ++++++ src/obikmer2/src/cmd/utils/mod.rs | 76 ++++ src/obikmer2/src/cmd/utils/partition_stats.rs | 206 +++++++++++ src/obikmer2/src/main.rs | 6 + src/obikquery/Cargo.toml | 13 +- src/obikquery/src/batch.rs | 79 ++++ src/obikquery/src/chunk.rs | 279 ++++++++++++++ src/obikquery/src/findere.rs | 75 ++++ src/obikquery/src/lib.rs | 20 +- src/obikquery/src/output.rs | 52 +++ src/obikquery/src/query_layer.rs | 155 +++----- src/obikquery/src/smer_index.rs | 41 +++ src/obikquery/src/tests/batch.rs | 127 +++++++ src/obikquery/src/tests/findere.rs | 105 ++++++ src/obikquery/src/tests/query_layer.rs | 47 ++- src/obikrebuild/src/compact_layer.rs | 7 +- src/obikstats/src/stats.rs | 25 +- 31 files changed, 1689 insertions(+), 178 deletions(-) create mode 100644 src/obikmer2/src/cmd/query/mod.rs create mode 100644 src/obikmer2/src/cmd/utils/maintenance.rs create mode 100644 src/obikmer2/src/cmd/utils/mod.rs create mode 100644 src/obikmer2/src/cmd/utils/partition_stats.rs create mode 100644 src/obikquery/src/batch.rs create mode 100644 src/obikquery/src/chunk.rs create mode 100644 src/obikquery/src/findere.rs create mode 100644 src/obikquery/src/output.rs create mode 100644 src/obikquery/src/smer_index.rs create mode 100644 src/obikquery/src/tests/batch.rs create mode 100644 src/obikquery/src/tests/findere.rs diff --git a/src/Cargo.lock b/src/Cargo.lock index 89441f67..a2205fc6 100644 --- a/src/Cargo.lock +++ b/src/Cargo.lock @@ -1673,12 +1673,16 @@ dependencies = [ "obikalgorithm", "obikdump", "obikfilter", + "obikidxcache", "obikindex", "obikindexer", "obikmerge", + "obikquery", "obikrebuild", + "obikrope", "obikselect", "obikseq", + "obikstats", "obipipeline", "obiread", "obiskbuilder", @@ -1732,7 +1736,15 @@ dependencies = [ name = "obikquery" version = "0.1.0" dependencies = [ + "obikidxcache", "obikindex", + "obikrope", + "obikseq", + "obiread", + "obiskbuilder", + "serde_json", + "tempfile", + "tracing", ] [[package]] diff --git a/src/obikdump/src/dump.rs b/src/obikdump/src/dump.rs index 960a0772..b2dce101 100644 --- a/src/obikdump/src/dump.rs +++ b/src/obikdump/src/dump.rs @@ -1,4 +1,5 @@ use std::io::Write; +use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; use rayon::prelude::*; @@ -39,7 +40,7 @@ pub trait IndexDump { ) -> OKIResult<()>; } -impl IndexDump for KmerIndex { +impl IndexDump for Arc { fn dump( &self, out: &mut W, @@ -87,7 +88,7 @@ impl IndexDump for KmerIndex { Ok(_) => { write_row(buf, row, prefix); true } } }; - let cache = IndexCache::new(self, Some(vec![i])); + let cache = IndexCache::new(Arc::clone(self), Some(vec![i])); if debug { cache.iter_partition_kmers_located(i, use_counts, n_genomes, filters, |part, layer, kmer, row| { let seq = String::from_utf8(kmer.to_ascii()).unwrap_or_else(|_| "?".repeat(kmer_size)); @@ -106,7 +107,7 @@ impl IndexDump for KmerIndex { // ── Unbounded: no atomic, no contention ─────────────────────────── (0..n).into_par_iter().map(|i| { let mut buf = Vec::::new(); - let cache = IndexCache::new(self, Some(vec![i])); + let cache = IndexCache::new(Arc::clone(self), Some(vec![i])); if debug { cache.iter_partition_kmers_located(i, use_counts, n_genomes, filters, |part, layer, kmer, row| { let seq = String::from_utf8(kmer.to_ascii()).unwrap_or_else(|_| "?".repeat(kmer_size)); diff --git a/src/obikdump/src/unitig.rs b/src/obikdump/src/unitig.rs index b4a3c527..5f68b594 100644 --- a/src/obikdump/src/unitig.rs +++ b/src/obikdump/src/unitig.rs @@ -5,6 +5,7 @@ //! foreign type, so this is an extension trait rather than an inherent `impl`. use std::io::Write; +use std::sync::Arc; use obidebruinj::GraphDeBruijn; use obifastwrite::write_unitig; @@ -25,7 +26,7 @@ pub trait IndexUnitigs { ) -> OKIResult; } -impl IndexUnitigs for KmerIndex { +impl IndexUnitigs for Arc { fn write_unitigs( &self, out: &mut W, @@ -40,7 +41,7 @@ impl IndexUnitigs for KmerIndex { let g = (0..n) .into_par_iter() .try_fold(GraphDeBruijn::new, |mut local_g, i| -> OKIResult { - let cache = IndexCache::new(self, Some(vec![i])); + let cache = IndexCache::new(Arc::clone(self), Some(vec![i])); cache.iter_partition_kmers(i, use_counts, n_genomes, filters, |kmer, _row| { local_g.push(kmer); true diff --git a/src/obikfilter/src/filter_algo.rs b/src/obikfilter/src/filter_algo.rs index 2db2664f..0b202c95 100644 --- a/src/obikfilter/src/filter_algo.rs +++ b/src/obikfilter/src/filter_algo.rs @@ -9,6 +9,7 @@ use std::io; use std::path::PathBuf; +use std::sync::Arc; use obikalgorithm::Algorithm; use obikindex::{IndexBuilder, IndexState, KmerIndex}; @@ -19,7 +20,7 @@ use crate::filter::KmerFilter; use crate::filter_partition::filter_partition; pub struct Filter<'a> { - src: &'a KmerIndex, + src: Arc, output: PathBuf, filters: &'a [Box], presence: bool, @@ -30,7 +31,7 @@ pub struct Filter<'a> { } impl<'a> Filter<'a> { - pub fn new(src: &'a KmerIndex, output: impl Into, filters: &'a [Box]) -> Self { + pub fn new(src: Arc, output: impl Into, filters: &'a [Box]) -> Self { Self { src, output: output.into(), @@ -80,7 +81,7 @@ impl Algorithm for Filter<'_> { type Output = KmerIndex; fn run(&mut self) -> obikalgorithm::Result { - let src = self.src; + let src = &self.src; let output = self.output.clone(); if src.state()? != IndexState::Indexed { diff --git a/src/obikfilter/src/filter_partition.rs b/src/obikfilter/src/filter_partition.rs index 6c2c1957..88e923cb 100644 --- a/src/obikfilter/src/filter_partition.rs +++ b/src/obikfilter/src/filter_partition.rs @@ -5,6 +5,7 @@ //! copied — removing kmers changes MPHF slot assignment layer-wide. use std::collections::HashMap; +use std::sync::Arc; use obidebruinj::GraphDeBruijn; use obikidxcache::index_cache::IndexCache; @@ -25,7 +26,7 @@ use crate::partition_iter::FilteredPartitionIter; /// output from a count source. pub(crate) fn filter_partition( dst: &KmerIndex, - src: &KmerIndex, + src: &Arc, i: usize, filters: &[Box], use_counts: bool, @@ -35,7 +36,7 @@ pub(crate) fn filter_partition( ) -> OKIResult<()> { let n_genomes = src.meta().genomes().map_err(OKIError::Io)?.len(); - let cache = IndexCache::new(src, Some(vec![i])); + let cache = IndexCache::new(Arc::clone(src), Some(vec![i])); let dst_partition = dst.partition(i)?; let n_layers = src.partition(i)?.n_layers(); diff --git a/src/obikfilter/src/partition_iter.rs b/src/obikfilter/src/partition_iter.rs index 6e4d640e..0d69f570 100644 --- a/src/obikfilter/src/partition_iter.rs +++ b/src/obikfilter/src/partition_iter.rs @@ -52,7 +52,7 @@ pub trait FilteredPartitionIter { ) -> OKIResult; } -impl FilteredPartitionIter for IndexCache<'_> { +impl FilteredPartitionIter for IndexCache { fn iter_partition_kmers( &self, part: usize, diff --git a/src/obikidxcache/src/index_cache.rs b/src/obikidxcache/src/index_cache.rs index f4ed54d0..c5ef6661 100644 --- a/src/obikidxcache/src/index_cache.rs +++ b/src/obikidxcache/src/index_cache.rs @@ -1,18 +1,23 @@ use std::collections::HashMap; +use std::sync::Arc; use obikindex::{KmerIndex, layer::KmerLayer}; use obikseq::CanonicalKmer; use crate::meta_cache::MetaCache; -pub struct IndexCache<'a> { - /// Ties this cache's lifetime to the `KmerIndex` it was opened from — - /// borrow-checker-enforced, not just documentation: an `IndexCache` - /// cannot outlive its index, so it can never be read against an index - /// that has since been mutated or dropped. `index()` is a genuine - /// accessor, but the field's real job is this lifetime bound, kept even - /// on paths that never call `index()`. - raw_index: &'a KmerIndex, +pub struct IndexCache { + /// Owns a strong reference to the `KmerIndex` it was opened from — the + /// cache cannot outlive its index (the `Arc` keeps it alive at least as + /// long as the cache itself, enforced at runtime by refcounting rather + /// than by a borrow-checked lifetime). `Arc`, not `&'a KmerIndex`, + /// because `IndexCache` must itself be `'static`-capable to be shared + /// across `obipipeline`'s `thread::spawn`-based workers (see + /// `obikquery::query_layer`) — a plain borrow can never satisfy that, + /// regardless of how it's threaded through. `index()` is a genuine + /// accessor, but the field's real job is this ownership guarantee, kept + /// even on paths that never call `index()`. + raw_index: Arc, meta: MetaCache, /// Keyed by each partition's own index number, not renumbered — @@ -24,7 +29,7 @@ pub struct IndexCache<'a> { layer_cache: HashMap>, } -impl<'a> IndexCache<'a> { +impl IndexCache { /// Opens every layer of the given `partitions` (every partition of the /// index, if `None`) once, up front, and keeps them alive for the /// cache's lifetime — each partition's own layer count is read @@ -41,7 +46,7 @@ impl<'a> IndexCache<'a> { /// cache can only be built on partitions that actually exist and are /// complete, so any failure here is a caller error, not a case to /// design around with a `Result`. - pub fn new(index: &'a KmerIndex, partitions: Option>) -> Self { + pub fn new(index: Arc, partitions: Option>) -> Self { let selected = partitions.unwrap_or_else(|| (0..index.n_partitions()).collect()); let mut layer_cache = HashMap::with_capacity(selected.len()); @@ -62,17 +67,18 @@ impl<'a> IndexCache<'a> { layer_cache.insert(p, layers); } + let meta = MetaCache::from_meta(&index.meta()); IndexCache { raw_index: index, - meta: MetaCache::from_meta(&index.meta()), + meta, layer_cache, } } /// The index this cache was opened from. #[inline] - pub fn index(&self) -> &'a KmerIndex { - self.raw_index + pub fn index(&self) -> &KmerIndex { + &self.raw_index } /// The index's metadata, snapshotted once at construction — see diff --git a/src/obikmer2/Cargo.toml b/src/obikmer2/Cargo.toml index 9cc1f074..97bdacd9 100644 --- a/src/obikmer2/Cargo.toml +++ b/src/obikmer2/Cargo.toml @@ -20,6 +20,10 @@ obikfilter = { path = "../obikfilter" } obikselect = { path = "../obikselect" } obikdump = { path = "../obikdump" } obikrebuild = { path = "../obikrebuild" } +obikstats = { path = "../obikstats" } +obikquery = { path = "../obikquery" } +obikidxcache = { path = "../obikidxcache" } +obikrope = { path = "../obikrope" } obifastwrite = { path = "../obifastwrite" } obiskbuilder = { path = "../obiskbuilder" } clap = { version = "4", features = ["derive"] } diff --git a/src/obikmer2/src/cmd/dump/mod.rs b/src/obikmer2/src/cmd/dump/mod.rs index abe28a2b..63bcfebb 100644 --- a/src/obikmer2/src/cmd/dump/mod.rs +++ b/src/obikmer2/src/cmd/dump/mod.rs @@ -1,5 +1,6 @@ use std::io::{self, BufWriter}; use std::path::PathBuf; +use std::sync::Arc; use clap::Args; use obikdump::IndexDump; @@ -32,10 +33,10 @@ pub struct DumpArgs { } pub fn run(args: DumpArgs) { - let idx = KmerIndex::open(&args.index).unwrap_or_else(|e| { + let idx = Arc::new(KmerIndex::open(&args.index).unwrap_or_else(|e| { eprintln!("error opening index: {e}"); std::process::exit(1); - }); + })); let n_genomes = idx.meta().genomes().unwrap_or_else(|e| { eprintln!("error reading index metadata: {e}"); diff --git a/src/obikmer2/src/cmd/filter/mod.rs b/src/obikmer2/src/cmd/filter/mod.rs index b54d6c70..831f4b51 100644 --- a/src/obikmer2/src/cmd/filter/mod.rs +++ b/src/obikmer2/src/cmd/filter/mod.rs @@ -1,4 +1,5 @@ use std::path::PathBuf; +use std::sync::Arc; use clap::Args; use obikalgorithm::Algorithm; @@ -51,10 +52,10 @@ pub struct FilterArgs { } pub fn run(args: FilterArgs) { - let src = KmerIndex::open(&args.source).unwrap_or_else(|e| { + let src = Arc::new(KmerIndex::open(&args.source).unwrap_or_else(|e| { eprintln!("error opening source index: {e}"); std::process::exit(1); - }); + })); let mut filters: Vec> = vec![Box::new(args.group_filter.build_filter(&src.meta()))]; @@ -80,7 +81,7 @@ pub fn run(args: FilterArgs) { let mut rep = Reporter::new(); let t = Stage::start("filter"); let pb = progress_bar("filter", src.n_partitions() as u64, "partitions"); - let mut alg = Filter::new(&src, &args.output, &filters) + let mut alg = Filter::new(Arc::clone(&src), &args.output, &filters) .presence(args.presence) .force(args.force) .sparse(!args.dense) diff --git a/src/obikmer2/src/cmd/mod.rs b/src/obikmer2/src/cmd/mod.rs index 0baca15c..e4b61300 100644 --- a/src/obikmer2/src/cmd/mod.rs +++ b/src/obikmer2/src/cmd/mod.rs @@ -6,6 +6,8 @@ pub mod index; pub mod merge; pub mod pack; mod predicate; +pub mod query; pub mod select; pub mod superkmer; pub mod unitig; +pub mod utils; diff --git a/src/obikmer2/src/cmd/pack/mod.rs b/src/obikmer2/src/cmd/pack/mod.rs index 1dbcac3a..cdf85ba0 100644 --- a/src/obikmer2/src/cmd/pack/mod.rs +++ b/src/obikmer2/src/cmd/pack/mod.rs @@ -1,4 +1,5 @@ use std::path::PathBuf; +use std::sync::Arc; use clap::Args; use obikindex::KmerIndex; @@ -32,10 +33,10 @@ pub fn run(args: PackArgs) { std::process::exit(1); }); - let idx = KmerIndex::open(&args.index).unwrap_or_else(|e| { + let idx = Arc::new(KmerIndex::open(&args.index).unwrap_or_else(|e| { eprintln!("error opening index: {e}"); std::process::exit(1); - }); + })); let n_genomes = idx.meta().genomes().unwrap_or_else(|e| { eprintln!("error reading index metadata: {e}"); diff --git a/src/obikmer2/src/cmd/query/mod.rs b/src/obikmer2/src/cmd/query/mod.rs new file mode 100644 index 00000000..a86a6583 --- /dev/null +++ b/src/obikmer2/src/cmd/query/mod.rs @@ -0,0 +1,343 @@ +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 obikidxcache::index_cache::IndexCache; +use obikindex::KmerIndex; +use obikindex::layer::IndexMode; +use obikquery::process_chunk; +use obikrope::Rope; +use obipipeline::{Throttled, ThrottleGuard, throttle}; +use obiread::chunk::read_sequence_chunks_sized; +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) + } +} + +// ── 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 genomes = idx.meta().genomes().unwrap_or_else(|e| { + eprintln!("error reading index metadata: {e}"); + std::process::exit(1); + }); + let n_genomes = genomes.len(); + let genomes = Arc::new(genomes); + let n_partitions = idx.n_partitions(); + let with_counts = idx.meta().config.with_counts; + let n_workers = args.threads.max(1); + + // Every partition/layer the query might touch is opened once, up front, + // and shared (via Arc) across every `obipipeline` worker — a query pass + // is then pure in-memory lookups, never a per-chunk disk open (the + // previous `obikindex`-based design's cost). `IndexCache` owns its + // `Arc`, so it's itself `'static`-capable, satisfying + // obipipeline's `Send + Sync + 'static` requirement on pipeline data — + // see `obikquery::query_layer`'s doc comment for why that rules out a + // borrow-based cache here. + let cache = Arc::new(IndexCache::new(Arc::clone(&idx), None)); + + // Chunk size: each chunk stays in memory for its entire processing lifetime. + // + // Per-chunk memory is not a dense n_genomes-wide buffer — it 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 (obikquery::chunk::process_chunk) + // both end up holding one (seq_idx, pos, value) entry — 3 × u32 = 12 + // bytes — 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, 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 + // a Flat pipeline stage, executed across the `n_workers` pool. + 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 cache = Arc::clone(&cache); + let genomes = Arc::clone(&genomes); + 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( + &cache, rope, k, n_genomes, n_partitions, with_counts, + effective_z, detail, count_missing, force_presence, presence_threshold, + &genomes, + ); + 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(); +} diff --git a/src/obikmer2/src/cmd/unitig/mod.rs b/src/obikmer2/src/cmd/unitig/mod.rs index 8a725b48..cc1baf34 100644 --- a/src/obikmer2/src/cmd/unitig/mod.rs +++ b/src/obikmer2/src/cmd/unitig/mod.rs @@ -1,5 +1,6 @@ use std::io::{self, BufWriter}; use std::path::PathBuf; +use std::sync::Arc; use clap::Args; use obikdump::IndexUnitigs; @@ -20,10 +21,10 @@ pub struct UnitigArgs { } pub fn run(args: UnitigArgs) { - let idx = KmerIndex::open(&args.index).unwrap_or_else(|e| { + let idx = Arc::new(KmerIndex::open(&args.index).unwrap_or_else(|e| { eprintln!("error opening index: {e}"); std::process::exit(1); - }); + })); info!( "unitig: building de Bruijn graph from {} partition(s) (k={})", diff --git a/src/obikmer2/src/cmd/utils/maintenance.rs b/src/obikmer2/src/cmd/utils/maintenance.rs new file mode 100644 index 00000000..d360131f --- /dev/null +++ b/src/obikmer2/src/cmd/utils/maintenance.rs @@ -0,0 +1,113 @@ +use std::path::PathBuf; +use std::sync::Arc; + +use obikalgorithm::Algorithm; +use obikindex::{GenomeInfo, KmerIndex}; +use obikstats::{BitsPerKmer, GenomeKmerCounts}; +use tracing::info; + +pub(super) fn run_stats(index_path: &PathBuf) { + let idx = Arc::new(KmerIndex::open(index_path).unwrap_or_else(|e| { + eprintln!("error opening index: {e}"); + std::process::exit(1); + })); + let genomes = idx.meta().genomes().unwrap_or_else(|e| { + eprintln!("error reading index metadata: {e}"); + std::process::exit(1); + }); + + let (total, per_genome) = GenomeKmerCounts::new(Arc::clone(&idx)).run().unwrap_or_else(|e| { + eprintln!("error computing stats: {e}"); + std::process::exit(1); + }); + + println!("genome,n_kmers"); + for (g, &n) in 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 = Arc::new(KmerIndex::open(index_path).unwrap_or_else(|e| { + eprintln!("error opening index: {e}"); + std::process::exit(1); + })); + + let stats = BitsPerKmer::new(idx).run().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_rename(index_path: &PathBuf, spec: &str) { + let (old_label, new_label) = parse_rename_spec(spec); + + let idx = KmerIndex::open(index_path).unwrap_or_else(|e| { + eprintln!("error opening index: {e}"); + std::process::exit(1); + }); + + let genomes = idx.meta().genomes().unwrap_or_else(|e| { + eprintln!("error reading index metadata: {e}"); + std::process::exit(1); + }); + + let pos = genomes + .iter() + .position(|g| g.label == old_label) + .unwrap_or_else(|| { + eprintln!("error: genome '{old_label}' not found in index"); + std::process::exit(1); + }); + + GenomeInfo::validate_label(&new_label).unwrap_or_else(|e| { + eprintln!("error: --new-label: {e}"); + std::process::exit(1); + }); + + if genomes.iter().any(|g| g.label == new_label) { + eprintln!("error: label '{new_label}' already exists in index"); + std::process::exit(1); + } + + idx.meta().rename_genome(pos, new_label.clone()).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/obikmer2/src/cmd/utils/mod.rs b/src/obikmer2/src/cmd/utils/mod.rs new file mode 100644 index 00000000..fdba755b --- /dev/null +++ b/src/obikmer2/src/cmd/utils/mod.rs @@ -0,0 +1,76 @@ +mod maintenance; +mod partition_stats; + +use std::path::PathBuf; + +use clap::Args; + +use maintenance::{run_bits_per_kmer, run_stats, 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, + + /// 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.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, --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/obikmer2/src/cmd/utils/partition_stats.rs b/src/obikmer2/src/cmd/utils/partition_stats.rs new file mode 100644 index 00000000..6cf6de24 --- /dev/null +++ b/src/obikmer2/src/cmd/utils/partition_stats.rs @@ -0,0 +1,206 @@ +use std::io::{self, Write}; +use std::path::PathBuf; + +use obikindex::KmerIndex; + +/// Per-partition, per-source byte count of all unitigs.bin files summed across layers. +struct PartRow { + partition: usize, + source: String, + bytes: u64, +} + +fn collect_rows(indexes: &[PathBuf]) -> Vec { + let mut rows = Vec::new(); + for path in indexes { + let idx = KmerIndex::open(path).unwrap_or_else(|e| { + eprintln!("error opening index {}: {e}", path.display()); + std::process::exit(1); + }); + let name = path + .file_name() + .map(|n| n.to_string_lossy().into_owned()) + .unwrap_or_else(|| path.display().to_string()); + let n_parts = idx.n_partitions(); + for i in 0..n_parts { + let mut bytes = 0u64; + let n_layers = idx.n_layers(i).unwrap_or_else(|e| { + eprintln!("error reading partition {i} of {}: {e}", path.display()); + std::process::exit(1); + }); + for l in 0..n_layers { + let p = idx.layer_unitigs_path(i, l).unwrap_or_else(|e| { + eprintln!("error reading layer {l} of partition {i} of {}: {e}", path.display()); + std::process::exit(1); + }); + if let Ok(m) = std::fs::metadata(&p) { + bytes += m.len(); + } + } + rows.push(PartRow { partition: i, source: name.clone(), bytes }); + } + } + rows +} + +/// Sum bytes per partition across all sources. +fn partition_totals(rows: &[PartRow], n_parts: usize) -> Vec { + let mut totals = vec![0u64; n_parts]; + for r in rows { + totals[r.partition] += r.bytes; + } + totals +} + +fn stats_summary(totals: &[u64]) -> (u64, u64, f64, f64, u64, u64, u64) { + let mut sorted = totals.to_vec(); + sorted.sort_unstable(); + let n = sorted.len(); + let min = sorted[0]; + let max = sorted[n - 1]; + let mean = sorted.iter().sum::() as f64 / n as f64; + let median = if n % 2 == 0 { + (sorted[n / 2 - 1] + sorted[n / 2]) as f64 / 2.0 + } else { + sorted[n / 2] as f64 + }; + let p95 = sorted[(n as f64 * 0.95) as usize]; + let p99 = sorted[(n as f64 * 0.99) as usize]; + let variance = sorted + .iter() + .map(|&v| (v as f64 - mean).powi(2)) + .sum::() + / n as f64; + let std_dev = variance.sqrt(); + (min, max, mean, median, p95, p99, std_dev as u64) +} + +fn human_bytes(b: u64) -> String { + if b >= 1 << 30 { + format!("{:.1} GB", b as f64 / (1u64 << 30) as f64) + } else if b >= 1 << 20 { + format!("{:.1} MB", b as f64 / (1u64 << 20) as f64) + } else if b >= 1 << 10 { + format!("{:.1} KB", b as f64 / (1u64 << 10) as f64) + } else { + format!("{b} B") + } +} + +fn ascii_histogram(totals: &[u64], n_buckets: usize, bar_width: usize) -> String { + let min = *totals.iter().min().unwrap(); + let max = *totals.iter().max().unwrap(); + if min == max { + return format!(" (all partitions identical: {})\n", human_bytes(min)); + } + + let bucket_size = (max - min).max(1) as f64 / n_buckets as f64; + let mut counts = vec![0usize; n_buckets]; + for &v in totals { + let b = (((v - min) as f64 / bucket_size) as usize).min(n_buckets - 1); + counts[b] += 1; + } + + let max_count = *counts.iter().max().unwrap(); + let mut out = String::new(); + for (i, &c) in counts.iter().enumerate() { + let lo = min + (i as f64 * bucket_size) as u64; + let hi = min + ((i + 1) as f64 * bucket_size) as u64; + let bar_len = if max_count > 0 { c * bar_width / max_count } else { 0 }; + let bar = "█".repeat(bar_len); + out.push_str(&format!( + " {:>8} – {:>8} │{:) { + let rows = collect_rows(indexes); + if rows.is_empty() { + eprintln!("partition-stats: no data found"); + std::process::exit(1); + } + + let n_parts = rows.iter().map(|r| r.partition).max().unwrap() + 1; + let totals = partition_totals(&rows, n_parts); + let (min, max, mean, median, p95, p99, std_dev) = stats_summary(&totals); + + // outliers: > median + 1.5 × IQR (approximate via > 1.5 × median as fallback) + let mut sorted_t = totals.clone(); + sorted_t.sort_unstable(); + let q1 = sorted_t[n_parts / 4] as f64; + let q3 = sorted_t[3 * n_parts / 4] as f64; + let iqr = q3 - q1; + let outlier_threshold = q3 + 1.5 * iqr; + + let mut out = String::new(); + out.push_str("# Partition size report\n\n"); + out.push_str(&format!( + "Sources: {} \nPartitions: {} \n\n", + indexes.len(), + n_parts + )); + + out.push_str("## Summary statistics (total unitigs.bin bytes per partition, sum across sources)\n\n"); + out.push_str("| Stat | Value |\n|---|---|\n"); + out.push_str(&format!("| min | {} |\n", human_bytes(min))); + out.push_str(&format!("| max | {} |\n", human_bytes(max))); + out.push_str(&format!("| mean | {} |\n", human_bytes(mean as u64))); + out.push_str(&format!("| median | {} |\n", human_bytes(median as u64))); + out.push_str(&format!("| p95 | {} |\n", human_bytes(p95))); + out.push_str(&format!("| p99 | {} |\n", human_bytes(p99))); + out.push_str(&format!("| std | {} |\n", human_bytes(std_dev))); + out.push_str(&format!("| max/median ratio | {:.2}× |\n\n", max as f64 / median)); + + out.push_str("## Histogram\n\n```\n"); + out.push_str(&ascii_histogram(&totals, 30, 40)); + out.push_str("```\n\n"); + + let outliers: Vec<(usize, u64)> = totals + .iter() + .enumerate() + .filter(|(_, v)| **v as f64 > outlier_threshold) + .map(|(i, v)| (i, *v)) + .collect(); + + if outliers.is_empty() { + out.push_str("## Outliers\n\nNone (threshold: Q3 + 1.5×IQR = "); + out.push_str(&human_bytes(outlier_threshold as u64)); + out.push_str(").\n"); + } else { + out.push_str(&format!( + "## Outliers (> Q3 + 1.5×IQR = {})\n\n| Partition | Total size | Ratio to median |\n|---|---|---|\n", + human_bytes(outlier_threshold as u64) + )); + for (i, v) in &outliers { + out.push_str(&format!( + "| {} | {} | {:.2}× |\n", + i, + human_bytes(*v), + *v as f64 / median + )); + } + out.push('\n'); + } + + print!("{out}"); + + if let Some(csv_out) = csv_path { + let file = std::fs::File::create(csv_out).unwrap_or_else(|e| { + eprintln!("error creating CSV file {}: {e}", csv_out.display()); + std::process::exit(1); + }); + let mut w = io::BufWriter::new(file); + writeln!(w, "partition,source,bytes").unwrap(); + for r in &rows { + writeln!(w, "{},{},{}", r.partition, r.source, r.bytes).unwrap(); + } + eprintln!("CSV written to {}", csv_out.display()); + } +} diff --git a/src/obikmer2/src/main.rs b/src/obikmer2/src/main.rs index 7dcb2d4e..a1fec405 100644 --- a/src/obikmer2/src/main.rs +++ b/src/obikmer2/src/main.rs @@ -25,6 +25,8 @@ enum Commands { Select(cmd::select::SelectArgs), /// Dump an index's kmers as a CSV table Dump(cmd::dump::DumpArgs), + /// Query sequences against an index, annotating each with per-genome matches + Query(cmd::query::QueryArgs), /// Assemble an index's kmers into unitigs and write them as FASTA Unitig(cmd::unitig::UnitigArgs), /// Pack an index's matrices into single-file, sparse format by default (--dense to opt out), in place @@ -33,6 +35,8 @@ enum Commands { Estimate(cmd::estimate::EstimateArgs), /// Read/write genome metadata (CSV) on an already-built index Annotate(cmd::annotate::AnnotateArgs), + /// Maintenance/inspection operations on already-built indexes + Utils(cmd::utils::UtilsArgs), } fn main() { @@ -51,9 +55,11 @@ fn main() { Commands::Filter(args) => cmd::filter::run(args), Commands::Select(args) => cmd::select::run(args), Commands::Dump(args) => cmd::dump::run(args), + Commands::Query(args) => cmd::query::run(args), Commands::Unitig(args) => cmd::unitig::run(args), Commands::Pack(args) => cmd::pack::run(args), Commands::Estimate(args) => cmd::estimate::run(args), Commands::Annotate(args) => cmd::annotate::run(args), + Commands::Utils(args) => cmd::utils::run(args), } } diff --git a/src/obikquery/Cargo.toml b/src/obikquery/Cargo.toml index df532dc6..685c50ca 100644 --- a/src/obikquery/Cargo.toml +++ b/src/obikquery/Cargo.toml @@ -4,4 +4,15 @@ version = "0.1.0" edition = "2024" [dependencies] -obikindex = { path = "../obikindex" } +obikindex = { path = "../obikindex" } +obikidxcache = { path = "../obikidxcache" } +obikseq = { path = "../obikseq" } +obikrope = { path = "../obikrope" } +obiread = { path = "../obiread" } +obiskbuilder = { path = "../obiskbuilder" } +serde_json = "1" +tracing = "0.1.44" + +[dev-dependencies] +obikseq = { path = "../obikseq", features = ["test-utils"] } +tempfile = "3" diff --git a/src/obikquery/src/batch.rs b/src/obikquery/src/batch.rs new file mode 100644 index 00000000..1d56e86f --- /dev/null +++ b/src/obikquery/src/batch.rs @@ -0,0 +1,79 @@ +use std::collections::HashMap; + +use obikseq::CanonicalKmer; +use obiread::record::SeqRecord; +use obiskbuilder::SuperKmerIter; + +use crate::query_layer::KmerDesc; + +/// 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(crate) struct QueryBatch { + /// Sequence ids in batch order. + pub(crate) ids: Vec, + /// Raw sequence bytes (for output), in batch order. + pub(crate) seqs: Vec>, + /// Total kmer count per sequence (used for `--detail` coverage allocation). + pub(crate) n_kmers: Vec, + /// Deduplicated k-mer occurrences, one map per partition. + pub(crate) 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(crate) 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, + } + } +} + +#[cfg(test)] +#[path = "tests/batch.rs"] +mod tests; diff --git a/src/obikquery/src/chunk.rs b/src/obikquery/src/chunk.rs new file mode 100644 index 00000000..4fc4383b --- /dev/null +++ b/src/obikquery/src/chunk.rs @@ -0,0 +1,279 @@ +use std::time::Instant; + +use obikidxcache::index_cache::IndexCache; +use obikindex::GenomeInfo; +use obikrope::Rope; +use obikseq::CanonicalKmer; +use obiread::record::parse_chunk; +use tracing::debug; + +use crate::batch::QueryBatch; +use crate::findere::{ConfirmedHit, sparse_findere_for_genome}; +use crate::output::emit_batch; +use crate::query_layer::{QueryHit, QueryPartition, QueryStats}; +use crate::smer_index::SmerIndex; + +pub(crate) struct SeqAcc { + pub(crate) kmer_count: u32, + pub(crate) kmer_missing: u32, + pub(crate) genome_totals: Vec, +} + +impl SeqAcc { + fn new(n_genomes: usize) -> Self { + Self { + kmer_count: 0, + kmer_missing: 0, + genome_totals: vec![0u32; n_genomes], + } + } +} + +/// Turn one chunk of raw sequence bytes into an obitools-style annotated +/// FASTA byte buffer, matching its k-mers against `cache`'s index. +/// +/// `cache` is expected to already hold every partition/layer this chunk +/// might touch (built once, up front, for the whole run — see +/// `obikidxcache::IndexCache::new`), so every lookup here is a plain +/// in-memory operation, never a disk open. +#[allow(clippy::too_many_arguments)] +pub fn process_chunk( + cache: &IndexCache, + 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, + genomes: &[GenomeInfo], +) -> 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 — 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. 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. 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 = cache.query_partition_with(part_idx, kmers, n_genomes, |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)); + } + } + }); + 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). + 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. + 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, + genomes, + 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/obikquery/src/findere.rs b/src/obikquery/src/findere.rs new file mode 100644 index 00000000..a5042276 --- /dev/null +++ b/src/obikquery/src/findere.rs @@ -0,0 +1,75 @@ +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(crate) 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(crate) 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) +} + +#[cfg(test)] +#[path = "tests/findere.rs"] +mod tests; diff --git a/src/obikquery/src/lib.rs b/src/obikquery/src/lib.rs index 4201bd0a..e8d0fb4c 100644 --- a/src/obikquery/src/lib.rs +++ b/src/obikquery/src/lib.rs @@ -1,10 +1,16 @@ -//! Query-side operations on an `obikindex::KmerIndex`: staging ground for -//! code currently living in `obikindex::index`'s query path, to be migrated -//! here to lighten that crate — see `DevDocMD/` for the rationale. Kept as -//! a separate crate — not a module of `obikindex` — so the dependency runs -//! one way only (query code depends on the data model, never the reverse), -//! same pattern as `obikindexer` for the build side. +//! Query-side operations on an `obikidxcache::IndexCache`: matching query +//! sequences' k-mers against an already-built `obikindex::KmerIndex` and +//! formatting the result. Kept as a separate crate — not a module of +//! `obikindex` — so the dependency runs one way only (query code depends on +//! the data model, never the reverse), same pattern as `obikindexer` for the +//! build side. mod query_layer; +mod batch; +mod chunk; +mod findere; +mod smer_index; +mod output; -pub use query_layer::{KmerDesc, QueryHit, QueryStats}; +pub use query_layer::{KmerDesc, QueryHit, QueryPartition, QueryStats}; +pub use chunk::process_chunk; diff --git a/src/obikquery/src/output.rs b/src/obikquery/src/output.rs new file mode 100644 index 00000000..36babdf0 --- /dev/null +++ b/src/obikquery/src/output.rs @@ -0,0 +1,52 @@ +use std::io::Write; + +use obikindex::GenomeInfo; + +use crate::batch::QueryBatch; +use crate::chunk::SeqAcc; + +pub(crate) fn emit_batch( + batch: &QueryBatch, + accs: &[SeqAcc], + genomes: &[GenomeInfo], + 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 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 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/obikquery/src/query_layer.rs b/src/obikquery/src/query_layer.rs index 2684408e..b662e0e4 100644 --- a/src/obikquery/src/query_layer.rs +++ b/src/obikquery/src/query_layer.rs @@ -1,78 +1,7 @@ use std::collections::HashMap; -use std::path::Path; -use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix}; +use obikidxcache::index_cache::IndexCache; use obikseq::CanonicalKmer; -use obikindex::layer::MphfLayer; -use obikindex::{OKIError, OKIResult}; - -use obikindex::KmerIndex; - -// ── per-layer query handle ──────────────────────────────────────────────────── - -enum QueryLayer { - Presence(MphfLayer, PersistentBitMatrix), - Count(MphfLayer, PersistentCompactIntMatrix), -} - -impl QueryLayer { - fn open(layer_dir: &Path, with_counts: bool) -> OKIResult { - let mphf = MphfLayer::open(layer_dir)?; - let counts_dir = layer_dir.join("counts"); - let presence_dir = layer_dir.join("presence"); - - if with_counts && counts_dir.exists() { - let mat = PersistentCompactIntMatrix::open(layer_dir).map_err(OKIError::Io)?; - Ok(QueryLayer::Count(mphf, mat)) - } else if presence_dir.exists() || !counts_dir.exists() { - // presence mode, or no matrix at all → Implicit handled inside open() - let mat = PersistentBitMatrix::open(layer_dir).map_err(OKIError::Io)?; - Ok(QueryLayer::Presence(mphf, mat)) - } else { - // counts exist but not presence — count layer, no presence requested - let mat = PersistentCompactIntMatrix::open(layer_dir).map_err(OKIError::Io)?; - Ok(QueryLayer::Count(mphf, mat)) - } - } - - /// MPHF lookup only — no matrix access. `Some(slot)` on hit. - fn find_slot(&self, kmer: CanonicalKmer) -> Option { - match self { - QueryLayer::Presence(mphf, _) | QueryLayer::Count(mphf, _) => mphf.find(kmer), - } - } - - /// Number of genome columns this layer's matrix actually has. Bounds - /// column-major iteration — usually equal to the index's `n_genomes`, but - /// `PersistentBitMatrix::Implicit` (the documented mono-genome fast path) - /// always reports exactly `1`, regardless of the index's real genome - /// count, so callers must use this rather than assuming `n_genomes`. - fn n_cols(&self) -> usize { - match self { - QueryLayer::Presence(_, mat) => mat.n_cols(), - QueryLayer::Count(_, mat) => mat.n_cols(), - } - } - - /// Every nonzero `(idx into slots, col, value)` triple among `slots`. - /// Format-agnostic: each matrix picks its own natural traversal - /// (`PersistentBitMatrix::nonzero_iter` dispatches to a genuinely - /// row-major decode on `Sparse`, not a column-major point-probe loop — - /// see `DevDocMD/architecture/siblings.md`, "`query` never benefits - /// from sparse row-major access"). Replaces the old per-`(genome, - /// slot)` `col_value` point lookup, which this layer's `Sparse` - /// presence matrices paid for badly: each such lookup rebuilt the - /// entire row just to return one cell. - fn nonzero_iter<'a>( - &'a self, - slots: &'a [usize], - ) -> Box + 'a> { - match self { - QueryLayer::Presence(_, mat) => mat.nonzero_iter(slots), - QueryLayer::Count(_, mat) => Box::new(mat.nonzero_iter(slots)), - } - } -} // ── KmerDesc — one occurrence of a k-mer in the query batch ────────────────── @@ -85,7 +14,7 @@ pub struct KmerDesc { } /// Aggregate counters for one `query_partition_with` call — feeds the -/// dedup-ratio and column-scan logging in `obikmer::cmd::query` (occurrences +/// dedup-ratio and column-scan logging in `obikquery::chunk` (occurrences /// vs. unique k-mers is the whole justification for k-mer-level /// dereplication; columns scanned / `get()` calls quantify the column-major /// fetch's locality claim). @@ -93,7 +22,7 @@ pub struct KmerDesc { pub struct QueryStats { /// Distinct canonical k-mers queried in this partition. pub n_unique_kmers: usize, - /// Total `MphfLayer::find` calls issued (a k-mer tried against more than + /// Total MPHF-membership checks issued (a k-mer tried against more than /// one layer before a hit, or against all layers on a miss, counts once /// per layer attempted). pub n_mphf_calls: usize, @@ -102,7 +31,7 @@ pub struct QueryStats { /// Total genome columns scanned across all hit layers (sum of /// `layer.n_cols()` over layers with at least one hit). pub n_columns_scanned: usize, - /// Total `col_value` calls issued during the column-major fetch pass + /// Total nonzero-cell fetches issued during the column-major fetch pass /// (`n_columns_scanned` × hits-per-layer, summed over layers). pub n_col_get_calls: usize, } @@ -119,7 +48,7 @@ impl std::ops::AddAssign for QueryStats { // ── QueryHit — one event delivered to query_partition_with's callback ─────── -/// One event from [`KmerPartition::query_partition_with`]'s two-stage query: +/// One event from [`QueryPartition::query_partition_with`]'s two-stage query: /// a `Found` event once per hit k-mer (stage 1, MPHF-only — mark the k-mer as /// indexed regardless of any genome's value), then a `Value` event per /// `(hit k-mer, genome)` pair with a nonzero matrix value (stage 2, @@ -131,9 +60,21 @@ pub enum QueryHit<'a> { Value(&'a [KmerDesc], usize, u32), } -// ── KmerPartition::query_partition_with ────────────────────────────────────── +// ── QueryPartition — extension trait over obikidxcache::IndexCache ───────── -impl KmerIndex { +/// Query-side operations on an already-opened [`IndexCache`] — `IndexCache` +/// is a foreign type (`obikidxcache`), so this is an extension trait rather +/// than an inherent `impl` (`obikindex`'s original `impl KmerIndex { .. }` +/// shape doesn't compile outside `obikindex` itself — the orphan rule). +/// +/// Built on `IndexCache` rather than re-opening each layer's MPHF/matrix +/// files per call (the previous, per-call-`QueryLayer::open` design): every +/// layer touched by a query is already open, for the whole lifetime of the +/// cache, so a query pass costs zero disk I/O beyond the first chunk. This +/// also drops the `with_counts` parameter the old design needed: whether a +/// layer holds a count or a presence matrix is no longer guessed from an +/// index-wide flag, it's read straight off `KmerLayer`'s own variant. +pub trait QueryPartition { /// Query a single partition for a pre-deduplicated map of canonical /// k-mers → their occurrences (`seq_idx`, `pos`) in the query batch. /// @@ -151,43 +92,54 @@ impl KmerIndex { /// matrix formats are column-oriented on disk (one `mmap`'d region per /// genome), so scanning one column at a time touches far fewer /// distinct mmap regions than fetching one full row per hit. - pub fn query_partition_with( + /// + /// Infallible: every layer this can touch was already opened (and + /// validated) when the cache was built — see `IndexCache::new`'s own + /// panic-on-failure contract. A `part_idx` this cache doesn't hold, or + /// an empty `kmers` map, both just return default (all-zero) stats. + fn query_partition_with( + &self, + part_idx: usize, + kmers: &HashMap>, + n_genomes: usize, + on_event: F, + ) -> QueryStats + where + F: FnMut(QueryHit); +} + +impl QueryPartition for IndexCache { + fn query_partition_with( &self, part_idx: usize, kmers: &HashMap>, n_genomes: usize, - with_counts: bool, mut on_event: F, - ) -> OKIResult + ) -> QueryStats where F: FnMut(QueryHit), { let mut stats = QueryStats::default(); if kmers.is_empty() { - return Ok(stats); + return stats; } - let index_dir = self.index_dir(part_idx); - if !index_dir.exists() { - return Ok(stats); - } - - let meta = self.partition_meta(part_idx)?; - let layers: Vec = (0..meta.n_layers) - .map(|i| QueryLayer::open(&self.layer_dir(part_idx, i), with_counts)) - .collect::>()?; + let n_layer = match self.n_layer(part_idx) { + Some(n) if n > 0 => n, + _ => return stats, + }; // ── Stage 1: MPHF-only pass, bucket hits by (layer_idx, slot) ──────── let mut by_layer: Vec>> = - (0..layers.len()).map(|_| HashMap::new()).collect(); + (0..n_layer).map(|_| HashMap::new()).collect(); for (kmer, descs) in kmers { stats.n_unique_kmers += 1; - for (layer_idx, layer) in layers.iter().enumerate() { + for l in 0..n_layer { stats.n_mphf_calls += 1; - if let Some(slot) = layer.find_slot(*kmer) { - by_layer[layer_idx].insert(slot, descs); + if let Some(slot) = self.find_in_layer(part_idx, l, *kmer) { + by_layer[l].insert(slot, descs); on_event(QueryHit::Found(descs)); stats.n_hits += 1; break; @@ -196,16 +148,13 @@ impl KmerIndex { } // ── Stage 2: nonzero-cell fetch, per layer ──────────────────────────── - // Format-agnostic — see `QueryLayer::nonzero_iter`. `n_cols` still - // bounds accepted genome columns (Implicit reports fewer than - // `n_genomes`; see `n_cols`'s doc), cells beyond it are dropped - // rather than ever produced, since `nonzero_iter` only knows the - // matrix's own column count, not the caller's `n_genomes`. - for (layer_idx, slots) in by_layer.iter().enumerate() { + for (l, slots) in by_layer.iter().enumerate() { if slots.is_empty() { continue; } - let layer = &layers[layer_idx]; + let layer = self + .get_layer(part_idx, l) + .expect("layer within n_layer(part_idx)"); let n_cols = layer.n_cols().min(n_genomes); stats.n_columns_scanned += n_cols; @@ -221,7 +170,7 @@ impl KmerIndex { } } - Ok(stats) + stats } } diff --git a/src/obikquery/src/smer_index.rs b/src/obikquery/src/smer_index.rs new file mode 100644 index 00000000..178e5780 --- /dev/null +++ b/src/obikquery/src/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 `chunk::process_chunk`). +pub(crate) struct SmerIndex { + in_index: Vec, // total_smers + offsets: Vec, // offsets[i]..offsets[i+1] = s-mer range for sequence i +} + +impl SmerIndex { + pub(crate) 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(crate) fn mark_found(&mut self, seq: usize, kmer: usize) { + let abs = self.offsets[seq] + kmer; + self.in_index[abs] = true; + } + + #[inline] + pub(crate) fn is_in_index(&self, seq: usize, kmer: usize) -> bool { + self.in_index[self.offsets[seq] + kmer] + } +} diff --git a/src/obikquery/src/tests/batch.rs b/src/obikquery/src/tests/batch.rs new file mode 100644 index 00000000..6a289bd3 --- /dev/null +++ b/src/obikquery/src/tests/batch.rs @@ -0,0 +1,127 @@ +use super::*; +use obikrope::Rope; +use obikseq::CanonicalKmer; +use obiread::record::parse_chunk; + +const K: usize = 11; +const M: usize = 5; + +/// Build a `QueryBatch` from raw FASTA text, going through the same +/// `Rope` + `parse_chunk` path `process_chunk` uses — avoids hand-building a +/// `normalized` `Rope`, which is an implementation detail of `obiread`. +/// +/// `obikseq`'s global K/M params are thread-local under `test-utils` (see +/// `obikseq::params`), so setting them here is per-test-thread and does not +/// need coordination with other tests. +fn batch_from_fasta(fasta: &str, k: usize, n_partitions: usize) -> QueryBatch { + obikseq::set_k(k); + obikseq::set_m(M); + let mut rope = Rope::new(Some("text/fasta")); + rope.push(fasta.as_bytes().to_vec()); + let records = parse_chunk(&rope, k); + QueryBatch::from_records(records, k, 6, 0.7, n_partitions) +} + +fn total_occurrences(batch: &QueryBatch) -> u64 { + batch.n_kmers.iter().map(|&n| n as u64).sum() +} + +fn total_unique_kmers(batch: &QueryBatch) -> u64 { + batch.by_partition.iter().map(|m| m.len() as u64).sum() +} + +// A 60 bp sequence, arbitrary but fixed — no attempt is made to prove it is +// free of internal k=11 repeats; the tests below only rely on inequalities +// that hold regardless (see each test's comment). +const SEQ: &str = "CATTAGCGTACCTGATCAGGTTACAGCTTAGGCATCCAGTTGACCATGACTGGACTTAGC"; + +#[test] +fn single_sequence_yields_plausible_kmer_counts() { + // A single record can still contain internal repeats (SEQ isn't + // guaranteed repeat-free at k=11) — this only checks the batch is + // internally consistent, not a specific dedup ratio. The cross-record + // tests below make the actual, unconditional dedup claims. + let fasta = format!(">r1\n{SEQ}\n"); + let batch = batch_from_fasta(&fasta, K, 1); + + assert_eq!(batch.ids, vec!["r1".to_string()]); + let occurrences = total_occurrences(&batch); + let unique = total_unique_kmers(&batch); + assert!(occurrences > 0, "sequence should yield at least one k-mer"); + assert!(unique > 0 && unique <= occurrences); +} + +#[test] +fn duplicated_sequence_across_records_deduplicates() { + // Two records with byte-identical sequences: every k-mer in record 1 + // exactly duplicates one in record 0, so unique kmers <= n_kmers[0], + // strictly less than the summed occurrences (2 * n_kmers[0]) as long as + // the sequence yields at least one k-mer. This holds regardless of + // whether SEQ has internal repeats. + let fasta = format!(">r1\n{SEQ}\n>r2\n{SEQ}\n"); + let batch = batch_from_fasta(&fasta, K, 1); + + assert_eq!(batch.ids.len(), 2); + let occurrences = total_occurrences(&batch); + let unique = total_unique_kmers(&batch); + + assert!(batch.n_kmers[0] > 0); + assert_eq!(occurrences, batch.n_kmers[0] as u64 + batch.n_kmers[1] as u64); + assert!( + unique <= batch.n_kmers[0] as u64, + "identical sequences must not produce more unique k-mers than one copy has" + ); + assert!( + unique < occurrences, + "k-mer-level dedup must collapse at least the cross-record duplication" + ); +} + +#[test] +fn duplicated_sequence_broadcasts_to_both_seq_indices() { + // Stronger than the ratio check above: pick any k-mer that hit in both + // records and confirm its occurrence list actually references both + // seq_idx 0 and seq_idx 1 — this is the specific new capability (dedup + // reaching across records/superkmers), not just a smaller unique count. + let fasta = format!(">r1\n{SEQ}\n>r2\n{SEQ}\n"); + let batch = batch_from_fasta(&fasta, K, 1); + + let shared = batch.by_partition[0] + .values() + .find(|descs| descs.iter().any(|d| d.seq_idx == 0) && descs.iter().any(|d| d.seq_idx == 1)); + + assert!( + shared.is_some(), + "expected at least one k-mer shared between the two identical records" + ); +} + +#[test] +fn empty_records_yield_empty_batch() { + let batch = batch_from_fasta("", K, 1); + assert!(batch.ids.is_empty()); + assert_eq!(total_occurrences(&batch), 0); + assert_eq!(total_unique_kmers(&batch), 0); +} + +#[test] +fn partition_routing_is_a_pure_function_of_the_kmer() { + // With n_partitions=4, every occurrence of a given k-mer must land in + // the same partition bucket as every other occurrence of that k-mer + // (partition routing is derived from the minimizer, shared by + // definition among instances of the same k-mer's containing superkmer + // in this test's single-sequence-pair setup). + let fasta = format!(">r1\n{SEQ}\n>r2\n{SEQ}\n"); + let batch = batch_from_fasta(&fasta, K, 4); + + let total_unique: u64 = batch.by_partition.iter().map(|m| m.len() as u64).sum(); + assert!(total_unique > 0); + + // No k-mer key appears in more than one partition's map. + let mut seen: std::collections::HashSet = std::collections::HashSet::new(); + for map in &batch.by_partition { + for kmer in map.keys() { + assert!(seen.insert(*kmer), "k-mer routed to more than one partition"); + } + } +} diff --git a/src/obikquery/src/tests/findere.rs b/src/obikquery/src/tests/findere.rs new file mode 100644 index 00000000..76b2b36e --- /dev/null +++ b/src/obikquery/src/tests/findere.rs @@ -0,0 +1,105 @@ +use super::*; + +// ── sparse_findere_for_genome vs. a dense reference implementation ────────── +// +// No property-testing crate (proptest/quickcheck) is a workspace dependency +// (checked before writing this — not adding one for a single test module, +// per this project's dependency-approval rule). A tiny deterministic xorshift +// PRNG, std-only, stands in for one. + +/// Faithful reimplementation of the pre-phase-5 dense sliding-window scan — +/// the algorithm `sparse_findere_for_genome` replaced — used here only as a +/// correctness oracle, not in production code. Operates on one genome's +/// hits across possibly many sequences, exactly like the sparse version. +fn dense_reference_findere( + hits: &[(u32, u32, u32)], + seq_lens: &[usize], + z: usize, + presence: bool, + threshold: u32, +) -> Vec<(u32, u32, u32)> { + let mut by_seq: Vec> = seq_lens.iter().map(|&n| vec![0u32; n]).collect(); + for &(seq, pos, val) in hits { + by_seq[seq as usize][pos as usize] = val; + } + + let mut confirmed = Vec::new(); + for (seq_idx, values) in by_seq.iter().enumerate() { + let n = values.len(); + let mut dq: std::collections::VecDeque<(usize, u32)> = std::collections::VecDeque::new(); + for i in 0..n { + let v_i = values[i]; + while dq.front().map_or(false, |&(f, _)| f + z <= i) { + dq.pop_front(); + } + while dq.back().map_or(false, |&(_, v)| v >= v_i) { + dq.pop_back(); + } + dq.push_back((i, v_i)); + if i + 1 >= z { + let win_min = dq.front().unwrap().1; + if win_min > 0 { + let pos_out = (i + 1 - z) as u32; + let c = if presence { u32::from(win_min >= threshold) } else { win_min }; + confirmed.push((seq_idx as u32, pos_out, c)); + } + } + } + } + confirmed +} + +/// Minimal std-only xorshift64 PRNG — deterministic, seedable, no dependency. +struct Xorshift64(u64); +impl Xorshift64 { + fn next(&mut self) -> u64 { + self.0 ^= self.0 << 13; + self.0 ^= self.0 >> 7; + self.0 ^= self.0 << 17; + self.0 + } + fn range(&mut self, n: u32) -> u32 { + (self.next() % n as u64) as u32 + } +} + +#[test] +fn sparse_findere_matches_dense_reference_on_random_inputs() { + let mut rng = Xorshift64(0x5eed_5eed_5eed_5eedu64); + + for case in 0..200 { + let n_seqs = 1 + rng.range(4) as usize; + let seq_lens: Vec = (0..n_seqs).map(|_| 1 + rng.range(30) as usize).collect(); + let z = 1 + rng.range(4) as usize; + let presence = rng.range(2) == 0; + let threshold = 1 + rng.range(3); + + // Sparse density varies across cases, including edge cases (empty, + // fully dense) — deliberately not uniform, to stress both few-hits + // and many-overlapping-runs scenarios. + let density = rng.range(101); + let mut hits: Vec<(u32, u32, u32)> = Vec::new(); + for (seq_idx, &len) in seq_lens.iter().enumerate() { + for pos in 0..len { + if rng.range(100) < density { + let val = 1 + rng.range(5); // never 0 — matches QueryHit::Value's invariant + hits.push((seq_idx as u32, pos as u32, val)); + } + } + } + + let mut sparse_input = hits.clone(); + let (mut sparse_result, _, _) = + sparse_findere_for_genome(&mut sparse_input, z, presence, threshold); + let mut dense_result = dense_reference_findere(&hits, &seq_lens, z, presence, threshold); + + sparse_result.sort_unstable(); + dense_result.sort_unstable(); + + assert_eq!( + sparse_result, dense_result, + "case {case}: n_seqs={n_seqs} seq_lens={seq_lens:?} z={z} presence={presence} \ + threshold={threshold} density={density} hits={hits:?}" + ); + } +} diff --git a/src/obikquery/src/tests/query_layer.rs b/src/obikquery/src/tests/query_layer.rs index 5dca0266..55ba4517 100644 --- a/src/obikquery/src/tests/query_layer.rs +++ b/src/obikquery/src/tests/query_layer.rs @@ -1,5 +1,7 @@ use super::*; -use obikindex::IndexConfig; +use obikidxcache::index_cache::IndexCache; +use obikindex::{IndexConfig, KmerIndex}; +use std::sync::Arc; // ── QueryStats::AddAssign ─────────────────────────────────────────────────── @@ -39,33 +41,39 @@ fn query_stats_default_is_zero() { // ── query_partition_with on a not-yet-indexed partition ───────────────────── /// A `KmerPartition` created but never taken through `build_layers` has no -/// `index/` subdirectory under any partition — `query_partition_with` must -/// recognise this and return default (all-zero) stats rather than erroring, -/// exactly like an empty `kmers` map. +/// `index/` subdirectory under any partition — `IndexCache::new` still opens +/// fine on it (zero layers found), and `query_partition_with` must recognise +/// this and return default (all-zero) stats rather than panicking, exactly +/// like an empty `kmers` map. #[test] fn query_partition_with_missing_index_dir_returns_default_stats() { let tmp = tempfile::tempdir().expect("tempdir"); + // Same k/m pair as `batch`'s tests (11, 5) — `obikseq`'s global k/m + // params are process-wide atomics in test builds (see + // `obikseq::params`'s own doc comment), shared by every test in this + // crate's single test binary regardless of thread; a different pair + // here would race with `batch`'s tests running concurrently. let config = IndexConfig { - kmer_size: 21, - minimizer_size: 9, + kmer_size: 11, + minimizer_size: 5, n_bits: 2, with_counts: false, evidence: obikindex::layer::IndexMode::Exact, block_bits: 0, }; let index = KmerIndex::create(tmp.path().join("idx"), config, None).expect("create index"); + let cache = IndexCache::new(Arc::new(index), Some(vec![0])); let mut kmers: HashMap> = HashMap::new(); // Any well-formed canonical k-mer works here — the call must return - // before ever attempting an MPHF lookup, since `index/` doesn't exist. + // before ever attempting an MPHF lookup, since the partition has zero + // layers. let kmer = CanonicalKmer::from_raw_unchecked(0u64); kmers.insert(kmer, vec![KmerDesc { seq_idx: 0, pos: 0 }]); - let stats = index - .query_partition_with(0, &kmers, 1, false, |_event| { - panic!("on_event must not be called: no index was built"); - }) - .expect("query_partition_with should not error on a missing index dir"); + let stats = cache.query_partition_with(0, &kmers, 1, |_event| { + panic!("on_event must not be called: no index was built"); + }); assert_eq!(stats, QueryStats::default()); } @@ -73,22 +81,23 @@ fn query_partition_with_missing_index_dir_returns_default_stats() { #[test] fn query_partition_with_empty_kmers_is_a_noop() { let tmp = tempfile::tempdir().expect("tempdir"); + // Same k/m pair as `batch`'s tests — see the comment in the previous + // test for why this must match across every test in this crate. let config = IndexConfig { - kmer_size: 21, - minimizer_size: 9, + kmer_size: 11, + minimizer_size: 5, n_bits: 2, with_counts: false, evidence: obikindex::layer::IndexMode::Exact, block_bits: 0, }; let index = KmerIndex::create(tmp.path().join("idx"), config, None).expect("create index"); + let cache = IndexCache::new(Arc::new(index), Some(vec![0])); let kmers: HashMap> = HashMap::new(); - let stats = index - .query_partition_with(0, &kmers, 1, false, |_event| { - panic!("on_event must not be called on an empty kmer map"); - }) - .expect("query_partition_with on an empty map should not error"); + let stats = cache.query_partition_with(0, &kmers, 1, |_event| { + panic!("on_event must not be called on an empty kmer map"); + }); assert_eq!(stats, QueryStats::default()); } diff --git a/src/obikrebuild/src/compact_layer.rs b/src/obikrebuild/src/compact_layer.rs index 5d603dba..47453250 100644 --- a/src/obikrebuild/src/compact_layer.rs +++ b/src/obikrebuild/src/compact_layer.rs @@ -8,6 +8,7 @@ use std::collections::HashMap; use std::fs; +use std::sync::Arc; use obidebruinj::GraphDeBruijn; use obikfilter::FilteredPartitionIter; @@ -24,7 +25,7 @@ pub trait IndexCompact { fn compact_layers(&self, on_partition: impl Fn() + Send + Sync) -> OKIResult<()>; } -impl IndexCompact for KmerIndex { +impl IndexCompact for Arc { fn compact_layers(&self, on_partition: impl Fn() + Send + Sync) -> OKIResult<()> { let n_genomes = self.meta().genomes().map_err(OKIError::Io)?.len(); let use_counts = self.meta().config.with_counts; @@ -47,7 +48,7 @@ impl IndexCompact for KmerIndex { /// existing layers in place. No-op if the partition already has one layer /// (or none). fn compact_partition( - idx: &KmerIndex, + idx: &Arc, i: usize, use_counts: bool, n_genomes: usize, @@ -63,7 +64,7 @@ fn compact_partition( // ── Read: every layer's kmers, combined — no filter, nothing removed ── let mut survivors: HashMap> = HashMap::new(); { - let cache = IndexCache::new(idx, Some(vec![i])); + let cache = IndexCache::new(Arc::clone(idx), Some(vec![i])); cache.iter_partition_kmers(i, use_counts, n_genomes, &[], |kmer, row| { survivors.insert(kmer, row); true diff --git a/src/obikstats/src/stats.rs b/src/obikstats/src/stats.rs index a89b8878..7324c166 100644 --- a/src/obikstats/src/stats.rs +++ b/src/obikstats/src/stats.rs @@ -1,5 +1,6 @@ use std::fs; use std::path::Path; +use std::sync::Arc; use rayon::prelude::*; @@ -88,22 +89,22 @@ fn layer_bytes(layer: &KmerLayer) -> LayerBytes { /// once up front — this is diagnostic tooling, not a hot path, so trading /// that eager open for not having to hand-reconstruct every layer's file /// paths is the right tradeoff here. -pub struct BitsPerKmer<'a> { - index: &'a KmerIndex, +pub struct BitsPerKmer { + index: Arc, } -impl<'a> BitsPerKmer<'a> { - pub fn new(index: &'a KmerIndex) -> Self { +impl BitsPerKmer { + pub fn new(index: Arc) -> Self { Self { index } } } -impl Algorithm for BitsPerKmer<'_> { +impl Algorithm for BitsPerKmer { type Output = IndexBitsPerKmer; fn run(&mut self) -> obikalgorithm::Result { let n_genomes = self.index.meta().genomes()?.len().max(1); - let cache = IndexCache::new(self.index, None); + let cache = IndexCache::new(Arc::clone(&self.index), None); let layers: Vec<&KmerLayer> = cache.iter().collect(); let (n_kmers, mphf_b, evidence_b, matrix_b) = layers @@ -148,23 +149,23 @@ impl Algorithm for BitsPerKmer<'_> { /// genome has a non-zero value (presence = 1, count > 0) — via /// `KmerLayer::col_weights` (`obicompactvec::ColumnWeights`), one call per /// layer instead of hand-opening each layer's matrix and summing rows. -pub struct GenomeKmerCounts<'a> { - index: &'a KmerIndex, +pub struct GenomeKmerCounts { + index: Arc, } -impl<'a> GenomeKmerCounts<'a> { - pub fn new(index: &'a KmerIndex) -> Self { +impl GenomeKmerCounts { + pub fn new(index: Arc) -> Self { Self { index } } } -impl Algorithm for GenomeKmerCounts<'_> { +impl Algorithm for GenomeKmerCounts { /// `(total_distinct_kmers, per_genome_kmer_counts)`. type Output = (usize, Vec); fn run(&mut self) -> obikalgorithm::Result<(usize, Vec)> { let n_genomes = self.index.meta().genomes()?.len(); - let cache = IndexCache::new(self.index, None); + let cache = IndexCache::new(Arc::clone(&self.index), None); let layers: Vec<&KmerLayer> = cache.iter().collect(); let total_kmers: usize = layers.iter().map(|l| l.n()).sum();