Introduce concurrency-bounded parallel processing for index cache layers

Implement an `obipipeline::Throttle` with an RAII guard to acquire and release concurrency slots. Expose new bounded parallel methods on `IndexCache` to process cached layers with a configurable cap. Refactor downstream aggregation logic to use single-pass parallel map-reduce instead of manual collect-map-reduce sequences, enforcing a maximum of 8 concurrent layer scans to bound memory usage.

src/Cargo.lock

@@ -1567,6 +1567,7 @@ dependencies = [
  "obicompactvec",
  "obikindex",
  "obikseq",
+ "obipipeline",
  "rayon",
  "tempfile",
 ]

src/obikidxcache/Cargo.toml

@@ -7,6 +7,7 @@ edition = "2024"
 obikindex     = { path = "../obikindex"}
 obikseq     = { path = "../obikseq"}
 obicompactvec = { path = "../obicompactvec" }
+obipipeline   = { path = "../obipipeline" }
 ndarray       = "0.17"
 rayon         = "1"
 

src/obikidxcache/src/index_cache.rs

@@ -3,6 +3,8 @@ use std::sync::Arc;
 
 use obikindex::{KmerIndex, layer::KmerLayer};
 use obikseq::CanonicalKmer;
+use obipipeline::Throttle;
+use rayon::prelude::*;
 
 use crate::meta_cache::MetaCache;
 
@@ -188,4 +190,86 @@ impl IndexCache {
         let partition = kmer.partition(self.raw_index.n_partitions());
         self.find_in_partition(partition, kmer)
     }
+
+    /// Parallel map-reduce over every cached `KmerLayer` (across every
+    /// cached partition), bounded to `max_concurrent` layers scanned at
+    /// once — a cap on how many layers are actively being mapped
+    /// simultaneously, independent of `rayon`'s own thread count (a large
+    /// per-layer working set, e.g. a whole matrix scan, can make "one
+    /// layer per thread" too much resident memory at once even when the
+    /// thread pool itself is much bigger).
+    ///
+    /// Safe only when `map` is independent per layer (no cross-layer or
+    /// cross-partition state) — when a caller's own per-layer work already
+    /// parallelizes internally at a finer grain (e.g. across partitions,
+    /// like `obikphylo`'s sibling-annex traversal), stacking this on top
+    /// fights that inner parallelism for cache/mmap locality instead of
+    /// helping; use the plain sequential [`iter`](Self::iter) there
+    /// instead.
+    pub fn par_map_reduce<T, M, R>(
+        &self,
+        max_concurrent: usize,
+        identity: impl Fn() -> T + Sync,
+        map: M,
+        reduce: R,
+    ) -> T
+    where
+        T: Send,
+        M: Fn(&KmerLayer) -> T + Sync,
+        R: Fn(T, T) -> T + Sync,
+    {
+        let throttle = Arc::new(Throttle::new(max_concurrent));
+        self.layer_cache
+            .par_iter()
+            .flat_map(|(_, layers)| layers.par_iter())
+            .map(|layer| {
+                let _guard = throttle.acquire_guard();
+                map(layer)
+            })
+            .reduce(&identity, &reduce)
+    }
+
+    /// Parallel, side-effecting foreach over every cached `KmerLayer`,
+    /// bounded to `max_concurrent` layers at once — same bounding
+    /// reasoning and same "independent per layer" safety requirement as
+    /// [`par_map_reduce`](Self::par_map_reduce), just without a value to
+    /// combine back.
+    pub fn par_for_each<F>(&self, max_concurrent: usize, f: F)
+    where
+        F: Fn(&KmerLayer) + Sync,
+    {
+        let throttle = Arc::new(Throttle::new(max_concurrent));
+        self.layer_cache
+            .par_iter()
+            .flat_map(|(_, layers)| layers.par_iter())
+            .for_each(|layer| {
+                let _guard = throttle.acquire_guard();
+                f(layer);
+            });
+    }
+
+    /// Bounded parallel iterator over every cached `KmerLayer` — for
+    /// callers that want to chain their own `rayon` combinators rather
+    /// than fit their computation into [`par_map_reduce`](Self::par_map_reduce)/
+    /// [`par_for_each`](Self::par_for_each)'s shape.
+    ///
+    /// Each item is paired with a [`obipipeline::ThrottleGuard`] (same
+    /// `Throttled<T>` carrier `obipipeline::throttle` uses for a plain
+    /// sequential source) that releases its slot on drop — keep the whole
+    /// `Throttled` alive for as long as you're "using" that layer (e.g.
+    /// through your own `.map()`/`.for_each()`), the same discipline
+    /// `obipipeline`'s own consumers already follow.
+    pub fn par_iter_bounded(
+        &self,
+        max_concurrent: usize,
+    ) -> impl ParallelIterator<Item = obipipeline::Throttled<&KmerLayer>> {
+        let throttle = Arc::new(Throttle::new(max_concurrent));
+        self.layer_cache
+            .par_iter()
+            .flat_map(|(_, layers)| layers.par_iter())
+            .map(move |layer| obipipeline::Throttled {
+                item: layer,
+                guard: throttle.acquire_guard(),
+            })
+    }
 }

