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
-        // variable family, appended in layer order for a single
-        // deterministic column order across the whole index.
-        let mut partials: Vec<Vec<Vec<u8>>> = Vec::with_capacity(layer_dirs.len());
+        // the grouping. Columns appended straight into `sequences` as each
+        // family comes back from `scan_layer_families` (bounded-batch, not
+        // whole-layer) — no intermediate per-layer buffer, layer order
+        // 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)?;
-            let columns = families
-                .into_iter()
-                .filter(|f| f.mask.family_size() >= 2) // monomorphic family — no signal, skip
-                .map(|f| f.genome_mask.iter().map(|&m| iupac_code(m)).collect())
-                .collect();
-            partials.push(columns);
+            scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, |mask, genome_mask| {
+                if mask.family_size() < 2 {
+                    return; // monomorphic family — no signal, skip
+                }
+                for (g, &m) in genome_mask.iter().enumerate() {
+                    sequences[g].push(iupac_code(m));
+                }
+            })?;
             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.
-        let mut partials: Vec<[[u64; 5]; 5]> = Vec::with_capacity(layer_dirs.len());
+        // partition-grouped locality. Tallied straight into `total` as each
+        // 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)?;
-            let mut counts = [[0u64; 5]; 5];
-
-            for family in &families {
-                if family.mask.family_size() < 2 {
+            scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, |mask, genome_mask| {
+                if 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 {
-                let variable = family.mask.family_size() >= 2;
-                let genome_mask = &family.genome_mask;
+            scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, |mask, genome_mask| {
+                let variable = mask.family_size() >= 2;
 
                 // 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
-//! 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
-//! cache (near-continuous ~2 GB/s pagein), the same failure mode
-//! `build_sibling_annex` had already been fixed for.
+//! between partitions in family order.
+//!
+//! Processed in bounded batches of [`FAMILY_BATCH`] families, not the whole
+//! layer at once — same rationale and constant as `build_layer_sibling_annex`'s
+//! `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,
-/// the 4-bit base-presence mask resolved across the whole index (bit `b`
-/// set iff that genome carries the family member whose own canonical
-/// central base is `b`).
-pub(super) struct FamilyRow {
-    pub mask: FamilyMask,
-    pub genome_mask: Vec<u8>,
-}
+/// Families per batch — bounds `scan_layer_families`'s extra memory to
+/// `FAMILY_BATCH * n_genomes` bytes regardless of layer size. Same value and
+/// rationale as `build_layer_sibling_annex`'s `BATCH_SIZE`: large enough to
+/// amortise the per-batch partition-grouping/sync cost, small enough to keep
+/// working memory a rounding error next to the index itself.
+const FAMILY_BATCH: usize = 4096;
 
 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
-/// base-presence mask resolved against the whole (already-open) partition
-/// cache.
-///
-/// Two passes, not one inline pass: the first is purely local (already-open
-/// annex/mphf/matrix, no lookups) and both enumerates the layer's families
-/// and collects every cross-partition query they need, grouped by
-/// 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.
+/// Visits every minorant family of one layer, in slot order, batched
+/// [`FAMILY_BATCH`] at a time — see the module docs for why. `on_family` is
+/// 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
+/// member whose own canonical central base is `b`), backed by a scratch
+/// buffer reused across every call — callers that need to keep data past
+/// the call must copy it themselves.
 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
-    // cross-partition query they need, grouped by destination partition.
-    struct Family {
-        slot: usize,
-        mask: FamilyMask,
-    }
-    let mut families: Vec<Family> = Vec::new();
-    let mut outgoing: Vec<Vec<(CanonicalKmer, usize, u8)>> = (0..n_parts).map(|_| Vec::new()).collect();
+    let n_slots = annex.len();
+    let mut slot = 0usize;
+    let mut batch_slots: Vec<usize> = Vec::with_capacity(FAMILY_BATCH);
+    let mut scratch = vec![0u8; n_genomes]; // reused across every `on_family` call — no per-family allocation
 
-    for slot in 0..annex.len() {
-        let Some(mask) = annex.get(slot) else { continue };
-        let Some(kmer) = slot_kmer[slot] else { continue };
+    while slot < n_slots {
+        batch_slots.clear();
+        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();
-        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));
+            batch_slots.push(s);
         }
+        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
-    // 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");
+        for (family_idx, &fslot) in batch_slots.iter().enumerate() {
+            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, 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
-    // contiguous sweep — same rationale as `build_sibling_annex`'s
+        // ── Resolve this batch's queries, one contiguous sweep per
+        // 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
-        .into_iter()
-        .enumerate()
-        .map(|(family_idx, family)| FamilyRow {
-            mask: family.mask,
-            genome_mask: (0..n_genomes)
-                .map(|g| genome_mask[family_idx * n_genomes + g].load(Ordering::Relaxed))
-                .collect(),
-        })
-        .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)?;
-            for family in &families {
+            scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, |mask, genome_mask| {
                 // "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();