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:
@@ -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()
|
// `enumerate_kmers_batch` — never a full-layer `Vec` collect (a
|
||||||
.filter_map(|(slot, maybe_kmer)| maybe_kmer.map(|kmer| (slot, kmer)))
|
// layer can hold billions of k-mers; see the "no full collect"
|
||||||
.collect();
|
// rule). Each k-mer's own base is seeded into `mask` in this same
|
||||||
let batches: Vec<Vec<(usize, CanonicalKmer)>> = sources
|
// pass, since this is exactly the iteration order `mask` is keyed
|
||||||
.chunks(BATCH_SIZE)
|
// on — no separate seeding pass needed.
|
||||||
.map(|chunk| chunk.to_vec())
|
let seed_mask = Arc::clone(&mask);
|
||||||
.collect();
|
let batches = mphf.enumerate_kmers_batch(BATCH_SIZE).map(move |(start, kmers)| {
|
||||||
let throttled = obipipeline::throttle(batches.into_iter(), n_workers).map(|t| SourceBatch {
|
kmers
|
||||||
|
.into_iter()
|
||||||
|
.enumerate()
|
||||||
|
.map(|(i, kmer)| {
|
||||||
|
seed_mask[start + i].fetch_or(1 << central_base(kmer, k), Ordering::Relaxed);
|
||||||
|
(start + i, kmer)
|
||||||
|
})
|
||||||
|
.collect::<Vec<_>>()
|
||||||
|
});
|
||||||
|
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)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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())
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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;
|
||||||
|
|||||||
@@ -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)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -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};
|
||||||
|
|||||||
@@ -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.
|
||||||
|
|||||||
@@ -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>)> {
|
||||||
|
|||||||
Reference in New Issue
Block a user