refactor: replace eager family collection with callback processing

Refactor `scan_layer_families` across the siblings module to accept a closure callback instead of returning an intermediate collection. This eliminates eager materialization and per-layer buffering by streaming results directly into genome-specific buffers or tally matrices. The update introduces bounded batch processing and scratch buffer reuse to cap peak auxiliary memory, while preserving existing computational behavior, control flow, and error semantics.

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,53 +112,57 @@ 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 {
+            let s = slot;
+            slot += 1;
+            let Some(mask) = annex.get(s) else { continue };
+            let Some(kmer) = slot_kmer[s] else { continue };
             if !is_minorant(kmer, mask, k) {
                 continue; // family tallied once, at its minorant
             }
-        let family_idx = families.len();
+            batch_slots.push(s);
-        families.push(Family { slot, mask });
-
-        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);
-            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 only: seed each family's own minorant presence (already
+        // open, no lookup) and collect this batch's cross-partition
+        // queries, grouped by destination partition.
+        let mut outgoing: Vec<Vec<(CanonicalKmer, usize, u8)>> = (0..n_parts).map(|_| Vec::new()).collect();
+        let genome_mask: Vec<AtomicU8> = (0..batch_slots.len() * n_genomes).map(|_| AtomicU8::new(0)).collect();
 
-    // Local presence (the family's own minorant slot) — already open, no
+        for (family_idx, &fslot) in batch_slots.iter().enumerate() {
-    // lookup, so no locality concern either way.
+            let kmer = slot_kmer[fslot].expect("batch slot has a kmer");
-    for (family_idx, family) in families.iter().enumerate() {
+            let mask = annex.get(fslot).expect("batch slot has a mask");
-        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) {
+                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));
+            }
         }
 
-    // ── Pass 2: resolve each destination partition's whole batch in one
+        // ── Resolve this batch's queries, one contiguous sweep per
-    // contiguous sweep — same rationale as `build_sibling_annex`'s
+        // destination partition — same rationale as `build_sibling_annex`'s
         // `outgoing` grouping.
         outgoing.par_iter().enumerate().filter(|(_, q)| !q.is_empty()).for_each(|(dest, queries)| {
             for &(variant, family_idx, base) in queries {
@@ -165,14 +175,14 @@ pub(super) fn scan_layer_families(
             }
         });
 
-    Ok(families
+        for (family_idx, &fslot) in batch_slots.iter().enumerate() {
-        .into_iter()
+            let mask = annex.get(fslot).expect("batch slot has a mask");
-        .enumerate()
+            for (g, dst) in scratch.iter_mut().enumerate() {
-        .map(|(family_idx, family)| FamilyRow {
+                *dst = genome_mask[family_idx * n_genomes + g].load(Ordering::Relaxed);
-            mask: family.mask,
+            }
-            genome_mask: (0..n_genomes)
+            on_family(mask, &scratch);
-                .map(|g| genome_mask[family_idx * n_genomes + g].load(Ordering::Relaxed))
+        }
-                .collect(),
+    }
-        })
+
-        .collect())
+    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();