src/obikindex/src/siblings/alignment.rs

@@ -78,30 +78,24 @@ impl KmerIndex {
         // lookups by partition internally, but interleave those sweeps
         // across layers at the OS level, scattering page-cache access over
         // every partition at once again and defeating the whole point of
-        // the grouping. `Vec<Vec<u8>>` per layer, one entry (column) per
+        // the grouping. Columns appended straight into `sequences` as each
-        // variable family, appended in layer order for a single
+        // family comes back from `scan_layer_families` (bounded-batch, not
-        // deterministic column order across the whole index.
+        // whole-layer) — no intermediate per-layer buffer, layer order
-        let mut partials: Vec<Vec<Vec<u8>>> = Vec::with_capacity(layer_dirs.len());
+        // giving a single deterministic column order across the whole index.
+        let mut sequences: Vec<Vec<u8>> = vec![Vec::new(); n_genomes];
         for layer_dir in &layer_dirs {
-            let families = scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache)?;
+            scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, |mask, genome_mask| {
-            let columns = families
+                if mask.family_size() < 2 {
-                .into_iter()
+                    return; // monomorphic family — no signal, skip
-                .filter(|f| f.mask.family_size() >= 2) // monomorphic family — no signal, skip
+                }
-                .map(|f| f.genome_mask.iter().map(|&m| iupac_code(m)).collect())
+                for (g, &m) in genome_mask.iter().enumerate() {
-                .collect();
+                    sequences[g].push(iupac_code(m));
-            partials.push(columns);
+                }
+            })?;
             pb.inc(1);
         }
         pb.finish_and_clear();
 
-        let mut sequences: Vec<Vec<u8>> = vec![Vec::new(); n_genomes];
-        for layer_columns in partials {
-            for column in layer_columns {
-                for (g, &code) in column.iter().enumerate() {
-                    sequences[g].push(code);
-                }
-            }
-        }
         Ok(SnpAlignment { sequences })
     }
 }

src/obikindex/src/siblings/cardinality.rs

@@ -78,14 +78,12 @@ impl KmerIndex {
         let pb = progress_bar("cardinality_tally", layer_dirs.len() as u64, "layers");
         // One layer at a time — see `snp_pseudo_alignment`'s comment for why
         // `par_iter()` over layers would defeat `scan_layer_families`'s
-        // partition-grouped locality.
+        // partition-grouped locality. Tallied straight into `total` as each
-        let mut partials: Vec<[[u64; 5]; 5]> = Vec::with_capacity(layer_dirs.len());
+        // family comes back — no per-layer buffer.
+        let mut total = [[0u64; 5]; 5];
         for layer_dir in &layer_dirs {
-            let families = scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache)?;
+            scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, |mask, genome_mask| {
-            let mut counts = [[0u64; 5]; 5];
+                if mask.family_size() < 2 {
-
-            for family in &families {
-                if family.mask.family_size() < 2 {
                     // Fully invariant family (never varies anywhere in
                     // the index) — genome-wide background, not
                     // SNP-adjacent signal; would otherwise swamp the
@@ -94,10 +92,9 @@ impl KmerIndex {
                     // `+ASC`-corrected alignment/likelihood actually
                     // models. See `base_pair_tally`'s `variable` gate
                     // on its own `same` diagonal for the matching fix.
-                    continue;
+                    return;
                 }
 
-                let genome_mask = &family.genome_mask;
                 for i in 0..n_genomes {
                     let card_i = genome_mask[i].count_ones() as usize;
                     for j in (i + 1)..n_genomes {
@@ -105,27 +102,17 @@ impl KmerIndex {
                             continue;
                         }
                         let card_j = genome_mask[j].count_ones() as usize;
-                        counts[card_i][card_j] += 1;
+                        total[card_i][card_j] += 1;
                         if card_i != card_j {
-                            counts[card_j][card_i] += 1;
+                            total[card_j][card_i] += 1;
                         }
                     }
                 }
-            }
+            })?;
-
-            partials.push(counts);
             pb.inc(1);
         }
         pb.finish_and_clear();
 
-        let mut total = [[0u64; 5]; 5];
-        for partial in partials {
-            for a in 0..5 {
-                for b in 0..5 {
-                    total[a][b] += partial[a][b];
-                }
-            }
-        }
         Ok(CardinalityTally { counts: total })
     }
 }