src/obikphylo/src/distance.rs

@@ -2,16 +2,17 @@
 //! (summing per-layer partials into one global partial) and the final
 //! distance calculation are both phylogenetic domain logic, so both live
 //! here, not in `obikidxcache` (which only ever hands out already-opened
-//! `KmerLayer`s via `IndexCache::iter()` — plain exposition, no
-//! computation of its own).
+//! `KmerLayer`s, via plain exposition — `IndexCache::iter()` — or bounded
+//! parallel map-reduce over them — `IndexCache::par_map_reduce()` — neither
+//! of which computes anything phylo-specific itself).
 //!
 //! [`AggregatedCount`]/[`AggregatedBit`] exist purely to "hang"
-//! `obicompactvec`'s `CountPartials`/`BitPartials` traits on a borrowed
-//! collection of `KmerLayer`s, so their already-written finalisation
-//! methods (`bray_dist_matrix`, `jaccard_dist_matrix`, ...) can be reused
-//! as-is instead of re-derived here — see `KmerLayer`'s own
-//! `ColumnWeights`/`CountPartials`/`BitPartials` impls (`obikindex`) for
-//! the per-layer primitives these sum.
+//! `obicompactvec`'s `CountPartials`/`BitPartials` traits on a bounded
+//! parallel reduction over `IndexCache`'s layers, so their already-written
+//! finalisation methods (`bray_dist_matrix`, `jaccard_dist_matrix`, ...) can
+//! be reused as-is instead of re-derived here — see `KmerLayer`'s own
+//! `ColumnWeights`/`CountPartials`/`BitPartials` impls (`obikindex`) for the
+//! per-layer primitives these sum.
 
 use ndarray::{Array1, Array2};
 use obicompactvec::{BitPartials, ColumnWeights, CountPartials};
@@ -64,77 +65,118 @@ impl DistanceMetric {
     }
 }
 
-// ── Aggregators — sum every cached layer's own partials into one global
-// partial, then hand the result off to `CountPartials`/`BitPartials`'s own
-// provided finalisation methods.
+// ── Aggregators — bounded parallel reduction of every cached layer's own
+// partial (of the wanted content kind) into one global partial, then handed
+// off to `CountPartials`/`BitPartials`'s own provided finalisation methods.
+// Layers of the *other* content kind contribute their zero element, same
+// as a layer that's simply absent — cheaper than pre-collecting a filtered
+// `Vec` first, since `IndexCache::par_map_reduce` never allocates one.
 
