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<String>,
@@ -114,30 +114,40 @@ pub struct QueryBatch {
     pub seqs: Vec<Vec<u8>>,
     /// Total kmer count per sequence (used for `--detail` coverage allocation).
     pub n_kmers: Vec<u32>,
-    /// Deduplicated superkmer map.
-    pub map: HashMap<RoutableSuperKmer, Vec<SKDesc>>,
+    /// Deduplicated k-mer occurrences, one map per partition.
+    pub by_partition: Vec<HashMap<CanonicalKmer, Vec<KmerDesc>>>,
 }
 
 impl QueryBatch {
-    /// Build a batch from a vec of parsed sequence records.
-    pub fn from_records(records: Vec<SeqRecord>, 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<SeqRecord>,
+        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::<usize>() / k.max(1);
-        let mut map: HashMap<RoutableSuperKmer, Vec<SKDesc>> = HashMap::with_capacity(cap);
+        let mask = (n_partitions as u64) - 1;
+        let mut by_partition: Vec<HashMap<CanonicalKmer, Vec<KmerDesc>>> =
+            (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 n = (rsk.seql() - k + 1) as u32;
-                map.entry(rsk).or_default().push(SKDesc {
+                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,
-                    kmer_offset,
+                        pos: kmer_offset + j as u32,
                     });
+                }
+                let n = (rsk.seql() - k + 1) as u32;
                 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<Vec<&RoutableSuperKmer>> {
-        let mask = (n_partitions as u64) - 1;
-        let mut by_part: Vec<Vec<&RoutableSuperKmer>> = 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;

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<CanonicalKmer> = 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");
+        }
+    }
+}

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};

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<usize> {
         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<F>(
         &self,
         part_idx: usize,
-        superkmers: &[&RoutableSuperKmer],
-        _k: usize,
+        kmers: &HashMap<CanonicalKmer, Vec<KmerDesc>>,
         n_genomes: usize,
         with_counts: bool,
         mut on_hit: F,
-    ) -> SKResult<()>
+    ) -> SKResult<QueryStats>
     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 (kmer, descs) in kmers {
+            stats.n_unique_kmers += 1;
             for layer in &layers {
-                    if layer.find_into(kmer, n_genomes, &mut buf) {
-                        on_hit(sk_idx, kmer_idx, &buf);
+                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(stats)
+    }
 }
 
-        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<Vec<Vec<Option<Box<[u32]>>>>> {
-        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<QueryLayer> = (0..meta.n_layers)
-            .map(|i| QueryLayer::open(&index_dir.join(format!("layer_{i}")), with_counts, &meta.mode))
-            .collect::<SKResult<_>>()?;
-
-        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())
-    }
-}
+#[cfg(test)]
+#[path = "tests/query_layer.rs"]
+mod tests;

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<CanonicalKmer, Vec<KmerDesc>> = 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<CanonicalKmer, Vec<KmerDesc>> = 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());
+}