refactor(query): optimize mmap locality with column-major matrix fetch

Refactor the query pipeline into a two-stage MPHF hit-detection pass followed by a column-major matrix fetch to improve cache efficiency. Introduce a QueryHit enum for event-driven callbacks, decoupling hit detection from data population. Add scan/fetch metrics to QueryStats, update Phase 4 architecture docs, and align tests with the new callback signature.

src/obicompactvec/src/bitmatrix.rs

@@ -318,6 +318,19 @@ impl PersistentBitMatrix {
         }
     }
 
+    /// Column-major point lookup: value at column `c`, slot `slot`, as 0/1.
+    ///
+    /// Unlike [`col_view`](Self::col_view), this never panics on `Implicit`
+    /// (every column reads as present, per the mono-genome fast path) — safe
+    /// to call for any `c < self.n_cols()`.
+    pub fn get(&self, c: usize, slot: usize) -> u32 {
+        match self {
+            Self::Columnar(m) => m.col(c).get(slot) as u32,
+            Self::Packed(m)   => m.col_slice(c).get(slot) as u32,
+            Self::Implicit { .. } => 1,
+        }
+    }
+
     pub fn col_persist(&self, c: usize, path: &Path) -> io::Result<PersistentBitVecBuilder> {
         match self {
             Self::Columnar(m) => PersistentBitVecBuilder::build_from(m.col(c), path),

src/obikmer/src/cmd/query.rs

@@ -7,7 +7,7 @@ use std::time::Instant;
 
 use clap::Args;
 use obikindex::KmerIndex;
-use obikpartitionner::{KmerDesc, QueryStats};
+use obikpartitionner::{KmerDesc, QueryHit, QueryStats};
 use obikrope::Rope;
 use obikseq::CanonicalKmer;
 use obilayeredmap::IndexMode;
@@ -199,11 +199,20 @@ impl KmerResults {
         self.offsets[seq + 1] - self.offsets[seq]
     }
 
-    fn set(&mut self, seq: usize, kmer: usize, row: &[u32]) {
+    /// 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 fills in per-genome values.
+    fn mark_found(&mut self, seq: usize, kmer: usize) {
         let abs = self.offsets[seq] + kmer;
         self.in_index[abs] = true;
-        let base = abs * self.n_genomes;
-        self.data[base..base + self.n_genomes].copy_from_slice(row);
+    }
+
+    /// Set the value for one genome at (seq, kmer). Called once per nonzero
+    /// `(k-mer, genome)` pair delivered by the column-major fetch (stage 2).
+    fn set_one(&mut self, seq: usize, kmer: usize, g: usize, value: u32) {
+        let abs = self.offsets[seq] + kmer;
+        self.data[abs * self.n_genomes + g] = value;
     }
 
     #[inline]
@@ -284,13 +293,16 @@ fn process_chunk(
                 kmers,
                 n_genomes,
                 with_counts,
-                |descs, row| {
-                    for desc in descs {
-                        results.set(
-                            desc.seq_idx as usize,
-                            desc.pos as usize,
-                            row,
-                        );
+                |event| match event {
+                    QueryHit::Found(descs) => {
+                        for desc in descs {
+                            results.mark_found(desc.seq_idx as usize, desc.pos as usize);
+                        }
+                    }
+                    QueryHit::Value(descs, g, v) => {
+                        for desc in descs {
+                            results.set_one(desc.seq_idx as usize, desc.pos as usize, g, v);
+                        }
                     }
                 },
             )
@@ -306,7 +318,9 @@ fn process_chunk(
         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"
+        n_columns_scanned = query_stats.n_columns_scanned,
+        n_col_get_calls = query_stats.n_col_get_calls,
+        "k-mer dedup + column-major fetch"
     );
 
     // Sliding window minimum — one reusable buffer and one deque per batch.

src/obikpartitionner/src/lib.rs

@@ -14,5 +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 query_layer::{KmerDesc, QueryHit, QueryStats};
 pub use select_layer::{AggOp, OutputCol};

src/obikpartitionner/src/query_layer.rs

@@ -52,12 +52,25 @@ impl QueryLayer {
         }
     }
 
-    /// 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]) {
+    /// 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.fill_row(slot, &mut buf[..n_genomes]),
-            QueryLayer::Count(_, mat) => mat.fill_row(slot, &mut buf[..n_genomes]),
+            QueryLayer::Presence(_, mat) => mat.n_cols(),
+            QueryLayer::Count(_, mat) => mat.n_cols(),
+        }
+    }
+
+    /// Column-major point lookup: value for genome column `g` at `slot`.
+    /// `g` must be `< self.n_cols()`; `slot` must come from [`find_slot`] on
+    /// this same layer.
+    fn col_value(&self, g: usize, slot: usize) -> u32 {
+        match self {
+            QueryLayer::Presence(_, mat) => mat.get(g, slot),
+            QueryLayer::Count(_, mat) => mat.col_view(g).get(slot),
         }
     }
 }
@@ -73,8 +86,10 @@ pub struct KmerDesc {
 }
 
 /// 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).
+/// dedup-ratio and column-scan logging in `obikmer::cmd::query` (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).
 #[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
 pub struct QueryStats {
     /// Distinct canonical k-mers queried in this partition.
@@ -85,6 +100,12 @@ pub struct QueryStats {
     pub n_mphf_calls: usize,
     /// Distinct canonical k-mers that matched some layer.
     pub n_hits: usize,
+    /// 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
+    /// (`n_columns_scanned` × hits-per-layer, summed over layers).
+    pub n_col_get_calls: usize,
 }
 
 impl std::ops::AddAssign for QueryStats {
@@ -92,29 +113,55 @@ impl std::ops::AddAssign for QueryStats {
         self.n_unique_kmers += other.n_unique_kmers;
         self.n_mphf_calls += other.n_mphf_calls;
         self.n_hits += other.n_hits;
+        self.n_columns_scanned += other.n_columns_scanned;
+        self.n_col_get_calls += other.n_col_get_calls;
     }
 }
 
+// ── QueryHit — one event delivered to query_partition_with's callback ───────
+
+/// One event from [`KmerPartition::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,
+/// column-major fetch). Carried as one enum, not two separate callbacks, so
+/// the caller only needs one `FnMut` closure — passing two closures that each
+/// need to mutably borrow the same accumulator does not borrow-check.
+pub enum QueryHit<'a> {
+    Found(&'a [KmerDesc]),
+    Value(&'a [KmerDesc], usize, u32),
+}
+
 // ── KmerPartition::query_partition_with ──────────────────────────────────────
 
 impl KmerPartition {
     /// 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.
+    /// Two stages:
+    /// 1. **MPHF-only pass**: for each unique k-mer, try each layer's MPHF in
+    ///    turn (stopping at the first hit) and bucket confirmed hits by
+    ///    `(layer, slot)`. Emits one `QueryHit::Found` per hit k-mer. This
+    ///    stage's cost is independent of the index's genome count.
+    /// 2. **Column-major fetch**: for each layer with at least one hit, walk
+    ///    its matrix **column by column** (genome by genome) — for each
+    ///    genome, scan the slots bucketed in stage 1 and look up their value.
+    ///    Emits one `QueryHit::Value` per nonzero `(k-mer, genome)` pair.
+    ///    Total lookups are the same as a row-major pass (`n_hits × n_cols`
+    ///    in the worst case); the win is memory locality — both persistent
+    ///    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<F>(
         &self,
         part_idx: usize,
         kmers: &HashMap<CanonicalKmer, Vec<KmerDesc>>,
         n_genomes: usize,
         with_counts: bool,
-        mut on_hit: F,
+        mut on_event: F,
     ) -> SKResult<QueryStats>
     where
-        F: FnMut(&[KmerDesc], &[u32]),
+        F: FnMut(QueryHit),
     {
         let mut stats = QueryStats::default();
 
@@ -132,22 +179,43 @@ impl KmerPartition {
             .map(|i| QueryLayer::open(&index_dir.join(format!("layer_{i}")), with_counts, &meta.mode))
             .collect::<SKResult<_>>()?;
 
-        let mut buf = vec![0u32; n_genomes];
+        // ── Stage 1: MPHF-only pass, bucket hits by (layer_idx, slot) ────────
+        let mut by_layer: Vec<HashMap<usize, &Vec<KmerDesc>>> =
+            (0..layers.len()).map(|_| HashMap::new()).collect();
 
         for (kmer, descs) in kmers {
             stats.n_unique_kmers += 1;
-            for layer in &layers {
+            for (layer_idx, layer) in layers.iter().enumerate() {
                 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);
+                    by_layer[layer_idx].insert(slot, descs);
+                    on_event(QueryHit::Found(descs));
                     stats.n_hits += 1;
                     break;
                 }
             }
         }
 
+        // ── Stage 2: column-major fetch, per layer ───────────────────────────
+        for (layer_idx, slots) in by_layer.iter().enumerate() {
+            if slots.is_empty() {
+                continue;
+            }
+            let layer = &layers[layer_idx];
+            let n_cols = layer.n_cols().min(n_genomes);
+            stats.n_columns_scanned += n_cols;
+
+            for g in 0..n_cols {
+                for (&slot, descs) in slots {
+                    stats.n_col_get_calls += 1;
+                    let v = layer.col_value(g, slot);
+                    if v != 0 {
+                        on_event(QueryHit::Value(descs, g, v));
+                    }
+                }
+            }
+        }
+
         Ok(stats)
     }
 }

src/obikpartitionner/src/tests/query_layer.rs

@@ -4,11 +4,25 @@ use super::*;
 
 #[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 };
+    let mut total = QueryStats {
+        n_unique_kmers: 3,
+        n_mphf_calls: 5,
+        n_hits: 2,
+        n_columns_scanned: 1,
+        n_col_get_calls: 7,
+    };
+    total += QueryStats {
+        n_unique_kmers: 1,
+        n_mphf_calls: 4,
+        n_hits: 1,
+        n_columns_scanned: 2,
+        n_col_get_calls: 3,
+    };
     assert_eq!(total.n_unique_kmers, 4);
     assert_eq!(total.n_mphf_calls, 9);
     assert_eq!(total.n_hits, 3);
+    assert_eq!(total.n_columns_scanned, 3);
+    assert_eq!(total.n_col_get_calls, 10);
 }
 
 #[test]
@@ -17,6 +31,8 @@ fn query_stats_default_is_zero() {
     assert_eq!(s.n_unique_kmers, 0);
     assert_eq!(s.n_mphf_calls, 0);
     assert_eq!(s.n_hits, 0);
+    assert_eq!(s.n_columns_scanned, 0);
+    assert_eq!(s.n_col_get_calls, 0);
 }
 
 // ── query_partition_with on a not-yet-indexed partition ─────────────────────
@@ -38,14 +54,12 @@ fn query_partition_with_missing_index_dir_returns_default_stats() {
     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");
+        .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");
 
-    assert_eq!(stats.n_unique_kmers, 0);
-    assert_eq!(stats.n_mphf_calls, 0);
-    assert_eq!(stats.n_hits, 0);
+    assert_eq!(stats, QueryStats::default());
 }
 
 #[test]
@@ -56,8 +70,8 @@ fn query_partition_with_empty_kmers_is_a_noop() {
 
     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");
+        .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");