-fn sum_array1(n: usize, items: impl Iterator<Item = Array1<u64>>) -> Array1<u64> {
-    items.fold(Array1::zeros(n), |acc, x| acc + x)
-}
-fn sum_array2_u64(n: usize, items: impl Iterator<Item = Array2<u64>>) -> Array2<u64> {
-    items.fold(Array2::zeros((n, n)), |acc, x| acc + x)
-}
-fn sum_array2_f64(n: usize, items: impl Iterator<Item = Array2<f64>>) -> Array2<f64> {
-    items.fold(Array2::zeros((n, n)), |acc, x| acc + x)
-}
+/// Layers scanned at once per `par_map_reduce` call — bounds resident
+/// memory (each layer's own partial-computation working set), independent
+/// of how many `rayon` threads exist. See `IndexCache::par_map_reduce`'s
+/// own docs for why this is a separate knob from thread count.
+const MAX_CONCURRENT_LAYERS: usize = 8;
 
 struct AggregatedCount<'a> {
+    cache: &'a IndexCache,
     n_genomes: usize,
-    layers: Vec<&'a KmerLayer>,
+}
+
+impl AggregatedCount<'_> {
+    fn reduce<T: Send>(
+        &self,
+        identity: impl Fn() -> T + Sync,
+        per_layer: impl Fn(&KmerLayer) -> T + Sync,
+        combine: impl Fn(T, T) -> T + Sync,
+    ) -> T {
+        self.cache.par_map_reduce(
+            MAX_CONCURRENT_LAYERS,
+            &identity,
+            |layer| if layer.content() == LayerContent::Count { per_layer(layer) } else { identity() },
+            &combine,
+        )
+    }
 }
 
 impl ColumnWeights for AggregatedCount<'_> {
     fn col_weights(&self) -> Array1<u64> {
-        sum_array1(self.n_genomes, self.layers.iter().map(|l| l.col_weights()))
+        let n = self.n_genomes;
+        self.reduce(|| Array1::zeros(n), |l| l.col_weights(), |a, b| a + b)
     }
 }
 
 impl CountPartials for AggregatedCount<'_> {
     fn partial_bray(&self) -> Array2<u64> {
-        sum_array2_u64(self.n_genomes, self.layers.iter().map(|l| l.partial_bray()))
+        let n = self.n_genomes;
+        self.reduce(|| Array2::zeros((n, n)), |l| l.partial_bray(), |a, b| a + b)
     }
     fn partial_euclidean(&self) -> Array2<f64> {
-        sum_array2_f64(self.n_genomes, self.layers.iter().map(|l| l.partial_euclidean()))
+        let n = self.n_genomes;
+        self.reduce(|| Array2::zeros((n, n)), |l| l.partial_euclidean(), |a, b| a + b)
     }
     fn partial_threshold_jaccard(&self, threshold: u32) -> (Array2<u64>, Array2<u64>) {
         let n = self.n_genomes;
-        self.layers.iter().map(|l| l.partial_threshold_jaccard(threshold)).fold(
-            (Array2::zeros((n, n)), Array2::zeros((n, n))),
-            |(inter, union), (i, u)| (inter + i, union + u),
+        self.reduce(
+            || (Array2::zeros((n, n)), Array2::zeros((n, n))),
+            |l| l.partial_threshold_jaccard(threshold),
+            |(ia, ua), (ib, ub)| (ia + ib, ua + ub),
         )
     }
     fn partial_relfreq_bray(&self, global: &Array1<u64>) -> Array2<f64> {
-        sum_array2_f64(self.n_genomes, self.layers.iter().map(|l| l.partial_relfreq_bray(global)))
+        let n = self.n_genomes;
+        self.reduce(|| Array2::zeros((n, n)), |l| l.partial_relfreq_bray(global), |a, b| a + b)
     }
     fn partial_relfreq_euclidean(&self, global: &Array1<u64>) -> Array2<f64> {
-        sum_array2_f64(self.n_genomes, self.layers.iter().map(|l| l.partial_relfreq_euclidean(global)))
+        let n = self.n_genomes;
+        self.reduce(|| Array2::zeros((n, n)), |l| l.partial_relfreq_euclidean(global), |a, b| a + b)
     }
     fn partial_hellinger(&self, global: &Array1<u64>) -> Array2<f64> {
-        sum_array2_f64(self.n_genomes, self.layers.iter().map(|l| l.partial_hellinger(global)))
+        let n = self.n_genomes;
+        self.reduce(|| Array2::zeros((n, n)), |l| l.partial_hellinger(global), |a, b| a + b)
     }
 }
 
 struct AggregatedBit<'a> {