src/obikindex/src/siblings/distance.rs

@@ -88,12 +88,10 @@ impl KmerIndex {
         // partition-grouped locality.
         let mut total = zero();
         for layer_dir in &layer_dirs {
-            let families = scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache)?;
             let mut acc = zero();
 
-            for family in &families {
+            scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, |mask, genome_mask| {
-                let variable = family.mask.family_size() >= 2;
+                let variable = mask.family_size() >= 2;
-                let genome_mask = &family.genome_mask;
 
                 // Per genome: which single form (if exactly one) it
                 // carries — `None` (a `popcount != 1` mask) once a
@@ -111,7 +109,7 @@ impl KmerIndex {
                         on_pair(&mut acc, i, j, bi, bj, variable);
                     }
                 }
-            }
+            })?;
 
             total = combine(total, acc);
             pb.inc(1);

src/obikindex/src/siblings/family_scan.rs

@@ -5,11 +5,22 @@
 //! `build_sibling_annex` (see `cache.rs`/`build.rs`'s docs): cross-partition
 //! lookups are grouped by destination partition and resolved in one
 //! contiguous sweep per partition, instead of one lookup at a time jumping
-//! between partitions in family order. The naive, per-family inline lookup
+//! between partitions in family order.
-//! was the shape all four consumers used to have independently; sampling a
+//!
-//! real ~170 GB index against 137 GB of RAM showed it thrashing the page
+//! Processed in bounded batches of [`FAMILY_BATCH`] families, not the whole
-//! cache (near-continuous ~2 GB/s pagein), the same failure mode
+//! layer at once — same rationale and constant as `build_layer_sibling_annex`'s
-//! `build_sibling_annex` had already been fixed for.
+//! `BATCH_SIZE`. An earlier version of this function materialised every
+//! family of the whole layer up front (`Vec<FamilyRow>`, each with its own
+//! heap-allocated `Vec<u8>`) before handing anything back to the caller:
+//! O(layer's family count × genome count) extra memory, on top of a second
+//! full copy for the owned per-family vectors — measured pushing a real run
+//! into heavy VM compression/swap even though the per-lookup I/O pattern was
+//! already fixed. A batch of a few thousand families needs the same
+//! partition-grouping to keep lookups local; it does not need the entire
+//! layer resident at once. Callers get each family's `genome_mask` through
+//! `on_family`, backed by one reused scratch buffer — no per-family
+//! allocation at all, matching the O(genome count) working set the original,
+//! pre-locality-fix per-family code had.
 
 use std::path::{Path, PathBuf};
 use std::sync::atomic::{AtomicU8, Ordering};
@@ -29,14 +40,12 @@ use super::cache::{Mat, PartitionCache};
 use super::helpers::{central_base, is_minorant, partition_of};
 use super::{olm_to_ok, ANNEX_FILE_NAME, INDEX_SUBDIR};
 
-/// One layer's minorant family, with its own annex mask and, per genome,
+/// Families per batch — bounds `scan_layer_families`'s extra memory to
-/// the 4-bit base-presence mask resolved across the whole index (bit `b`
+/// `FAMILY_BATCH * n_genomes` bytes regardless of layer size. Same value and
-/// set iff that genome carries the family member whose own canonical
+/// rationale as `build_layer_sibling_annex`'s `BATCH_SIZE`: large enough to
-/// central base is `b`).
+/// amortise the per-batch partition-grouping/sync cost, small enough to keep
-pub(super) struct FamilyRow {
+/// working memory a rounding error next to the index itself.
-    pub mask: FamilyMask,
+const FAMILY_BATCH: usize = 4096;
-    pub genome_mask: Vec<u8>,
-}
 
 impl KmerIndex {
     /// Every (partition, layer) directory carrying a sibling annex, checked
@@ -66,17 +75,13 @@ impl KmerIndex {
     }
 }
 
