Replace slot-based indexing with iteration order and stream k-mers

Transitions the index from MPHF slot-based to physical iteration-order indexing, aligning with the unitig layout. Introduces a streaming-only pipeline for k-mer iteration that adheres to memory constraints by avoiding full in-memory collections. Updates layer and sibling iterators to own an Arc clone of the file reader, making them Send + 'static and safe for concurrent use without borrowing the parent. Exposes batch and k-mer iterator types publicly while simplifying signature syntax with modern lifetime elision.
This commit is contained in:
Eric Coissac
2026-08-16 14:07:22 +02:00
parent b1f54b7d2f
commit 519195d4a1
7 changed files with 219 additions and 84 deletions
+75 -56
View File
@@ -21,13 +21,17 @@ use super::{olm_to_ok, ANNEX_FILE_NAME, INDEX_SUBDIR};
// ── obipipeline data types ───────────────────────────────────────────────── // ── obipipeline data types ─────────────────────────────────────────────────
/// A batch of this layer's distinct k-mers (local MPHF slot + k-mer), the /// A batch of this layer's distinct k-mers (iteration-order index + k-mer),
/// pipeline's source item — batched, not one k-mer per item, so that /// the pipeline's source item — batched, not one k-mer per item, so that
/// pipeline messages and their synchronisation cost stay amortised over /// pipeline messages and their synchronisation cost stay amortised over
/// thousands of lookups (see `build_layer_sibling_annex`'s comment on /// thousands of lookups (see `build_layer_sibling_annex`'s comment on
/// `BATCH_SIZE`). Carries the throttle permit for the whole batch, moved /// `BATCH_SIZE`). Carries the throttle permit for the whole batch, moved
/// (not cloned) into the corresponding `VariantBatch` — a 1-to-1 transform, /// (not cloned) into the corresponding `VariantBatch` — a 1-to-1 transform,
/// unlike the fan-out `Flat` stage this replaced, needs no `Arc` sharing. /// unlike the fan-out `Flat` stage this replaced, needs no `Arc` sharing.
///
/// The `usize` is this k-mer's position in `iter_kmers()`'s enumeration
/// order, **not** an MPHF slot — see `docmd/architecture/siblings.md`. It
/// is the same index the annex file is written under.
struct SourceBatch { struct SourceBatch {
items: Vec<(usize, CanonicalKmer)>, items: Vec<(usize, CanonicalKmer)>,
_permit: ThrottleGuard, _permit: ThrottleGuard,
@@ -36,7 +40,9 @@ struct SourceBatch {
/// One batch's worth of central-substitution variants (up to 3 per source /// One batch's worth of central-substitution variants (up to 3 per source
/// k-mer), each already routed to its destination partition and carrying /// k-mer), each already routed to its destination partition and carrying
/// its own central base (0=A/1=C/2=G/3=T) — the mask bit it will set on a /// its own central base (0=A/1=C/2=G/3=T) — the mask bit it will set on a
/// hit. `(dest_partition, variant, source_slot, base)` per entry. /// hit. `(dest_partition, variant, source_order, base)` per entry, where
/// `source_order` is the source k-mer's iteration-order index (see
/// `SourceBatch`), not an MPHF slot.
struct VariantBatch { struct VariantBatch {
items: Vec<(usize, CanonicalKmer, usize, u8)>, items: Vec<(usize, CanonicalKmer, usize, u8)>,
_permit: ThrottleGuard, _permit: ThrottleGuard,
@@ -103,7 +109,7 @@ impl KmerIndex {
Ok(()) Ok(())
} }
/// Returns the number of distinct k-mers (annex slots) processed, for /// Returns the number of distinct k-mers (annex entries) processed, for
/// progress reporting. /// progress reporting.
fn build_layer_sibling_annex( fn build_layer_sibling_annex(
&self, &self,
@@ -114,29 +120,31 @@ impl KmerIndex {
let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir"); 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 meta = PartitionMeta::load(index_dir).map_err(olm_to_ok)?;
let mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?; let mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?;
let n_slots = mphf.n();
// ── Enumerate this layer's distinct k-mers, one per slot — direct
// slot -> k-mer reconstruction (evidence + direct-access unitigs, no
// MPHF hashing, no file scan), not scan-and-hash-every-k-mer-forward.
let slot_kmer: Vec<Option<CanonicalKmer>> = (0..n_slots).map(|slot| mphf.kmer_at(slot)).collect();
let k = self.kmer_size(); let k = self.kmer_size();
let n = mphf.n();
// ── Reconciliation state, initialised with each slot's own base — // ── Reconciliation state, indexed by this layer's k-mer iteration
// that member is trivially present, no lookup needed. Built before // order (the physical layout of `unitigs.bin`), never by MPHF slot
// the pipeline runs, from the same enumeration, since `sources` // — this layer's own k-mers are known members by construction, so
// below is consumed as a throttled iterator, not collected. // no evidence check, no MPHF slot, and no slot -> k-mer
// `AtomicU8`, not `FamilyMask`, because the gather phase below // reconstruction is needed or legitimate here (see
// parallelises across destination partitions (independent // `docmd/architecture/siblings.md`: the MPHF is not invertible, and
// `query_partition_with` calls, safe to run concurrently) and their // evidence answers membership, not identity). The annex is written
// `Found` hits can land on arbitrary, possibly-shared slots — a // under this same iteration order end to end, so a reader can later
// lock-free `fetch_or` avoids needing any synchronisation beyond // zip `iter_kmers()` with the annex file directly, with no
// that. ───────────────────────────────────────────────────────── // MPHF/slot indirection at read time either.
let mask: Vec<AtomicU8> = (0..n_slots).map(|_| AtomicU8::new(0)).collect(); //
for (slot, kmer) in slot_kmer.iter().enumerate().filter_map(|(s, k)| k.map(|k| (s, k))) { // `Arc<Vec<AtomicU8>>`, not `Vec<AtomicU8>` — `Pipe::apply` requires
mask[slot].fetch_or(1 << central_base(kmer, k), Ordering::Relaxed); // its source iterator to be `Send + 'static` (its items are
} // dispatched to worker threads that outlive this call), so the
// batch-generating closure below needs an owned handle it can move
// in, not a borrow of a local. `AtomicU8`, not `FamilyMask`,
// because the gather phase below parallelises across destination
// partitions (independent `query_partition_with` calls, safe to run
// concurrently) and their `Found` hits can land on arbitrary,
// possibly-shared source entries — a lock-free `fetch_or` avoids
// needing any synchronisation beyond that. ─────────────────────────
let mask: Arc<Vec<AtomicU8>> = Arc::new((0..n).map(|_| AtomicU8::new(0)).collect());
// ── obipipeline: a *batch* transform, not a per-k-mer `Flat` one — // ── obipipeline: a *batch* transform, not a per-k-mer `Flat` one —
// the actual cross-partition lookup reuses // the actual cross-partition lookup reuses
@@ -169,16 +177,25 @@ impl KmerIndex {
// the pipeline by the accumulation loop below. See // the pipeline by the accumulation loop below. See
// `obipipeline::throttle`'s docs for why this is required, not // `obipipeline::throttle`'s docs for why this is required, not
// optional, once a `Flat`-style stage sits in the pipeline. // optional, once a `Flat`-style stage sits in the pipeline.
let sources: Vec<(usize, CanonicalKmer)> = slot_kmer //
.iter() // Batches stream straight from `unitigs.bin` via
// `enumerate_kmers_batch` — never a full-layer `Vec` collect (a
// layer can hold billions of k-mers; see the "no full collect"
// rule). Each k-mer's own base is seeded into `mask` in this same
// pass, since this is exactly the iteration order `mask` is keyed
// on — no separate seeding pass needed.
let seed_mask = Arc::clone(&mask);
let batches = mphf.enumerate_kmers_batch(BATCH_SIZE).map(move |(start, kmers)| {
kmers
.into_iter()
.enumerate() .enumerate()
.filter_map(|(slot, maybe_kmer)| maybe_kmer.map(|kmer| (slot, kmer))) .map(|(i, kmer)| {
.collect(); seed_mask[start + i].fetch_or(1 << central_base(kmer, k), Ordering::Relaxed);
let batches: Vec<Vec<(usize, CanonicalKmer)>> = sources (start + i, kmer)
.chunks(BATCH_SIZE) })
.map(|chunk| chunk.to_vec()) .collect::<Vec<_>>()
.collect(); });
let throttled = obipipeline::throttle(batches.into_iter(), n_workers).map(|t| SourceBatch { let throttled = obipipeline::throttle(batches, n_workers).map(|t| SourceBatch {
items: t.item, items: t.item,
_permit: t.guard, _permit: t.guard,
}); });
@@ -188,7 +205,7 @@ impl KmerIndex {
| { | {
move |batch: SourceBatch| -> VariantBatch { move |batch: SourceBatch| -> VariantBatch {
let mut items = Vec::with_capacity(batch.items.len() * 3); let mut items = Vec::with_capacity(batch.items.len() * 3);
for (slot, kmer) in batch.items { for (order, kmer) in batch.items {
for variant in kmer.central_canonical_neighbors() { for variant in kmer.central_canonical_neighbors() {
if variant == kmer { if variant == kmer {
continue; continue;
@@ -196,7 +213,7 @@ impl KmerIndex {
items.push(( items.push((
partition_of(variant, n_parts), partition_of(variant, n_parts),
variant, variant,
slot, order,
central_base(variant, k), central_base(variant, k),
)); ));
} }
@@ -216,8 +233,8 @@ impl KmerIndex {
// them. Each batch's throttle permit drops here, once accumulated. // them. Each batch's throttle permit drops here, once accumulated.
let mut outgoing: Vec<Vec<(CanonicalKmer, usize, u8)>> = (0..n_parts).map(|_| Vec::new()).collect(); let mut outgoing: Vec<Vec<(CanonicalKmer, usize, u8)>> = (0..n_parts).map(|_| Vec::new()).collect();
for vb in pipe.apply(throttled, n_workers, capacity) { for vb in pipe.apply(throttled, n_workers, capacity) {
for (dest_partition, variant, source_slot, base) in vb.items { for (dest_partition, variant, source_order, base) in vb.items {
outgoing[dest_partition].push((variant, source_slot, base)); outgoing[dest_partition].push((variant, source_order, base));
} }
} }
@@ -226,33 +243,35 @@ impl KmerIndex {
// read-only) so this keeps using multiple cores without giving up // read-only) so this keeps using multiple cores without giving up
// the per-partition locality above. ───────────────────────────── // the per-partition locality above. ─────────────────────────────
outgoing.par_iter().enumerate().filter(|(_, q)| !q.is_empty()).for_each(|(dest, queries)| { outgoing.par_iter().enumerate().filter(|(_, q)| !q.is_empty()).for_each(|(dest, queries)| {
for &(variant, source_slot, base) in queries { for &(variant, source_order, base) in queries {
if cache.find(dest, variant) { if cache.find(dest, variant) {
mask[source_slot].fetch_or(1 << base, Ordering::Relaxed); mask[source_order].fetch_or(1 << base, Ordering::Relaxed);
} }
} }
}); });
// ── Write the layer's annex file ───────────────────────────────────── // ── Write the layer's annex file, indexed by iteration order — a
// The minorant flag is computed here, not by every later reader: // second streamed pass over `unitigs.bin` (via `enumerate_kmers`),
// this is the one place the whole family's *final* mask and this // now that every entry's mask is final; never a `Vec` hold of the
// slot's own k-mer (already in hand, no extra lookup) are both // whole layer. The minorant flag is computed here, not by every
// available together. Every consumer that only needs to know // later reader: this is the one place the whole family's *final*
// "is this slot the family's minorant" — the common case, since a // mask and this entry's own k-mer (already in hand, no extra
// family is tallied once, at its minorant — reads the bit straight // lookup) are both available together. Every consumer that only
// back instead of re-deriving it (re-scanning `unitigs.bin` and // needs to know "is this k-mer the family's minorant" — the common
// re-hashing through the MPHF, the cost `is_minorant` was cheap to // case, since a family is tallied once, at its minorant — reads
// *compute* but expensive to *get the inputs for* every time). // the bit straight back instead of re-deriving it (re-scanning
// `unitigs.bin` and re-hashing through the MPHF, the cost
// `is_minorant` was cheap to *compute* but expensive to *get the
// inputs for* every time).
let annex_path = layer_dir.join(ANNEX_FILE_NAME); let annex_path = layer_dir.join(ANNEX_FILE_NAME);
let mut builder = SiblingAnnexBuilder::new(n_slots, &annex_path)?; let mut builder = SiblingAnnexBuilder::new(n, &annex_path)?;
for (slot, m) in mask.iter().enumerate() { for (order, kmer) in mphf.enumerate_kmers() {
let Some(kmer) = slot_kmer[slot] else { continue }; // unused MPHF slot, if any — leave at the sentinel let final_mask = FamilyMask::from_bits(mask[order].load(Ordering::Relaxed));
let final_mask = FamilyMask::from_bits(m.load(Ordering::Relaxed));
let minorant = is_minorant(kmer, final_mask, k); let minorant = is_minorant(kmer, final_mask, k);
builder.set(slot, final_mask.with_minorant(minorant)); builder.set(order, final_mask.with_minorant(minorant));
} }
builder.close()?; builder.close()?;
Ok(n_slots as u64) Ok(n as u64)
} }
} }
+78
View File
@@ -0,0 +1,78 @@
//! Phylo/sibling-domain iteration over a layer — an extension trait, not a
//! new field on `MphfLayer`/`Layer<D>`: "family"/"minorant" are phylo
//! concepts, `obilayeredmap` stays kmer/slot-mapping only (see
//! `docmd/architecture/siblings.md`).
//!
//! The sibling annex is persisted in the same order as `iter_kmers()`
//! (`build_sibling_annex`, see `build.rs`), so pairing them is a plain zip —
//! no MPHF, no slot, no `kmer_at`. Both sides are already `Send + 'static`
//! (`KmerIter` owns an `Arc<UnitigFileReader>` clone; `SiblingAnnex` is
//! mmap-backed and handed in as an `Arc` by the caller), so `SiblingIter`
//! streams straight from disk and can be fed to `obipipeline` batch by
//! batch — never collected whole into memory (see the project's "no full
//! collect" rule).
use std::sync::Arc;
use obicompactvec::{FamilyMask, SiblingAnnex};
use obikseq::CanonicalKmer;
use obilayeredmap::{KmerIter, MphfLayer};
/// One layer entry: a k-mer's position in the layer's iteration order (the
/// same index the sibling annex is keyed on — not an MPHF slot), the k-mer
/// itself, and its family mask.
#[derive(Debug, Clone, Copy)]
pub struct SiblingEntry {
pub order: usize,
pub kmer: CanonicalKmer,
pub mask: FamilyMask,
}
/// Streams `(order, kmer, mask)` triples for one layer, in iteration order.
/// Produced by [`SiblingLayerExt::iter_siblings`].
pub struct SiblingIter {
kmers: KmerIter,
annex: Arc<SiblingAnnex>,
order: usize,
}
impl Iterator for SiblingIter {
type Item = SiblingEntry;
fn next(&mut self) -> Option<Self::Item> {
loop {
let kmer = self.kmers.next()?;
let order = self.order;
self.order += 1;
// `None` means "not yet computed" (see `SiblingAnnex` module
// docs) — shouldn't happen against a fully-built annex, but
// skip rather than misalign the two streams if it does.
if let Some(mask) = self.annex.get(order) {
return Some(SiblingEntry { order, kmer, mask });
}
}
}
}
/// Adds phylo/sibling iteration to `MphfLayer`.
pub trait SiblingLayerExt {
/// Zip this layer's k-mers with their sibling-annex entry, in iteration
/// order. `annex` must have been built from this same layer (its length
/// must match the layer's k-mer count).
fn iter_siblings(&self, annex: Arc<SiblingAnnex>) -> SiblingIter;
/// Like [`iter_siblings`](Self::iter_siblings), filtered to the
/// minorant of each family — the common case, since a family is
/// tallied once, at its minorant.
fn iter_minorants(&self, annex: Arc<SiblingAnnex>) -> impl Iterator<Item = SiblingEntry>;
}
impl SiblingLayerExt for MphfLayer {
fn iter_siblings(&self, annex: Arc<SiblingAnnex>) -> SiblingIter {
SiblingIter { kmers: self.iter_kmers(), annex, order: 0 }
}
fn iter_minorants(&self, annex: Arc<SiblingAnnex>) -> impl Iterator<Item = SiblingEntry> {
self.iter_siblings(annex).filter(|e| e.mask.is_minorant())
}
}
+2
View File
@@ -49,6 +49,7 @@ mod cardinality;
mod distance; mod distance;
mod family_scan; mod family_scan;
mod helpers; mod helpers;
mod iter;
mod stats; mod stats;
#[cfg(test)] #[cfg(test)]
@@ -57,6 +58,7 @@ mod tests;
pub use alignment::SnpAlignment; pub use alignment::SnpAlignment;
pub use cardinality::CardinalityTally; pub use cardinality::CardinalityTally;
pub use distance::{BasePairTally, RawSnpDistanceOutput}; pub use distance::{BasePairTally, RawSnpDistanceOutput};
pub use iter::{SiblingEntry, SiblingIter, SiblingLayerExt};
pub use stats::SiblingAnnexStats; pub use stats::SiblingAnnexStats;
use obilayeredmap::OLMError; use obilayeredmap::OLMError;
+4 -4
View File
@@ -85,24 +85,24 @@ impl<D: LayerData> Layer<D> {
} }
/// Iterate over all canonical kmers in the layer, in deterministic order. /// Iterate over all canonical kmers in the layer, in deterministic order.
pub fn iter_kmers(&self) -> crate::mphf_layer::KmerIter<'_> { pub fn iter_kmers(&self) -> crate::mphf_layer::KmerIter {
self.mphf.iter_kmers() self.mphf.iter_kmers()
} }
/// Iterate over all canonical kmers, each paired with its zero-based /// Iterate over all canonical kmers, each paired with its zero-based
/// sequence index in `unitigs.bin`. /// sequence index in `unitigs.bin`.
pub fn enumerate_kmers(&self) -> std::iter::Enumerate<crate::mphf_layer::KmerIter<'_>> { pub fn enumerate_kmers(&self) -> std::iter::Enumerate<crate::mphf_layer::KmerIter> {
self.mphf.enumerate_kmers() self.mphf.enumerate_kmers()
} }
/// Iterate over the layer's canonical kmers in batches of `n`. /// Iterate over the layer's canonical kmers in batches of `n`.
pub fn iter_kmers_batch(&self, n: usize) -> crate::mphf_layer::KmerBatchIter<'_> { pub fn iter_kmers_batch(&self, n: usize) -> crate::mphf_layer::KmerBatchIter {
self.mphf.iter_kmers_batch(n) self.mphf.iter_kmers_batch(n)
} }
/// Iterate over batches, each paired with the zero-based index of the /// Iterate over batches, each paired with the zero-based index of the
/// first kmer in the batch. /// first kmer in the batch.
pub fn enumerate_kmers_batch(&self, n: usize) -> std::iter::Enumerate<crate::mphf_layer::KmerBatchIter<'_>> { pub fn enumerate_kmers_batch(&self, n: usize) -> std::iter::Enumerate<crate::mphf_layer::KmerBatchIter> {
self.mphf.enumerate_kmers_batch(n) self.mphf.enumerate_kmers_batch(n)
} }
+1 -1
View File
@@ -12,4 +12,4 @@ pub use layer::{Hit, Layer, LayerData};
pub use layered_store::LayeredStore; pub use layered_store::LayeredStore;
pub use map::LayeredMap; pub use map::LayeredMap;
pub use meta::{IndexMode, PartitionMeta}; pub use meta::{IndexMode, PartitionMeta};
pub use mphf_layer::{MphfLayer, MphfOnly}; pub use mphf_layer::{KmerBatchIter, KmerIter, MphfLayer, MphfOnly};
+29 -22
View File
@@ -1,6 +1,7 @@
use std::fs; use std::fs;
use std::iter::Enumerate; use std::iter::Enumerate;
use std::path::{Path, PathBuf}; use std::path::{Path, PathBuf};
use std::sync::Arc;
use cacheline_ef::{CachelineEf, CachelineEfVec}; use cacheline_ef::{CachelineEf, CachelineEfVec};
use epserde::prelude::*; use epserde::prelude::*;
@@ -29,9 +30,9 @@ type MphfEps = PtrHash<u64, CubicEps, CachelineEfVec<&'static [CachelineEf]>, Xx
// ── LayerEvidence ───────────────────────────────────────────────────────────── // ── LayerEvidence ─────────────────────────────────────────────────────────────
enum LayerEvidence { enum LayerEvidence {
Exact { evidence: Evidence, unitigs: UnitigFileReader }, Exact { evidence: Evidence, unitigs: Arc<UnitigFileReader> },
Approx { fingerprint: FingerprintVec, unitigs: UnitigFileReader, unitigs_path: PathBuf }, Approx { fingerprint: FingerprintVec, unitigs: Arc<UnitigFileReader>, unitigs_path: PathBuf },
Hybrid { evidence: Evidence, unitigs: UnitigFileReader, fingerprint: FingerprintVec }, Hybrid { evidence: Evidence, unitigs: Arc<UnitigFileReader>, fingerprint: FingerprintVec },
} }
// ── MphfLayer ───────────────────────────────────────────────────────────────── // ── MphfLayer ─────────────────────────────────────────────────────────────────
@@ -58,13 +59,13 @@ impl MphfLayer {
IndexMode::Exact => { IndexMode::Exact => {
let evidence = Evidence::open(&dir.join(EVIDENCE_FILE))?; let evidence = Evidence::open(&dir.join(EVIDENCE_FILE))?;
let n = evidence.len(); let n = evidence.len();
let unitigs = UnitigFileReader::open(&dir.join(UNITIGS_FILE))?; let unitigs = Arc::new(UnitigFileReader::open(&dir.join(UNITIGS_FILE))?);
(LayerEvidence::Exact { evidence, unitigs }, n) (LayerEvidence::Exact { evidence, unitigs }, n)
} }
IndexMode::Approx { .. } => { IndexMode::Approx { .. } => {
let fingerprint = FingerprintVec::open(&dir.join(FINGERPRINT_FILE))?; let fingerprint = FingerprintVec::open(&dir.join(FINGERPRINT_FILE))?;
let n = fingerprint.n(); let n = fingerprint.n();
let unitigs = UnitigFileReader::open(&dir.join(UNITIGS_FILE))?; let unitigs = Arc::new(UnitigFileReader::open(&dir.join(UNITIGS_FILE))?);
let unitigs_path = dir.join(UNITIGS_FILE); let unitigs_path = dir.join(UNITIGS_FILE);
(LayerEvidence::Approx { fingerprint, unitigs, unitigs_path }, n) (LayerEvidence::Approx { fingerprint, unitigs, unitigs_path }, n)
} }
@@ -72,7 +73,7 @@ impl MphfLayer {
let evidence = Evidence::open(&dir.join(EVIDENCE_FILE))?; let evidence = Evidence::open(&dir.join(EVIDENCE_FILE))?;
let fingerprint = FingerprintVec::open(&dir.join(FINGERPRINT_FILE))?; let fingerprint = FingerprintVec::open(&dir.join(FINGERPRINT_FILE))?;
let n = evidence.len(); let n = evidence.len();
let unitigs = UnitigFileReader::open(&dir.join(UNITIGS_FILE))?; let unitigs = Arc::new(UnitigFileReader::open(&dir.join(UNITIGS_FILE))?);
(LayerEvidence::Hybrid { evidence, unitigs, fingerprint }, n) (LayerEvidence::Hybrid { evidence, unitigs, fingerprint }, n)
} }
}; };
@@ -184,16 +185,20 @@ impl MphfLayer {
/// Iterate over all canonical kmers in the layer, in deterministic order. /// Iterate over all canonical kmers in the layer, in deterministic order.
/// ///
/// The iteration order follows the physical layout of `unitigs.bin` and is /// The iteration order follows the physical layout of `unitigs.bin` and is
/// **not** correlated with MPHF slot numbers. Multiple `KmerIter` instances /// **not** correlated with MPHF slot numbers. Owns a clone of the
/// can be held concurrently for as long as the layer lives, because each /// underlying `Arc<UnitigFileReader>` rather than borrowing `self` — `Send
/// + 'static`, so it can be handed to a threaded consumer (e.g.
/// `obipipeline`) directly, streamed from disk, with no need to collect
/// the layer's kmers into memory first. Multiple instances can be held
/// concurrently for as long as the underlying file exists, because each
/// carries its own cursor. /// carries its own cursor.
pub fn iter_kmers(&self) -> KmerIter<'_> { pub fn iter_kmers(&self) -> KmerIter {
let reader = match &self.ev { let reader = match &self.ev {
LayerEvidence::Exact { unitigs, .. } => unitigs, LayerEvidence::Exact { unitigs, .. } => unitigs,
LayerEvidence::Approx { unitigs, .. } => unitigs, LayerEvidence::Approx { unitigs, .. } => unitigs,
LayerEvidence::Hybrid { unitigs, .. } => unitigs, LayerEvidence::Hybrid { unitigs, .. } => unitigs,
}; };
KmerIter { inner: Box::new(reader.iter_indexed_canonical_kmers()) } KmerIter { inner: Box::new(reader.iter_indexed_canonical_kmers_owned()) }
} }
/// Iterate over all canonical kmers, each paired with its zero-based /// Iterate over all canonical kmers, each paired with its zero-based
@@ -202,7 +207,7 @@ impl MphfLayer {
/// Yields `(index, kmer)` where `index` starts at `0` for the first kmer /// Yields `(index, kmer)` where `index` starts at `0` for the first kmer
/// stored in the layer. This is the standard Rust `enumerate` adapter /// stored in the layer. This is the standard Rust `enumerate` adapter
/// applied to [`iter_kmers`](Self::iter_kmers). /// applied to [`iter_kmers`](Self::iter_kmers).
pub fn enumerate_kmers(&self) -> Enumerate<KmerIter<'_>> { pub fn enumerate_kmers(&self) -> Enumerate<KmerIter> {
self.iter_kmers().enumerate() self.iter_kmers().enumerate()
} }
@@ -210,7 +215,7 @@ impl MphfLayer {
/// ///
/// Each call to [`next`](Iterator::next) returns a [`Vec`] of up to `n` /// Each call to [`next`](Iterator::next) returns a [`Vec`] of up to `n`
/// kmers. The final batch may be shorter when the layer is exhausted. /// kmers. The final batch may be shorter when the layer is exhausted.
pub fn iter_kmers_batch(&self, n: usize) -> KmerBatchIter<'_> { pub fn iter_kmers_batch(&self, n: usize) -> KmerBatchIter {
KmerBatchIter { inner: self.iter_kmers(), batch_size: n } KmerBatchIter { inner: self.iter_kmers(), batch_size: n }
} }
@@ -220,7 +225,7 @@ impl MphfLayer {
/// Yields `(batch_start_index, Vec<CanonicalKmer>)` where /// Yields `(batch_start_index, Vec<CanonicalKmer>)` where
/// `batch_start_index` is a multiple of `n`. This is the standard Rust /// `batch_start_index` is a multiple of `n`. This is the standard Rust
/// `enumerate` adapter applied to [`iter_kmers_batch`](Self::iter_kmers_batch). /// `enumerate` adapter applied to [`iter_kmers_batch`](Self::iter_kmers_batch).
pub fn enumerate_kmers_batch(&self, n: usize) -> Enumerate<KmerBatchIter<'_>> { pub fn enumerate_kmers_batch(&self, n: usize) -> Enumerate<KmerBatchIter> {
self.iter_kmers_batch(n).enumerate() self.iter_kmers_batch(n).enumerate()
} }
} }
@@ -229,14 +234,16 @@ impl MphfLayer {
/// Iterator over the canonical kmers stored in a layer. /// Iterator over the canonical kmers stored in a layer.
/// ///
/// Produced by [`MphfLayer::iter_kmers`]. Multiple `KmerIter` instances can /// Produced by [`MphfLayer::iter_kmers`]. Owns an `Arc<UnitigFileReader>`
/// coexist concurrently for as long as the parent layer lives, because each /// clone internally (via `iter_indexed_canonical_kmers_owned`) instead of
/// holds its own cursor over the underlying [`UnitigFileReader`]. /// borrowing the parent layer — `Send + 'static`, streamed from disk one
pub struct KmerIter<'a> { /// kmer at a time, never materialised as a whole. Multiple `KmerIter`
inner: Box<dyn Iterator<Item = (CanonicalKmer, usize, usize)> + 'a>, /// instances can coexist concurrently, because each holds its own cursor.
pub struct KmerIter {
inner: Box<dyn Iterator<Item = (CanonicalKmer, usize, usize)> + Send>,
} }
impl<'a> Iterator for KmerIter<'a> { impl Iterator for KmerIter {
type Item = CanonicalKmer; type Item = CanonicalKmer;
/// Return the next canonical kmer in iteration order. /// Return the next canonical kmer in iteration order.
@@ -250,12 +257,12 @@ impl<'a> Iterator for KmerIter<'a> {
/// Produced by [`MphfLayer::iter_kmers_batch`]. Each call to [`next`](Self::next) /// Produced by [`MphfLayer::iter_kmers_batch`]. Each call to [`next`](Self::next)
/// returns a [`Vec`] of up to `batch_size` kmers. The last batch may be shorter /// returns a [`Vec`] of up to `batch_size` kmers. The last batch may be shorter
/// than `batch_size` when the layer is exhausted. /// than `batch_size` when the layer is exhausted.
pub struct KmerBatchIter<'a> { pub struct KmerBatchIter {
inner: KmerIter<'a>, inner: KmerIter,
batch_size: usize, batch_size: usize,
} }
impl<'a> Iterator for KmerBatchIter<'a> { impl Iterator for KmerBatchIter {
type Item = Vec<CanonicalKmer>; type Item = Vec<CanonicalKmer>;
/// Return the next batch of kmers, or `None` if the layer is exhausted. /// Return the next batch of kmers, or `None` if the layer is exhausted.
+29
View File
@@ -1,5 +1,6 @@
use std::fs::File; use std::fs::File;
use std::path::Path; use std::path::Path;
use std::sync::Arc;
use memmap2::Mmap; use memmap2::Mmap;
use obikseq::{CanonicalKmer, Kmer, Unitig}; use obikseq::{CanonicalKmer, Kmer, Unitig};
@@ -203,6 +204,34 @@ impl UnitigFileReader {
.map(move |(rank, kmer)| (kmer, chunk_id, rank)) .map(move |(rank, kmer)| (kmer, chunk_id, rank))
}) })
} }
/// Same streamed sequence as [`iter_indexed_canonical_kmers`](Self::iter_indexed_canonical_kmers),
/// but owning a clone of `self` instead of borrowing it — `Send + 'static`,
/// so it can be handed to a threaded consumer (e.g. `obipipeline`) without
/// first collecting the layer's k-mers into memory. Reads `mmap` fresh
/// through the `Arc` on every step; no data is duplicated up front.
pub fn iter_indexed_canonical_kmers_owned(
self: &Arc<Self>,
) -> impl Iterator<Item = (CanonicalKmer, usize, usize)> + Send + 'static {
let this = Arc::clone(self);
let k = this.k;
let n = this.n_unitigs;
let mut offset = 0usize;
(0..n)
.map(move |chunk_id| {
let mmap = &*this.mmap;
let seql = mmap[offset] as usize + k;
let byte_len = (seql + 3) / 4;
let bytes = mmap[offset + 1..offset + 1 + byte_len].to_vec().into_boxed_slice();
offset += 1 + byte_len;
(chunk_id, Unitig::new((seql % 4) as u8, bytes))
})
.flat_map(|(chunk_id, u)| {
u.into_canonical_kmers()
.enumerate()
.map(move |(rank, kmer)| (kmer, chunk_id, rank))
})
}
} }
fn read_idx(path: &Path) -> SKResult<(usize, usize, u8, Vec<u32>)> { fn read_idx(path: &Path) -> SKResult<(usize, usize, u8, Vec<u32>)> {