From a348637f3b6c1b34248e08cf47c375bb9e4a3981 Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Tue, 7 Jul 2026 16:14:10 +0200 Subject: [PATCH] refactor(query): deduplicate k-mers upfront and split MPHF lookup Refactor `QueryBatch` construction to perform canonical k-mer deduplication and partition routing during initialization, eliminating post-batch splitting. Split the `QueryLayer` MPHF lookup into separate `find_slot` and `fill_row` methods, updating callbacks to pass occurrence descriptors instead of indices. Introduce `QueryStats` for tracking MPHF calls and dereplication ratios, and add comprehensive unit tests for batch construction, stats arithmetic, and safe partition handling. Expose new query-layer types in the public API. --- src/obikmer/src/cmd/query.rs | 109 +++++++----- src/obikmer/src/cmd/tests/query.rs | 124 +++++++++++++ src/obikpartitionner/src/lib.rs | 1 + src/obikpartitionner/src/query_layer.rs | 165 +++++++++--------- src/obikpartitionner/src/tests/query_layer.rs | 65 +++++++ 5 files changed, 330 insertions(+), 134 deletions(-) create mode 100644 src/obikmer/src/cmd/tests/query.rs create mode 100644 src/obikpartitionner/src/tests/query_layer.rs diff --git a/src/obikmer/src/cmd/query.rs b/src/obikmer/src/cmd/query.rs index ab03b87..4b2d128 100644 --- a/src/obikmer/src/cmd/query.rs +++ b/src/obikmer/src/cmd/query.rs @@ -7,8 +7,9 @@ use std::time::Instant; use clap::Args; use obikindex::KmerIndex; +use obikpartitionner::{KmerDesc, QueryStats}; use obikrope::Rope; -use obikseq::RoutableSuperKmer; +use obikseq::CanonicalKmer; use obilayeredmap::IndexMode; use obipipeline::{Throttled, ThrottleGuard, throttle}; use obiread::chunk::read_sequence_chunks_sized; @@ -94,19 +95,18 @@ impl QueryArgs { } } -// ── SKDesc — one occurrence of a superkmer in the batch ─────────────────────── - -/// Describes one occurrence of a superkmer in the query batch. -pub struct SKDesc { - /// Index of the source sequence within the batch. - pub seq_idx: u32, - /// Kmer offset of the first kmer of this superkmer within its sequence. - pub kmer_offset: u32, -} - // ── QueryBatch ──────────────────────────────────────────────────────────────── -/// A batch of query sequences with their superkmers deduplicated. +/// 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, @@ -114,30 +114,40 @@ pub struct QueryBatch { pub seqs: Vec>, /// Total kmer count per sequence (used for `--detail` coverage allocation). pub n_kmers: Vec, - /// Deduplicated superkmer map. - pub map: HashMap>, + /// Deduplicated k-mer occurrences, one map per partition. + pub by_partition: Vec>>, } impl QueryBatch { - /// Build a batch from a vec of parsed sequence records. - pub fn from_records(records: Vec, k: usize, level_max: usize, theta: f64) -> Self { + /// 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()); - // Upper-bound estimate: at most one superkmer per k bases. - // Avoids repeated rehash on large chunks. - let cap = records.iter().map(|r| r.normalized.len()).sum::() / k.max(1); - let mut map: HashMap> = HashMap::with_capacity(cap); + 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; - map.entry(rsk).or_default().push(SKDesc { - seq_idx: seq_idx as u32, - kmer_offset, - }); kmer_offset += n; } @@ -150,20 +160,9 @@ impl QueryBatch { ids, seqs, n_kmers, - map, + by_partition, } } - - /// Split the superkmer map by partition index. - pub fn split_by_partition(&self, n_partitions: usize) -> Vec> { - let mask = (n_partitions as u64) - 1; - let mut by_part: Vec> = vec![Vec::new(); n_partitions]; - for rsk in self.map.keys() { - let part = (rsk.minimizer().seq_hash() & mask) as usize; - by_part[part].push(rsk); - } - by_part - } } // ── KmerResults — allocation-free ragged result matrix ─────────────────────── @@ -260,33 +259,36 @@ fn process_chunk( return Vec::new(); } - let batch = QueryBatch::from_records(records, k, 6, 0.7); + let batch = QueryBatch::from_records(records, k, 6, 0.7, n_partitions); let n_seqs = batch.ids.len(); // Flat result matrix — one allocation for the whole chunk. let mut results = KmerResults::new(&batch.n_kmers, n_genomes); - let by_part = batch.split_by_partition(n_partitions); + // 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, part_sks) in by_part.iter().enumerate() { - if part_sks.is_empty() { + for (part_idx, kmers) in batch.by_partition.iter().enumerate() { + if kmers.is_empty() { continue; } - idx.partition() + let stats = idx.partition() .query_partition_with( part_idx, - part_sks, - k, + kmers, n_genomes, with_counts, - |sk_idx, kmer_idx, row| { - let rsk = part_sks[sk_idx]; - let descs = batch.map.get(rsk).expect("rsk must be in map"); + |descs, row| { for desc in descs { results.set( desc.seq_idx as usize, - desc.kmer_offset as usize + kmer_idx, + desc.pos as usize, row, ); } @@ -296,8 +298,17 @@ fn process_chunk( 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, + "k-mer dedup" + ); + // Sliding window minimum — one reusable buffer and one deque per batch. // // win_min[pos * n_genomes + g] = min count across the z-window [pos, pos+z) @@ -686,3 +697,7 @@ fn emit_batch( let _ = out.write_all(b"\n"); } } + +#[cfg(test)] +#[path = "tests/query.rs"] +mod tests; diff --git a/src/obikmer/src/cmd/tests/query.rs b/src/obikmer/src/cmd/tests/query.rs new file mode 100644 index 0000000..37a2083 --- /dev/null +++ b/src/obikmer/src/cmd/tests/query.rs @@ -0,0 +1,124 @@ +use super::*; + +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/obikpartitionner/src/lib.rs b/src/obikpartitionner/src/lib.rs index a63e2c9..44189e4 100644 --- a/src/obikpartitionner/src/lib.rs +++ b/src/obikpartitionner/src/lib.rs @@ -14,4 +14,5 @@ mod select_layer; pub use filter::{GroupQuorumFilter, KmerFilter, passes_all}; pub use merge_layer::MergeMode; pub use partition::{KmerPartition, KmerSpectrum, PARTITIONS_SUBDIR}; +pub use query_layer::{KmerDesc, QueryStats}; pub use select_layer::{AggOp, OutputCol}; diff --git a/src/obikpartitionner/src/query_layer.rs b/src/obikpartitionner/src/query_layer.rs index e918ec6..8219bda 100644 --- a/src/obikpartitionner/src/query_layer.rs +++ b/src/obikpartitionner/src/query_layer.rs @@ -1,7 +1,8 @@ +use std::collections::HashMap; use std::path::Path; use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix}; -use obikseq::{CanonicalKmer, RoutableSuperKmer}; +use obikseq::CanonicalKmer; use obiskio::{SKError, SKResult}; use obilayeredmap::{IndexMode, MphfLayer, OLMError}; use obilayeredmap::meta::PartitionMeta; @@ -44,53 +45,86 @@ impl QueryLayer { } } - /// Write per-genome values into `buf` if `kmer` is indexed; returns true on hit. - fn find_into(&self, kmer: CanonicalKmer, n_genomes: usize, buf: &mut [u32]) -> bool { + /// MPHF lookup only — no matrix access. `Some(slot)` on hit. + fn find_slot(&self, kmer: CanonicalKmer) -> Option { match self { - QueryLayer::Presence(mphf, mat) => { - if let Some(slot) = mphf.find(kmer) { - mat.fill_row(slot, &mut buf[..n_genomes]); - true - } else { - false - } - } - QueryLayer::Count(mphf, mat) => { - if let Some(slot) = mphf.find(kmer) { - mat.fill_row(slot, &mut buf[..n_genomes]); - true - } else { - false - } - } + QueryLayer::Presence(mphf, _) | QueryLayer::Count(mphf, _) => mphf.find(kmer), + } + } + + /// Write per-genome values for `slot` into `buf`. `slot` must come from + /// [`find_slot`] on this same layer. + fn fill_row(&self, slot: usize, n_genomes: usize, buf: &mut [u32]) { + match self { + QueryLayer::Presence(_, mat) => mat.fill_row(slot, &mut buf[..n_genomes]), + QueryLayer::Count(_, mat) => mat.fill_row(slot, &mut buf[..n_genomes]), } } } -// ── KmerPartition::query_partition* ────────────────────────────────────────── +// ── KmerDesc — one occurrence of a k-mer in the query batch ────────────────── + +/// Describes one occurrence of a (deduplicated) k-mer in the query batch: +/// which sequence it came from, and its absolute s-mer position within it. +#[derive(Debug, Clone, Copy)] +pub struct KmerDesc { + pub seq_idx: u32, + pub pos: u32, +} + +/// Aggregate counters for one `query_partition_with` call — feeds the +/// dedup-ratio logging in `obikmer::cmd::query` (occurrences vs. unique +/// k-mers is the whole justification for k-mer-level dereplication). +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] +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 + /// one layer before a hit, or against all layers on a miss, counts once + /// per layer attempted). + pub n_mphf_calls: usize, + /// Distinct canonical k-mers that matched some layer. + pub n_hits: usize, +} + +impl std::ops::AddAssign for QueryStats { + fn add_assign(&mut self, other: Self) { + self.n_unique_kmers += other.n_unique_kmers; + self.n_mphf_calls += other.n_mphf_calls; + self.n_hits += other.n_hits; + } +} + +// ── KmerPartition::query_partition_with ────────────────────────────────────── impl KmerPartition { - /// Query a single partition, calling `on_hit(sk_idx, kmer_idx, row)` for - /// every found k-mer without allocating intermediate result vectors. + /// Query a single partition for a pre-deduplicated map of canonical + /// k-mers → their occurrences (`seq_idx`, `pos`) in the query batch. + /// + /// Each unique k-mer triggers at most one MPHF lookup per layer (stopping + /// at the first hit) and, on hit, one matrix row fetch — regardless of + /// how many times that k-mer occurs across the batch. `on_hit(descs, row)` + /// is called once per hit, with every occurrence to broadcast the row to. pub fn query_partition_with( &self, part_idx: usize, - superkmers: &[&RoutableSuperKmer], - _k: usize, + kmers: &HashMap>, n_genomes: usize, with_counts: bool, mut on_hit: F, - ) -> SKResult<()> + ) -> SKResult where - F: FnMut(usize, usize, &[u32]), + F: FnMut(&[KmerDesc], &[u32]), { - if superkmers.is_empty() { - return Ok(()); + let mut stats = QueryStats::default(); + + if kmers.is_empty() { + return Ok(stats); } let index_dir = self.part_dir(part_idx).join(INDEX_SUBDIR); if !index_dir.exists() { - return Ok(()); + return Ok(stats); } let meta = PartitionMeta::load(&index_dir).map_err(olm_to_sk)?; @@ -100,67 +134,24 @@ impl KmerPartition { let mut buf = vec![0u32; n_genomes]; - for (sk_idx, rsk) in superkmers.iter().enumerate() { - for (kmer_idx, kmer) in rsk.superkmer().iter_canonical_kmers().enumerate() { - for layer in &layers { - if layer.find_into(kmer, n_genomes, &mut buf) { - on_hit(sk_idx, kmer_idx, &buf); - buf.fill(0); - break; - } + for (kmer, descs) in kmers { + stats.n_unique_kmers += 1; + for layer in &layers { + stats.n_mphf_calls += 1; + if let Some(slot) = layer.find_slot(*kmer) { + layer.fill_row(slot, n_genomes, &mut buf); + on_hit(descs, &buf); + buf.fill(0); + stats.n_hits += 1; + break; } } } - Ok(()) - } - - /// Query a single partition for a slice of super-kmers, returning per-kmer rows. - /// Prefer [`query_partition_with`] to avoid per-kmer heap allocations. - pub fn query_partition( - &self, - part_idx: usize, - superkmers: &[&RoutableSuperKmer], - _k: usize, - n_genomes: usize, - with_counts: bool, - ) -> SKResult>>>> { - if superkmers.is_empty() { - return Ok(Vec::new()); - } - - let index_dir = self.part_dir(part_idx).join(INDEX_SUBDIR); - - if !index_dir.exists() { - return Ok(superkmers - .iter() - .map(|rsk| vec![None; rsk.seql()]) - .collect()); - } - - let meta = PartitionMeta::load(&index_dir).map_err(olm_to_sk)?; - let layers: Vec = (0..meta.n_layers) - .map(|i| QueryLayer::open(&index_dir.join(format!("layer_{i}")), with_counts, &meta.mode)) - .collect::>()?; - - let mut buf = vec![0u32; n_genomes]; - Ok(superkmers - .iter() - .map(|rsk| { - rsk.superkmer() - .iter_canonical_kmers() - .map(|kmer| { - for layer in &layers { - if layer.find_into(kmer, n_genomes, &mut buf) { - let row: Box<[u32]> = buf[..n_genomes].into(); - buf.fill(0); - return Some(row); - } - } - None - }) - .collect() - }) - .collect()) + Ok(stats) } } + +#[cfg(test)] +#[path = "tests/query_layer.rs"] +mod tests; diff --git a/src/obikpartitionner/src/tests/query_layer.rs b/src/obikpartitionner/src/tests/query_layer.rs new file mode 100644 index 0000000..06e6824 --- /dev/null +++ b/src/obikpartitionner/src/tests/query_layer.rs @@ -0,0 +1,65 @@ +use super::*; + +// ── QueryStats::AddAssign ─────────────────────────────────────────────────── + +#[test] +fn query_stats_add_assign_sums_fields() { + let mut total = QueryStats { n_unique_kmers: 3, n_mphf_calls: 5, n_hits: 2 }; + total += QueryStats { n_unique_kmers: 1, n_mphf_calls: 4, n_hits: 1 }; + assert_eq!(total.n_unique_kmers, 4); + assert_eq!(total.n_mphf_calls, 9); + assert_eq!(total.n_hits, 3); +} + +#[test] +fn query_stats_default_is_zero() { + let s = QueryStats::default(); + assert_eq!(s.n_unique_kmers, 0); + assert_eq!(s.n_mphf_calls, 0); + assert_eq!(s.n_hits, 0); +} + +// ── 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. +#[test] +fn query_partition_with_missing_index_dir_returns_default_stats() { + let tmp = tempfile::tempdir().expect("tempdir"); + let partition = KmerPartition::create(tmp.path().join("idx"), 2, 21, 9, false) + .expect("create partition"); + + 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. + let kmer = CanonicalKmer::from_raw_unchecked(0u64); + kmers.insert(kmer, vec![KmerDesc { seq_idx: 0, pos: 0 }]); + + let stats = partition + .query_partition_with(0, &kmers, 1, false, |_descs, _row| { + panic!("on_hit must not be called: no index was built"); + }) + .expect("query_partition_with should not error on a missing index dir"); + + assert_eq!(stats.n_unique_kmers, 0); + assert_eq!(stats.n_mphf_calls, 0); + assert_eq!(stats.n_hits, 0); +} + +#[test] +fn query_partition_with_empty_kmers_is_a_noop() { + let tmp = tempfile::tempdir().expect("tempdir"); + let partition = KmerPartition::create(tmp.path().join("idx"), 2, 21, 9, false) + .expect("create partition"); + + let kmers: HashMap> = HashMap::new(); + let stats = partition + .query_partition_with(0, &kmers, 1, false, |_, _| { + panic!("on_hit must not be called on an empty kmer map"); + }) + .expect("query_partition_with on an empty map should not error"); + + assert_eq!(stats, QueryStats::default()); +}