Introduce unified nonzero_iter API across matrix types

Replaces nested column-major point lookups with a batched iterator that delegates to format-native traversal strategies. The implementation enforces a single pass per matrix type, using row-major iteration for sparse formats and eager collection for packed/columnar layouts while preserving original slot ordering. Memory allocation is optimized by removing `n_cols`-wide buffers in favor of per-row buffering or lazy iteration. Correctness tests verify iterator output against dense baselines across all supported layouts, and architecture documentation is updated to reflect the new format-agnostic query pattern.

src/obicompactvec/src/bitmatrix/persistent.rs

@@ -160,24 +160,8 @@ impl PersistentBitMatrix {
     /// in the same order as `slots`.  Column access is sequential to maximize
     /// cache efficiency on the underlying mmap.
     pub fn sub_matrix(&self, slots: &[usize]) -> Vec<Vec<bool>> {
-        let n_cols = self.n_cols();
-        let mut out: Vec<Vec<bool>> = Vec::with_capacity(n_cols);
-        for c in 0..n_cols {
-            let mut col_buf = vec![false; slots.len()];
-            match self {
-                Self::Columnar(m) => m.col(c).view().fill_batch(slots, &mut col_buf),
-                Self::Packed(m)   => m.col_slice(c).fill_batch(slots, &mut col_buf),
-                Self::Sparse(m)   => {
-                    let mut row_buf = vec![false; m.n_cols()];
-                    for (i, &slot) in slots.iter().enumerate() {
-                        m.fill_row_bool(slot, &mut row_buf);
-                        col_buf[i] = row_buf[c];
-                    }
-                }
-                Self::Implicit { .. } => col_buf.iter_mut().for_each(|b| *b = true),
-            }
-            out.push(col_buf);
-        }
+        let mut out: Vec<Vec<bool>> = (0..self.n_cols()).map(|_| Vec::new()).collect();
+        self.fill_sub_matrix(slots, &mut out);
         out
     }
 
@@ -187,36 +171,68 @@ impl PersistentBitMatrix {
     /// `out` must have length `self.n_cols()`.  Each `out[c]` is cleared,
     /// resized to `slots.len()`, and filled with the values for column `c`
     /// in the same order as `slots`.
+    ///
+    /// Derived from [`nonzero_iter`](Self::nonzero_iter) — drains it,
+    /// scatters into `out[][]` — rather than a second, independently
+    /// maintained per-format traversal: that duplication is exactly how
+    /// this method's `Sparse` branch used to silently reimplement (worse)
+    /// what `PersistentSparseBitMatrix::fill_sub_matrix` already did
+    /// correctly one file over (see `DevDocMD/architecture/siblings.md`,
+    /// "`query` never benefits from sparse row-major access").
     pub fn fill_sub_matrix(&self, slots: &[usize], out: &mut [Vec<bool>]) {
         assert_eq!(out.len(), self.n_cols());
-        let n = slots.len();
-        if n == 0 {
-            for col in out.iter_mut() { col.clear(); }
-            return;
+        for col in out.iter_mut() {
+            col.clear();
+            col.resize(slots.len(), false);
         }
-        let mut perm: Vec<usize> = (0..n).collect();
-        perm.sort_by_key(|&i| slots[i]);
-        let sorted_slots: Vec<usize> = perm.iter().map(|&i| slots[i]).collect();
-        for (c, col) in out.iter_mut().enumerate() {
-            col.resize(n, false);
-            let mut tmp = vec![false; n];
-            match self {
-                Self::Columnar(m) => m.col(c).view().fill_batch_sorted(&sorted_slots, &mut tmp),
-                Self::Packed(m)   => m.col_slice(c).fill_batch_sorted(&sorted_slots, &mut tmp),
-                Self::Sparse(m)   => {
-                    for (i, &orig_idx) in perm.iter().enumerate() {
-                        let slot = sorted_slots[i];
-                        let mut row_buf = vec![false; m.n_cols()];
-                        m.fill_row_bool(slot, &mut row_buf);
-                        col[orig_idx] = row_buf[c];
-                    }
-                }
-                Self::Implicit { .. } => tmp.iter_mut().for_each(|b| *b = true),
+        for (i, c, v) in self.nonzero_iter(slots) {
+            if v != 0 {
+                out[c][i] = true;
             }
-            for (i, &orig_idx) in perm.iter().enumerate() {
-                if !matches!(self, Self::Sparse(_)) {
-                    col[orig_idx] = tmp[i];
+        }
+    }
+
+    /// Yields every nonzero `(idx into slots, col, value)` triple.
+    /// Order is implementation-defined.
+    ///
+    /// One primitive per format, each choosing its own natural traversal:
+    /// `Sparse` delegates to its native row-major decode (no `n_cols`-wide
+    /// buffer, ever — the whole point, see
+    /// `PersistentSparseBitMatrix::nonzero_iter`); `Columnar`/`Packed`
+    /// reuse the same sorted-slot batching `fill_sub_matrix` always used
+    /// (`BitSliceView::nonzero_among_sorted`), just eagerly collected here
+    /// rather than filled into a dense buffer — same `get` calls, same
+    /// mmap locality, no per-format duplication of that logic elsewhere.
+    /// `Implicit` is trivial (one column, always present).
+    ///
+    /// Boxed (not `impl Iterator`) because the match arms are genuinely
+    /// different concrete types — this method isn't on a trait (see
+    /// `traits.rs`'s `BinaryMatrix`, kept RPITIT-free because it's used as
+    /// `dyn BinaryMatrix` in `tests/sparse.rs`), so there's no dyn-safety
+    /// constraint against it; the allocation is one `Box` per call, not
+    /// per cell — negligible next to what it replaces (a fresh `Vec`
+    /// allocated per single-cell `get` in the old `Sparse` path).
+    pub fn nonzero_iter<'a>(&'a self, slots: &'a [usize]) -> Box<dyn Iterator<Item = (usize, usize, u32)> + 'a> {
+        match self {
+            Self::Sparse(m) => Box::new(m.nonzero_iter(slots)),
+            Self::Implicit { .. } => {
+                Box::new(slots.iter().enumerate().map(|(i, _)| (i, 0usize, 1u32)))
+            }
+            Self::Columnar(_) | Self::Packed(_) => {
+                let n = slots.len();
+                let mut perm: Vec<usize> = (0..n).collect();
+                perm.sort_by_key(|&i| slots[i]);
+                let sorted_slots: Vec<usize> = perm.iter().map(|&i| slots[i]).collect();
+                let n_cols = self.n_cols();
+                let mut hits: Vec<(usize, usize, u32)> = Vec::new();
+                for c in 0..n_cols {
+                    let view = self.col_view(c);
+                    hits.extend(
+                        view.nonzero_among_sorted(&sorted_slots)
+                            .map(|pos| (perm[pos], c, 1u32)),
+                    );
                 }
+                Box::new(hits.into_iter())
             }
         }
     }

src/obicompactvec/src/bitmatrix/sparse.rs

@@ -256,6 +256,40 @@ impl PersistentSparseBitMatrix {
             self.for_each_genome_in_row(slot, |g| out[g][i] = true);
         }
     }
+
+    /// Yields every `(idx into slots, genome, 1)` triple, row by row, in
+    /// `slots` order (no sort needed here — unlike the dense formats,
+    /// decode cost per row is independent of slot order). One row buffered
+    /// at a time (`buf`, sized to that row's own cardinality — 1 for a
+    /// singleton, `k̄` for a multi-genome row) via
+    /// [`for_each_genome_in_row`](Self::for_each_genome_in_row): the whole
+    /// reason this exists, vs. going through `fill_sub_matrix` — no
+    /// `n_cols`-wide buffer, ever. See
+    /// `DevDocMD/architecture/siblings.md`, "`query` never benefits from
+    /// sparse row-major access".
+    pub fn nonzero_iter<'a>(&'a self, slots: &'a [usize]) -> impl Iterator<Item = (usize, usize, u32)> + 'a {
+        let mut next_slot = 0usize;
+        let mut cur_idx = 0usize;
+        let mut buf: Vec<usize> = Vec::new();
+        let mut buf_pos = 0usize;
+
+        std::iter::from_fn(move || loop {
+            if buf_pos < buf.len() {
+                let g = buf[buf_pos];
+                buf_pos += 1;
+                return Some((cur_idx, g, 1u32));
+            }
+            if next_slot >= slots.len() {
+                return None;
+            }
+            cur_idx = next_slot;
+            let slot = slots[next_slot];
+            next_slot += 1;
+            buf.clear();
+            self.for_each_genome_in_row(slot, |g| buf.push(g));
+            buf_pos = 0;
+        })
+    }
 }
 
 impl crate::traits::BinaryMatrix for PersistentSparseBitMatrix {

src/obicompactvec/src/intmatrix.rs

@@ -404,6 +404,31 @@ impl PersistentCompactIntMatrix {
         }
     }
 
+    /// Yields every nonzero `(idx into slots, col, value)` triple. Same
+    /// primitive as `PersistentBitMatrix::nonzero_iter` (see
+    /// `DevDocMD/architecture/siblings.md`, "`query` never benefits from
+    /// sparse row-major access" — counts explicitly not excluded from that
+    /// design even though no sparse count format exists yet). No native
+    /// low-effort case here the way `PersistentSparseBitMatrix` has one —
+    /// both variants reuse the same sorted-slot batching `fill_sub_matrix`
+    /// already used, via `IntSliceView::nonzero_among_sorted`.
+    pub fn nonzero_iter<'a>(&'a self, slots: &'a [usize]) -> impl Iterator<Item = (usize, usize, u32)> + 'a {
+        let n = slots.len();
+        let mut perm: Vec<usize> = (0..n).collect();
+        perm.sort_by_key(|&i| slots[i]);
+        let sorted_slots: Vec<usize> = perm.iter().map(|&i| slots[i]).collect();
+        let n_cols = self.n_cols();
+        let mut hits: Vec<(usize, usize, u32)> = Vec::new();
+        for c in 0..n_cols {
+            hits.extend(
+                self.col_view(c)
+                    .nonzero_among_sorted(&sorted_slots)
+                    .map(|(pos, v)| (perm[pos], c, v)),
+            );
+        }
+        hits.into_iter()
+    }
+
     #[inline]
     pub fn sum(&self) -> Array1<u64> {
         match self { Self::Columnar(m) => m.sum(), Self::Packed(m) => m.sum() }

src/obicompactvec/src/tests/intmatrix.rs

@@ -345,6 +345,38 @@ fn col_view_packed_matches_columnar() {
     drop(dir_col);
 }
 
+/// `nonzero_iter` must agree with `row`/`sub_matrix` (already exercised
+/// elsewhere): every nonzero cell among `slots`, no more, no less, values
+/// matching. Compares against both `Columnar` and `Packed` — same data,
+/// two on-disk layouts, one shared `nonzero_iter` implementation.
+#[test]
+fn nonzero_iter_matches_row() {
+    let data: &[&[u32]] = &[&[0, 5, 0, 3, 7], &[2, 0, 0, 4, 0], &[0, 0, 9, 0, 1]];
+    let (dir_col, m_col) = make_matrix(data);
+    let (dir_pack, _) = make_matrix(data);
+    pack_compact_int_matrix(&dir_pack.path().join("counts")).unwrap();
+    let m_pack = PersistentCompactIntMatrix::open(dir_pack.path()).unwrap();
+
+    let slots = [4usize, 0, 3, 1];
+    let mut expected: Vec<(usize, usize, u32)> = Vec::new();
+    for (i, &slot) in slots.iter().enumerate() {
+        for c in 0..data.len() {
+            let v = data[c][slot];
+            if v != 0 {
+                expected.push((i, c, v));
+            }
+        }
+    }
+    expected.sort();
+
+    for (label, m) in [("columnar", &m_col), ("packed", &m_pack)] {
+        let mut got: Vec<(usize, usize, u32)> = m.nonzero_iter(&slots).collect();
+        got.sort();
+        assert_eq!(got, expected, "{label}");
+    }
+    drop(dir_col);
+}
+
 #[test]
 fn partial_relfreq_bray_additive_across_split() {
     // Split rows [1,2,3,4,5] between two matrices; partial sums should add up.

src/obicompactvec/src/tests/sparse.rs

@@ -141,6 +141,64 @@ fn fill_sub_matrix_matches_dense() {
     assert_eq!(sparse_sub, dense_sub);
 }
 
+#[test]
+fn nonzero_iter_matches_dense() {
+    let col0 = [true, false, true, false, true];
+    let col1 = [false, true, true, true, false];
+    let col2 = [true, true, false, false, true];
+    let (_dense_dir, dense) = make_dense(&[&col0, &col1, &col2]);
+    let sparse_root = tempdir().unwrap();
+    let sparse_dir = sparse_root.path().join("sparse");
+    let sparse = PersistentSparseBitMatrixBuilder::build_from_dense(&dense, &sparse_dir)
+        .unwrap()
+        .finish()
+        .unwrap();
+
+    let slots = [4usize, 0, 2];
+    let mut dense_hits: Vec<(usize, usize, u32)> = dense.nonzero_iter(&slots).collect();
+    let mut sparse_hits: Vec<(usize, usize, u32)> = sparse.nonzero_iter(&slots).collect();
+    dense_hits.sort();
+    sparse_hits.sort();
+    assert_eq!(dense_hits, sparse_hits);
+}
+
+/// Targets the dispatch bug directly, not just each concrete type's own
+/// correctness (`fill_sub_matrix_matches_dense` above calls
+/// `PersistentSparseBitMatrix::fill_sub_matrix` straight, never through the
+/// `PersistentBitMatrix::Sparse` enum wrapper — so it could pass even while
+/// the wrapper's own dispatch stayed broken, which is exactly what
+/// happened before `PersistentBitMatrix::fill_sub_matrix`/`nonzero_iter`
+/// were rewritten to drain `nonzero_iter` instead of hand-rolling a second
+/// traversal per format). `PersistentBitMatrix::open` only auto-detects
+/// `Sparse` from a `presence/sparse_meta.json` layout, so build the enum
+/// variant directly here instead of routing through `open`.
+#[test]
+fn enum_wrapper_dispatches_to_native_sparse() {
+    let col0 = [true, false, true, false, true];
+    let col1 = [false, true, true, true, false];
+    let col2 = [true, true, false, false, true];
+    let (_dense_dir, dense) = make_dense(&[&col0, &col1, &col2]);
+    let sparse_root = tempdir().unwrap();
+    let sparse_dir = sparse_root.path().join("sparse");
+    let sparse_inner = PersistentSparseBitMatrixBuilder::build_from_dense(&dense, &sparse_dir)
+        .unwrap()
+        .finish()
+        .unwrap();
+    let sparse_wrapped = PersistentBitMatrix::Sparse(sparse_inner);
+
+    let slots = [4usize, 0, 2];
+
+    let dense_sub = dense.sub_matrix(&slots);
+    let wrapped_sub = sparse_wrapped.sub_matrix(&slots);
+    assert_eq!(wrapped_sub, dense_sub);
+
+    let mut dense_hits: Vec<(usize, usize, u32)> = dense.nonzero_iter(&slots).collect();
+    let mut wrapped_hits: Vec<(usize, usize, u32)> = sparse_wrapped.nonzero_iter(&slots).collect();
+    dense_hits.sort();
+    wrapped_hits.sort();
+    assert_eq!(wrapped_hits, dense_hits);
+}
+
 #[test]
 fn single_genome_matrix() {
     // n_cols=1 — every row is necessarily a singleton (cardinality 1 or

src/obicompactvec/src/views.rs

@@ -58,6 +58,20 @@ impl<'a> BitSliceView<'a> {
             out[i] = self.get(slot);
         }
     }
+
+    /// Positions in `sorted_slots` (not the slots themselves) whose bit is
+    /// set — a `filter` over `get`, not a new traversal. `sorted_slots`
+    /// need not actually be sorted for correctness, but callers batching
+    /// several columns over the same slot set sort once and reuse it for
+    /// cache-friendlier `get` access (see `PersistentBitMatrix::nonzero_iter`).
+    pub(crate) fn nonzero_among_sorted<'b>(self, sorted_slots: &'b [usize]) -> impl Iterator<Item = usize> + 'b
+    where
+        Self: 'b,
+    {
+        sorted_slots.iter().enumerate()
+            .filter_map(move |(pos, &slot)| self.get(slot).then_some(pos))
+    }
+
     pub fn count_ones(&self) -> u64 {
         self.words.iter().map(|w| w.count_ones() as u64).sum()
     }
@@ -200,6 +214,19 @@ impl<'a> IntSliceView<'a> {
         }
     }
 
+    /// Like [`BitSliceView::nonzero_among_sorted`]: `(position in
+    /// sorted_slots, value)` for every nonzero value — a `filter_map` over
+    /// `get`, not a new traversal.
+    pub(crate) fn nonzero_among_sorted<'b>(self, sorted_slots: &'b [usize]) -> impl Iterator<Item = (usize, u32)> + 'b
+    where
+        Self: 'b,
+    {
+        sorted_slots.iter().enumerate().filter_map(move |(pos, &slot)| {
+            let v = self.get(slot);
+            (v != 0).then_some((pos, v))
+        })
+    }
+
     /// Sequential merge scan: yields all n values in slot order.
     #[inline]
     pub fn iter(&self) -> IntSliceViewIter<'a> {

src/obikpartitionner/src/query_layer.rs

@@ -64,13 +64,19 @@ impl QueryLayer {
         }
     }
 
-    /// 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 {
+    /// 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<dyn Iterator<Item = (usize, usize, u32)> + 'a> {
         match self {
-            QueryLayer::Presence(_, mat) => mat.get(g, slot),
-            QueryLayer::Count(_, mat) => mat.col_view(g).get(slot),
+            QueryLayer::Presence(_, mat) => mat.nonzero_iter(slots),
+            QueryLayer::Count(_, mat) => Box::new(mat.nonzero_iter(slots)),
         }
     }
 }
@@ -196,7 +202,12 @@ impl KmerPartition {
             }
         }
 
-        // ── Stage 2: column-major fetch, per layer ───────────────────────────
+        // ── 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() {
             if slots.is_empty() {
                 continue;
@@ -205,14 +216,15 @@ impl KmerPartition {
             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));
-                    }
+            let slot_list: Vec<usize> = slots.keys().copied().collect();
+            for (idx, g, v) in layer.nonzero_iter(&slot_list) {
+                if g >= n_cols {
+                    continue;
                 }
+                stats.n_col_get_calls += 1;
+                debug_assert_ne!(v, 0, "nonzero_iter must not yield zero-valued cells");
+                let descs = slots[&slot_list[idx]];
+                on_event(QueryHit::Value(descs, g, v));
             }
         }