-/// Every minorant family of one layer, each with its per-genome
+/// Visits every minorant family of one layer, in slot order, batched
-/// base-presence mask resolved against the whole (already-open) partition
+/// [`FAMILY_BATCH`] at a time — see the module docs for why. `on_family` is
-/// cache.
+/// called once per family with its own annex mask and its per-genome
-///
+/// base-presence (`genome_mask[g]`: bit `b` set iff genome `g` carries the
-/// Two passes, not one inline pass: the first is purely local (already-open
+/// member whose own canonical central base is `b`), backed by a scratch
-/// annex/mphf/matrix, no lookups) and both enumerates the layer's families
+/// buffer reused across every call — callers that need to keep data past
-/// and collects every cross-partition query they need, grouped by
+/// the call must copy it themselves.
-/// destination partition; the second resolves each partition's whole batch
-/// in one contiguous sweep, keeping that partition's mmap'd pages hot for
-/// its entire batch instead of faulting them in and out as lookups jump
-/// between partitions in family order.
 pub(super) fn scan_layer_families(
     layer_dir: &Path,
     n_parts: usize,
@@ -84,7 +89,8 @@ pub(super) fn scan_layer_families(
     with_counts: bool,
     k: usize,
     cache: &PartitionCache,
-) -> OKIResult<Vec<FamilyRow>> {
+    mut on_family: impl FnMut(FamilyMask, &[u8]),
+) -> OKIResult<()> {
     let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir");
     let meta = PartitionMeta::load(index_dir).map_err(olm_to_ok)?;
     let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?;
@@ -106,73 +112,77 @@ pub(super) fn scan_layer_families(
     };
     let n_cols = mat.n_cols().min(n_genomes);
 
-    // ── Pass 1: local only — enumerate minorant families and collect every
+    let n_slots = annex.len();
-    // cross-partition query they need, grouped by destination partition.
+    let mut slot = 0usize;
-    struct Family {
+    let mut batch_slots: Vec<usize> = Vec::with_capacity(FAMILY_BATCH);
-        slot: usize,
+    let mut scratch = vec![0u8; n_genomes]; // reused across every `on_family` call — no per-family allocation
-        mask: FamilyMask,
-    }
-    let mut families: Vec<Family> = Vec::new();
-    let mut outgoing: Vec<Vec<(CanonicalKmer, usize, u8)>> = (0..n_parts).map(|_| Vec::new()).collect();
 
-    for slot in 0..annex.len() {
+    while slot < n_slots {
-        let Some(mask) = annex.get(slot) else { continue };
+        batch_slots.clear();
-        let Some(kmer) = slot_kmer[slot] else { continue };
+        while slot < n_slots && batch_slots.len() < FAMILY_BATCH {
-        if !is_minorant(kmer, mask, k) {
+            let s = slot;
-            continue; // family tallied once, at its minorant
+            slot += 1;
-        }
+            let Some(mask) = annex.get(s) else { continue };
-        let family_idx = families.len();
+            let Some(kmer) = slot_kmer[s] else { continue };
-        families.push(Family { slot, mask });
+            if !is_minorant(kmer, mask, k) {
-
+                continue; // family tallied once, at its minorant
-        for other in kmer.central_canonical_neighbors() {
-            if other == kmer {
-                continue; // local — resolved below straight from `mat`, no lookup
             }
-            let base = central_base(other, k);
+            batch_slots.push(s);
-            if !mask.has(base) {
-                continue;
-            }
-            let dest = partition_of(other, n_parts);
-            outgoing[dest].push((other, family_idx, base));
         }
-    }
+        if batch_slots.is_empty() {
-
+            continue; // ran out of minorant slots before filling a batch — loop condition ends it
-    let genome_mask: Vec<AtomicU8> = (0..families.len() * n_genomes).map(|_| AtomicU8::new(0)).collect();
-
-    // Local presence (the family's own minorant slot) — already open, no
-    // lookup, so no locality concern either way.
-    for (family_idx, family) in families.iter().enumerate() {
-        let kmer = slot_kmer[family.slot].expect("family slot has a kmer");
-        let base = central_base(kmer, k);
-        for g in 0..n_cols {
-            if mat.carries(g, family.slot) {
-                genome_mask[family_idx * n_genomes + g].fetch_or(1 << base, Ordering::Relaxed);
-            }
         }
-    }
 
-    // ── Pass 2: resolve each destination partition's whole batch in one
+        // ── Local only: seed each family's own minorant presence (already
-    // contiguous sweep — same rationale as `build_sibling_annex`'s
+        // open, no lookup) and collect this batch's cross-partition
-    // `outgoing` grouping.
+        // queries, grouped by destination partition.
-    outgoing.par_iter().enumerate().filter(|(_, q)| !q.is_empty()).for_each(|(dest, queries)| {
+        let mut outgoing: Vec<Vec<(CanonicalKmer, usize, u8)>> = (0..n_parts).map(|_| Vec::new()).collect();
-        for &(variant, family_idx, base) in queries {
+        let genome_mask: Vec<AtomicU8> = (0..batch_slots.len() * n_genomes).map(|_| AtomicU8::new(0)).collect();
-            let Some(presence) = cache.find_presence(dest, variant, n_genomes) else { continue };
+
-            for (g, &present) in presence.iter().enumerate() {
+        for (family_idx, &fslot) in batch_slots.iter().enumerate() {
-                if present {
+            let kmer = slot_kmer[fslot].expect("batch slot has a kmer");
+            let mask = annex.get(fslot).expect("batch slot has a mask");
+            let base = central_base(kmer, k);
+            for g in 0..n_cols {
+                if mat.carries(g, fslot) {
                     genome_mask[family_idx * n_genomes + g].fetch_or(1 << base, Ordering::Relaxed);
                 }
             }
+            for other in kmer.central_canonical_neighbors() {
+                if other == kmer {
+                    continue; // local — resolved above straight from `mat`, no lookup
+                }
+                let b = central_base(other, k);
+                if !mask.has(b) {
+                    continue;
+                }
+                let dest = partition_of(other, n_parts);
+                outgoing[dest].push((other, family_idx, b));
+            }
         }
-    });
 
-    Ok(families
+        // ── Resolve this batch's queries, one contiguous sweep per
-        .into_iter()
+        // destination partition — same rationale as `build_sibling_annex`'s
-        .enumerate()
+        // `outgoing` grouping.
-        .map(|(family_idx, family)| FamilyRow {
+        outgoing.par_iter().enumerate().filter(|(_, q)| !q.is_empty()).for_each(|(dest, queries)| {
-            mask: family.mask,
+            for &(variant, family_idx, base) in queries {
-            genome_mask: (0..n_genomes)
+                let Some(presence) = cache.find_presence(dest, variant, n_genomes) else { continue };
-                .map(|g| genome_mask[family_idx * n_genomes + g].load(Ordering::Relaxed))
+                for (g, &present) in presence.iter().enumerate() {
-                .collect(),
+                    if present {
-        })
+                        genome_mask[family_idx * n_genomes + g].fetch_or(1 << base, Ordering::Relaxed);
-        .collect())
+                    }
+                }
+            }
+        });
+
+        for (family_idx, &fslot) in batch_slots.iter().enumerate() {
+            let mask = annex.get(fslot).expect("batch slot has a mask");
+            for (g, dst) in scratch.iter_mut().enumerate() {
+                *dst = genome_mask[family_idx * n_genomes + g].load(Ordering::Relaxed);
+            }
+            on_family(mask, &scratch);
+        }
+    }
+
+    Ok(())
 }

src/obikindex/src/siblings/stats.rs

@@ -65,19 +65,18 @@ impl KmerIndex {
             ..Default::default()
         };
         for layer_dir in &layer_dirs {
-            let families = scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache)?;
+            scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, |mask, genome_mask| {
-            for family in &families {
                 // "Genome g represents this family" means g carries
                 // *any* of its members, not just the minorant's own —
                 // `genome_mask[g] != 0` is exactly that.
-                let s = family.mask.siblings() as usize;
+                let s = mask.siblings() as usize;
                 stats.counts[s] += 1;
-                for (g, &m) in family.genome_mask.iter().enumerate() {
+                for (g, &m) in genome_mask.iter().enumerate() {
                     if m != 0 {
                         stats.per_genome[g][s] += 1;
                     }
                 }
-            }
+            })?;
             pb.inc(1);
         }
         pb.finish_and_clear();