+    cache: &'a IndexCache,
     n_genomes: usize,
-    layers: Vec<&'a KmerLayer>,
+}
+
+impl AggregatedBit<'_> {
+    fn reduce<T: Send>(
+        &self,
+        identity: impl Fn() -> T + Sync,
+        per_layer: impl Fn(&KmerLayer) -> T + Sync,
+        combine: impl Fn(T, T) -> T + Sync,
+    ) -> T {
+        self.cache.par_map_reduce(
+            MAX_CONCURRENT_LAYERS,
+            &identity,
+            |layer| if layer.content() == LayerContent::Presence { per_layer(layer) } else { identity() },
+            &combine,
+        )
+    }
 }
 
 impl ColumnWeights for AggregatedBit<'_> {
     fn col_weights(&self) -> Array1<u64> {
-        sum_array1(self.n_genomes, self.layers.iter().map(|l| l.col_weights()))
+        let n = self.n_genomes;
+        self.reduce(|| Array1::zeros(n), |l| l.col_weights(), |a, b| a + b)
     }
 }
 
 impl BitPartials for AggregatedBit<'_> {
     fn partial_jaccard(&self) -> (Array2<u64>, Array2<u64>) {
         let n = self.n_genomes;
-        self.layers.iter().map(|l| l.partial_jaccard()).fold(
-            (Array2::zeros((n, n)), Array2::zeros((n, n))),
-            |(inter, union), (i, u)| (inter + i, union + u),
+        self.reduce(
+            || (Array2::zeros((n, n)), Array2::zeros((n, n))),
+            |l| l.partial_jaccard(),
+            |(ia, ua), (ib, ub)| (ia + ib, ua + ub),
         )
     }
     fn partial_hamming(&self) -> Array2<u64> {
-        sum_array2_u64(self.n_genomes, self.layers.iter().map(|l| l.partial_hamming()))
+        let n = self.n_genomes;
+        self.reduce(|| Array2::zeros((n, n)), |l| l.partial_hamming(), |a, b| a + b)
     }
 }
 
