Add cache-optimized batch retrieval and sub-matrix methods

Introduces batch retrieval and sub-matrix extraction methods across vector, view, reader, and matrix types. These implementations optimize cache locality by sorting requested indices for sequential memory access before applying an inverse permutation to restore original order. Includes allocation-free variants that populate caller-provided buffers. Updates architecture documentation to define sibling annex persistence in iteration order and clarify pipeline separation.
This commit is contained in:
Eric Coissac
2026-08-16 14:01:35 +02:00
parent dae543fdfc
commit b1f54b7d2f
9 changed files with 375 additions and 0 deletions
+77
View File
@@ -0,0 +1,77 @@
# Sibling annex — architecture (discussion)
Status: architecture decided (2026-08-14). Implementation not yet mandated.
## Two index spaces, uncorrelated
Every kmer stored in a `Layer` lives in two independent index spaces:
- **Iteration order**: its position when enumerating `unitigs.bin` (the
superkmer file), deterministic but arbitrary with respect to slot.
- **MPHF slot**: `MphfLayer::index(kmer)`, the number the MPHF assigns.
The two are not correlated by any formula. Converting from one to the other
requires either recomputing the MPHF (kmer → slot) or scanning the iteration
stream (kmer → order). There is no `slot → kmer` operation: the MPHF is a
one-way function, not an invertible bijection with a stored inverse. Any
method that reconstructs a kmer from a bare slot number is wrong by
construction, regardless of the mechanism used (MPHF re-hash, or evidence
decode + direct unitig read). See `MphfLayer::kmer_at`
(`obilayeredmap/src/mphf_layer.rs`) — flagged for removal, currently called
from `obikindex/siblings/build.rs:122` and `family_scan.rs:173`.
## Two pipelines, never mixed
| | origin of the kmer | membership known? | correct mapping |
|---|---|---|---|
| **query pipeline** | external (caller-supplied) | no | `query`/`find`/`find_strict` — MPHF + evidence check |
| **iteration pipeline** | enumerated from this layer's own `unitigs.bin` | yes, by construction | `index`/`index_batch` — MPHF only, no evidence |
Evidence exists solely to answer "is this external kmer a member of the
layer" for the query pipeline. Using it (or the MPHF) to go the other way —
recover a kmer from a slot, or re-verify a kmer that was just produced by
iterating the layer — is a conceptual error: evidence can be probabilistic
(`Approx` mode), so any slot→kmer attempt is unsound in general, and
pointless even in `Exact`/`Hybrid` mode since the kmer was already known.
## Sibling annex: an iteration-pipeline artifact only
The sibling annex (`FamilyMask`/`SiblingAnnex`, `.psib`,
`obicompactvec/src/siblingannex.rs`) records, per kmer, whether it is a
family minorant and which family members are present in the index. Its only
consumers (`obikindex/siblings/stats.rs`, `family_scan.rs`) enumerate it
exhaustively (`0..annex.len()`); no query-pipeline code path touches it.
**Decision**: the annex must be persisted in iteration order, not slot
order. This lets readers zip-iterate `Layer::iter_kmers()` and the annex
file directly — one linear, cache-friendly pass, no MPHF/slot indirection,
no `kmer_at`. It also enables specialized iterators building on this zip:
minorants-only iteration, batch-of-kmers → batch-of-family-members, etc.
Today the annex is built and stored in **slot** order
(`build_layer_sibling_annex`, `siblings/build.rs`): `slot_kmer` is populated
via `(0..n_slots).map(|slot| mphf.kmer_at(slot))`, and the origin `slot` is
threaded through the whole cross-partition reconciliation pipeline (variant
generation, `query_partition_with`, final `mask[slot].fetch_or(...)`). This
must change to iterating `iter_kmers()`/`enumerate_kmers()` and threading
the **iteration index** instead of the slot end to end — eliminating
`kmer_at` from the build path entirely, not just the read path. No
slot-indexed intermediate is needed even during construction; the
iteration-order id is sufficient throughout.
The cross-partition side of the same pipeline is unaffected: checking
whether a generated family-variant kmer exists in another partition is a
genuine query-pipeline operation (the variant's membership in the *target*
partition is unknown) and must keep going through
`KmerPartition::query_partition_with` (MPHF + evidence), never a raw
`index()`.
## Pending work
- Remove `MphfLayer::kmer_at`.
- `siblings/build.rs`: build `slot_kmer`-equivalent via iteration, not
`kmer_at`; thread iteration index instead of slot through the
variant/reconciliation pipeline; persist the annex in iteration order.
- `siblings/family_scan.rs`, `stats.rs`: read the annex via zipped
iteration (`iter_kmers().zip(annex_iter)`) instead of `0..annex.len()` +
`kmer_at`.
+1
View File
@@ -58,6 +58,7 @@ nav:
- Architecture: - Architecture:
- Sequences: architecture/sequences/invariant.md - Sequences: architecture/sequences/invariant.md
- Kmer index: architecture/index_architecture.md - Kmer index: architecture/index_architecture.md
- Sibling annex (discussion): architecture/siblings.md
- NUMA-aware worker pools: architecture/numa_worker_pools.md - NUMA-aware worker pools: architecture/numa_worker_pools.md
- NUMA-aware partition runner: architecture/numa_partition_runner.md - NUMA-aware partition runner: architecture/numa_partition_runner.md
+26
View File
@@ -30,3 +30,29 @@ L'API du `Layer<D>` expose maintenant quatre itérateurs sur les kmers du layer
- `enumerate_kmers_batch(&self, n: usize) -> Enumerate<KmerBatchIter<'_>>` — variante indexée retournant `(usize, Vec<CanonicalKmer>)`, où l'index correspond au premier kmer du batch. Construit par composition sur `iter_kmers_batch()`. - `enumerate_kmers_batch(&self, n: usize) -> Enumerate<KmerBatchIter<'_>>` — variante indexée retournant `(usize, Vec<CanonicalKmer>)`, où l'index correspond au premier kmer du batch. Construit par composition sur `iter_kmers_batch()`.
L'itération est implémentée via des structs `KmerIter` et `KmerBatchIter` qui encapsulent l'itérateur interne de `UnitigFileReader::iter_indexed_canonical_kmers()`. Les structs sont publics et peuvent être stockés, transmis ou combinés avec d'autres adaptateurs d'itérateur. L'itération est implémentée via des structs `KmerIter` et `KmerBatchIter` qui encapsulent l'itérateur interne de `UnitigFileReader::iter_indexed_canonical_kmers()`. Les structs sont publics et peuvent être stockés, transmis ou combinés avec d'autres adaptateurs d'itérateur.
## Performance — itérateurs de batch
Les itérateurs `iter_kmers_batch` et `enumerate_kmers_batch` allouent un nouveau `Vec` à chaque batch. Pour des tailles de batch importantes ou des chemins critiques, cela peut générer une pression malloc significative.
**À concevoir** : un système de pool de vecteurs avec réallocation automatique dès qu'un buffer n'est plus référencé, pour réutiliser les allocations entre batches et réduire le nombre d'appels système. Ce pool pourrait être intégré à `KmerBatchIter` ou proposé comme un adaptateur d'itérateur générique.
## Accès aux matrices de présence/comptage
Les types de vecteurs persistants exposent maintenant des méthodes de lookup batch optimisées pour l'accès séquentiel au mmap :
- `PersistentCompactIntVec::get_batch(slots) -> Vec<u32>` et `fill_batch(slots, out)`
- `PersistentBitVec::get_batch(slots) -> Vec<bool>` et `fill_batch(slots, out)`
- `IntSliceView::get_batch(slots) -> Vec<u32>` et `fill_batch(slots, out)`
- `BitSliceView::get_batch(slots) -> Vec<bool>` et `fill_batch(slots, out)`
Toutes ces méthodes trient les slots en interne pour un accès mmap séquentiel, puis réordonnent les résultats selon l'ordre d'origine. `fill_batch` évite l'allocation en remplissant un buffer fourni par le caller.
Les matrices persistantes exposent `sub_matrix(slots) -> Vec<Vec<T>>` et `fill_sub_matrix(slots, out)` :
- `PersistentCompactIntMatrix::sub_matrix(slots)` — retourne `Vec<Vec<u32>>` column-first
- `PersistentCompactIntMatrix::fill_sub_matrix(slots, out: &mut [Vec<u32>])` — remplit des buffers fournis, réutilise les allocations existantes
- `PersistentBitMatrix::sub_matrix(slots)` — retourne `Vec<Vec<bool>>` column-first
- `PersistentBitMatrix::fill_sub_matrix(slots, out: &mut [Vec<bool>])` — remplit des buffers fournis, réutilise les allocations existantes
Le scan est colonne-first pour respecter la layout mémoire. `fill_sub_matrix` trie les slots une seule fois, puis appelle `fill_batch_sorted` sur chaque colonne pour éviter le tri redondant. Ces méthodes sont également disponibles sur `Layer<PersistentCompactIntMatrix>` et `Layer<PersistentBitMatrix>`.
@@ -126,6 +126,57 @@ impl PersistentBitMatrix {
} }
} }
/// Extract a sub-matrix containing only the rows at `slots`.
///
/// Returns a column-first `Vec<Vec<bool>>`: the outer `Vec` has one entry per
/// column, and each inner `Vec` contains the values for the requested rows
/// 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::Implicit { .. } => col_buf.iter_mut().for_each(|b| *b = true),
}
out.push(col_buf);
}
out
}
/// Like [`sub_matrix`](Self::sub_matrix), but fills caller-provided column
/// buffers to avoid allocating the outer `Vec`.
///
/// `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`.
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;
}
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::Implicit { .. } => tmp.iter_mut().for_each(|b| *b = true),
}
for (i, &orig_idx) in perm.iter().enumerate() {
col[orig_idx] = tmp[i];
}
}
}
pub fn count_ones(&self) -> Array1<u64> { pub fn count_ones(&self) -> Array1<u64> {
match self { match self {
Self::Columnar(m) => m.count_ones(), Self::Columnar(m) => m.count_ones(),
+36
View File
@@ -48,6 +48,42 @@ impl PersistentBitVec {
(self.mmap[HEADER_SIZE + (slot >> 3)] >> (slot & 7)) & 1 != 0 (self.mmap[HEADER_SIZE + (slot >> 3)] >> (slot & 7)) & 1 != 0
} }
/// Batch lookup: read bits at `slots` in an order that minimizes
/// cache misses.
///
/// The slots are sorted internally before reading so that accesses to
/// the underlying mmap are as sequential as possible, then the results
/// are reordered to match the input order.
pub fn get_batch(&self, slots: &[usize]) -> Vec<bool> {
let mut out = vec![false; slots.len()];
self.fill_batch(slots, &mut out);
out
}
/// Like [`get_batch`](Self::get_batch), but fills a caller-provided buffer.
pub fn fill_batch(&self, slots: &[usize], out: &mut [bool]) {
assert_eq!(slots.len(), out.len());
let n = slots.len();
if n == 0 { return; }
let mut perm: Vec<usize> = (0..n).collect();
perm.sort_by_key(|&i| slots[i]);
let sorted: Vec<usize> = perm.iter().map(|&i| slots[i]).collect();
let mut tmp = vec![false; n];
self.fill_batch_sorted(&sorted, &mut tmp);
for (i, &orig_idx) in perm.iter().enumerate() {
out[orig_idx] = tmp[i];
}
}
/// Fill `out` assuming `sorted_slots` is already in ascending order.
/// Results are written in `sorted_slots` order (no reordering).
pub(crate) fn fill_batch_sorted(&self, sorted_slots: &[usize], out: &mut [bool]) {
assert_eq!(sorted_slots.len(), out.len());
for (i, &slot) in sorted_slots.iter().enumerate() {
out[i] = self.get(slot);
}
}
// SAFETY: mmap is page-aligned, HEADER_SIZE=16 divisible by 8 → u64-aligned. // SAFETY: mmap is page-aligned, HEADER_SIZE=16 divisible by 8 → u64-aligned.
fn data_words(&self) -> &[u64] { fn data_words(&self) -> &[u64] {
let nw = n_words(self.n); let nw = n_words(self.n);
+44
View File
@@ -346,6 +346,50 @@ impl PersistentCompactIntMatrix {
pub fn fill_row(&self, slot: usize, buf: &mut [u32]) { pub fn fill_row(&self, slot: usize, buf: &mut [u32]) {
match self { Self::Columnar(m) => m.fill_row(slot, buf), Self::Packed(m) => m.fill_row(slot, buf) } match self { Self::Columnar(m) => m.fill_row(slot, buf), Self::Packed(m) => m.fill_row(slot, buf) }
} }
/// Extract a sub-matrix containing only the rows at `slots`.
///
/// Returns a column-first `Vec<Vec<u32>>`: the outer `Vec` has one entry per
/// column, and each inner `Vec` contains the values for the requested rows
/// 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<u32>> {
let n_cols = self.n_cols();
let mut out: Vec<Vec<u32>> = Vec::with_capacity(n_cols);
for c in 0..n_cols {
let mut col_buf = vec![0u32; slots.len()];
self.col_view(c).fill_batch(slots, &mut col_buf);
out.push(col_buf);
}
out
}
/// Like [`sub_matrix`](Self::sub_matrix), but fills caller-provided column
/// buffers to avoid allocating the outer `Vec`.
///
/// `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`.
pub fn fill_sub_matrix(&self, slots: &[usize], out: &mut [Vec<u32>]) {
assert_eq!(out.len(), self.n_cols());
let n = slots.len();
if n == 0 {
for col in out.iter_mut() { col.clear(); }
return;
}
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, 0);
let mut tmp = vec![0u32; n];
self.col_view(c).fill_batch_sorted(&sorted_slots, &mut tmp);
for (i, &orig_idx) in perm.iter().enumerate() {
col[orig_idx] = tmp[i];
}
}
}
pub fn sum(&self) -> Array1<u64> { pub fn sum(&self) -> Array1<u64> {
match self { Self::Columnar(m) => m.sum(), Self::Packed(m) => m.sum() } match self { Self::Columnar(m) => m.sum(), Self::Packed(m) => m.sum() }
} }
+36
View File
@@ -57,6 +57,42 @@ impl PersistentCompactIntVec {
} }
} }
/// Batch lookup: read values at `slots` in an order that minimizes
/// cache misses.
///
/// The slots are sorted internally before reading so that accesses to
/// the underlying mmap are as sequential as possible, then the results
/// are reordered to match the input order.
pub fn get_batch(&self, slots: &[usize]) -> Vec<u32> {
let mut out = vec![0u32; slots.len()];
self.fill_batch(slots, &mut out);
out
}
/// Like [`get_batch`](Self::get_batch), but fills a caller-provided buffer.
pub fn fill_batch(&self, slots: &[usize], out: &mut [u32]) {
assert_eq!(slots.len(), out.len());
let n = slots.len();
if n == 0 { return; }
let mut perm: Vec<usize> = (0..n).collect();
perm.sort_by_key(|&i| slots[i]);
let sorted: Vec<usize> = perm.iter().map(|&i| slots[i]).collect();
let mut tmp = vec![0u32; n];
self.fill_batch_sorted(&sorted, &mut tmp);
for (i, &orig_idx) in perm.iter().enumerate() {
out[orig_idx] = tmp[i];
}
}
/// Fill `out` assuming `sorted_slots` is already in ascending order.
/// Results are written in `sorted_slots` order (no reordering).
pub(crate) fn fill_batch_sorted(&self, sorted_slots: &[usize], out: &mut [u32]) {
assert_eq!(sorted_slots.len(), out.len());
for (i, &slot) in sorted_slots.iter().enumerate() {
out[i] = self.get(slot);
}
}
fn overflow_get(&self, slot: usize) -> u32 { fn overflow_get(&self, slot: usize) -> u32 {
let (pos_start, pos_end) = if self.step == 0 { let (pos_start, pos_end) = if self.step == 0 {
(0, self.n_overflow) (0, self.n_overflow)
+66
View File
@@ -23,6 +23,38 @@ impl<'a> BitSliceView<'a> {
(self.words[slot >> 6] >> (slot & 63)) & 1 != 0 (self.words[slot >> 6] >> (slot & 63)) & 1 != 0
} }
/// Batch lookup: read bits at `slots` with cache-friendly access pattern.
///
/// Slots are sorted internally before reading, then results are reordered
/// to match the input order.
pub fn get_batch(&self, slots: &[usize]) -> Vec<bool> {
let mut out = vec![false; slots.len()];
self.fill_batch(slots, &mut out);
out
}
/// Like [`get_batch`](Self::get_batch), but fills a caller-provided buffer.
pub fn fill_batch(&self, slots: &[usize], out: &mut [bool]) {
assert_eq!(slots.len(), out.len());
let n = slots.len();
if n == 0 { return; }
let mut perm: Vec<usize> = (0..n).collect();
perm.sort_by_key(|&i| slots[i]);
let sorted: Vec<usize> = perm.iter().map(|&i| slots[i]).collect();
let mut tmp = vec![false; n];
self.fill_batch_sorted(&sorted, &mut tmp);
for (i, &orig_idx) in perm.iter().enumerate() {
out[orig_idx] = tmp[i];
}
}
/// Fill `out` assuming `sorted_slots` is already in ascending order.
/// Results are written in `sorted_slots` order (no reordering).
pub(crate) fn fill_batch_sorted(&self, sorted_slots: &[usize], out: &mut [bool]) {
assert_eq!(sorted_slots.len(), out.len());
for (i, &slot) in sorted_slots.iter().enumerate() {
out[i] = self.get(slot);
}
}
pub fn count_ones(&self) -> u64 { pub fn count_ones(&self) -> u64 {
self.words.iter().map(|w| w.count_ones() as u64).sum() self.words.iter().map(|w| w.count_ones() as u64).sum()
} }
@@ -123,6 +155,40 @@ impl<'a> IntSliceView<'a> {
panic!("slot {slot} marked overflow but not found") panic!("slot {slot} marked overflow but not found")
} }
/// Batch lookup: read values at `slots` with cache-friendly access pattern.
///
/// Slots are sorted internally before reading, then results are reordered
/// to match the input order.
pub fn get_batch(&self, slots: &[usize]) -> Vec<u32> {
let mut out = vec![0u32; slots.len()];
self.fill_batch(slots, &mut out);
out
}
/// Like [`get_batch`](Self::get_batch), but fills a caller-provided buffer.
pub fn fill_batch(&self, slots: &[usize], out: &mut [u32]) {
assert_eq!(slots.len(), out.len());
let n = slots.len();
if n == 0 { return; }
let mut perm: Vec<usize> = (0..n).collect();
perm.sort_by_key(|&i| slots[i]);
let sorted: Vec<usize> = perm.iter().map(|&i| slots[i]).collect();
let mut tmp = vec![0u32; n];
self.fill_batch_sorted(&sorted, &mut tmp);
for (i, &orig_idx) in perm.iter().enumerate() {
out[orig_idx] = tmp[i];
}
}
/// Fill `out` assuming `sorted_slots` is already in ascending order.
/// Results are written in `sorted_slots` order (no reordering).
pub(crate) fn fill_batch_sorted(&self, sorted_slots: &[usize], out: &mut [u32]) {
assert_eq!(sorted_slots.len(), out.len());
for (i, &slot) in sorted_slots.iter().enumerate() {
out[i] = self.get(slot);
}
}
/// Sequential merge scan: yields all n values in slot order. /// Sequential merge scan: yields all n values in slot order.
pub fn iter(&self) -> IntSliceViewIter<'a> { pub fn iter(&self) -> IntSliceViewIter<'a> {
IntSliceViewIter { IntSliceViewIter {
+38
View File
@@ -186,6 +186,25 @@ impl Layer<PersistentCompactIntMatrix> {
PersistentCompactIntMatrix::append_column(&layer_dir.join(COUNTS_DIR), value_of) PersistentCompactIntMatrix::append_column(&layer_dir.join(COUNTS_DIR), value_of)
.map_err(OLMError::Io) .map_err(OLMError::Io)
} }
/// Extract a sub-matrix of counts for the rows at `slots`.
///
/// Returns a column-first `Vec<Vec<u32>>`: one inner `Vec` per genome
/// column, containing the counts for the requested slots in order.
/// Column access is sequential for cache efficiency.
pub fn sub_matrix(&self, slots: &[usize]) -> Vec<Vec<u32>> {
self.data.sub_matrix(slots)
}
/// Like [`sub_matrix`](Self::sub_matrix), but fills caller-provided column
/// buffers to avoid allocating the outer `Vec`.
///
/// `out` must have length equal to the number of genome columns. Each
/// `out[c]` is cleared, resized to `slots.len()`, and filled with the
/// counts for column `c` in the same order as `slots`.
pub fn fill_sub_matrix(&self, slots: &[usize], out: &mut [Vec<u32>]) {
self.data.fill_sub_matrix(slots, out)
}
} }
// ── Mode 3 — presence/absence matrix ───────────────────────────────────────── // ── Mode 3 — presence/absence matrix ─────────────────────────────────────────
@@ -199,6 +218,25 @@ impl Layer<PersistentBitMatrix> {
.map_err(OLMError::Io) .map_err(OLMError::Io)
} }
/// Extract a sub-matrix of presence/absence for the rows at `slots`.
///
/// Returns a column-first `Vec<Vec<bool>>`: one inner `Vec` per genome
/// column, containing the presence bits for the requested slots in order.
/// Column access is sequential for cache efficiency.
pub fn sub_matrix(&self, slots: &[usize]) -> Vec<Vec<bool>> {
self.data.sub_matrix(slots)
}
/// Like [`sub_matrix`](Self::sub_matrix), but fills caller-provided column
/// buffers to avoid allocating the outer `Vec`.
///
/// `out` must have length equal to the number of genome columns. Each
/// `out[c]` is cleared, resized to `slots.len()`, and filled with the
/// presence bits for column `c` in the same order as `slots`.
pub fn fill_sub_matrix(&self, slots: &[usize], out: &mut [Vec<bool>]) {
self.data.fill_sub_matrix(slots, out)
}
pub fn build_presence( pub fn build_presence(
out_dir: &Path, out_dir: &Path,
block_bits: u8, block_bits: u8,