@@ -167,8 +209,7 @@ impl Metrics for IndexCache {
         let kmer_size = self.meta().config().kmer_size;
 
         if use_counts {
-            let layers: Vec<&KmerLayer> = self.iter().filter(|l| l.content() == LayerContent::Count).collect();
-            let agg = AggregatedCount { n_genomes, layers };
+            let agg = AggregatedCount { cache: self, n_genomes };
 
             let matrix = match metric {
                 DistanceMetric::BrayCurtis         => agg.bray_dist_matrix(),
@@ -195,8 +236,7 @@ impl Metrics for IndexCache {
 
             Ok(DistanceOutput { matrix, shared_kmers: shared })
         } else {
-            let layers: Vec<&KmerLayer> = self.iter().filter(|l| l.content() == LayerContent::Presence).collect();
-            let agg = AggregatedBit { n_genomes, layers };
+            let agg = AggregatedBit { cache: self, n_genomes };
 
             let matrix = match metric {
                 DistanceMetric::Jaccard => agg.jaccard_dist_matrix(),

src/obikstats/src/stats.rs

@@ -2,13 +2,16 @@ use std::fs;
 use std::path::Path;
 use std::sync::Arc;
 
-use rayon::prelude::*;
-
 use obikalgorithm::Algorithm;
 use obikidxcache::index_cache::IndexCache;
 use obikindex::KmerIndex;
 use obikindex::layer::KmerLayer;
 
+/// Layers scanned at once per `par_map_reduce` call — see
+/// `obikidxcache::IndexCache::par_map_reduce`'s own docs for why this is a
+/// separate knob from thread count.
+const MAX_CONCURRENT_LAYERS: usize = 8;
+
 /// Bits per kmer broken down by index component.
 pub struct IndexBitsPerKmer {
     /// Total distinct k-mers across all partitions and layers.
@@ -105,16 +108,16 @@ impl Algorithm for BitsPerKmer {
     fn run(&mut self) -> obikalgorithm::Result<IndexBitsPerKmer> {
         let n_genomes = self.index.meta().genomes()?.len().max(1);
         let cache = IndexCache::new(Arc::clone(&self.index), None);
-        let layers: Vec<&KmerLayer> = cache.iter().collect();
 
-        let (n_kmers, mphf_b, evidence_b, matrix_b) = layers
-            .par_iter()
-            .map(|layer| layer_bytes(layer))
-            .map(|lb| (lb.n_kmers, lb.mphf, lb.evidence, lb.matrix))
-            .reduce(
-                || (0usize, 0u64, 0u64, 0u64),
-                |a, b| (a.0 + b.0, a.1 + b.1, a.2 + b.2, a.3 + b.3),
-            );
+        let (n_kmers, mphf_b, evidence_b, matrix_b) = cache.par_map_reduce(
+            MAX_CONCURRENT_LAYERS,
+            || (0usize, 0u64, 0u64, 0u64),
+            |layer| {
+                let lb = layer_bytes(layer);
+                (lb.n_kmers, lb.mphf, lb.evidence, lb.matrix)
+            },
+            |a, b| (a.0 + b.0, a.1 + b.1, a.2 + b.2, a.3 + b.3),
+        );
 
         if n_kmers == 0 {
             return Ok(IndexBitsPerKmer {
@@ -166,19 +169,26 @@ impl Algorithm for GenomeKmerCounts {
     fn run(&mut self) -> obikalgorithm::Result<(usize, Vec<u64>)> {
         let n_genomes = self.index.meta().genomes()?.len();
         let cache = IndexCache::new(Arc::clone(&self.index), None);
-        let layers: Vec<&KmerLayer> = cache.iter().collect();
 
-        let total_kmers: usize = layers.iter().map(|l| l.n()).sum();
-
-        let mut total_counts = vec![0u64; n_genomes];
-        let per_layer_weights: Vec<_> = layers.par_iter().map(|l| l.col_weights()).collect();
-        for weights in per_layer_weights {
-            for (g, &v) in weights.iter().enumerate() {
-                if g < n_genomes {
-                    total_counts[g] += v;
+        let (total_kmers, total_counts) = cache.par_map_reduce(
+            MAX_CONCURRENT_LAYERS,
+            || (0usize, vec![0u64; n_genomes]),
+            |layer| {
+                let mut counts = vec![0u64; n_genomes];
+                for (g, &w) in layer.col_weights().iter().enumerate() {
+                    if g < n_genomes {
+                        counts[g] = w;
+                    }
                 }
-            }
-        }
+                (layer.n(), counts)
+            },
+            |(ka, mut ca), (kb, cb)| {
+                for (x, y) in ca.iter_mut().zip(cb.iter()) {
+                    *x += y;
+                }
+                (ka + kb, ca)
+            },
+        );
 
         Ok((total_kmers, total_counts))
     }

src/obipipeline/src/throttle.rs

@@ -35,6 +35,19 @@ impl Throttle {
     }
 }
 
+impl Throttle {
+    /// Acquire one slot and return a guard that releases it on drop —
+    /// for callers with no natural "source iterator" to wrap via
+    /// [`throttle`] (e.g. bounding concurrency inside a `rayon` parallel
+    /// iterator's own closure, one call per item, rather than up front in
+    /// a single source thread). Blocks the calling thread until a slot is
+    /// free, same as [`throttle`]'s own per-item acquisition.
+    pub fn acquire_guard(self: &Arc<Self>) -> ThrottleGuard {
+        self.acquire();
+        ThrottleGuard(Arc::clone(self))
+    }
+}
+
 // ── ThrottleGuard ─────────────────────────────────────────────────────────────
 
 /// RAII guard: releases one slot in the `Throttle` when dropped.