256 lines
12 KiB
Rust
256 lines
12 KiB
Rust
|
use std::path::Path;
|
||
|
|
use std::sync::atomic::{AtomicU8, Ordering};
|
||
|
|
use std::sync::Arc;
|
||
|
|
|
||
|
|
use rayon::prelude::*;
|
||
|
|
|
||
|
|
use obicompactvec::{FamilyMask, SiblingAnnexBuilder};
|
||
|
|
use obikpartitionner::KmerPartition;
|
||
|
|
use obipipeline::ThrottleGuard;
|
||
|
|
use obikseq::CanonicalKmer;
|
||
|
|
use obilayeredmap::MphfLayer;
|
||
|
|
use obilayeredmap::meta::PartitionMeta;
|
||
|
|
use obiskio::UnitigFileReader;
|
||
|
|
use obisys::progress_bar;
|
||
|
|
|
||
|
|
use crate::error::{OKIError, OKIResult};
|
||
|
|
use crate::index::KmerIndex;
|
||
|
|
|
||
|
|
use super::cache::PartitionCache;
|
||
|
|
use super::helpers::{central_base, partition_of};
|
||
|
|
use super::{olm_to_ok, ANNEX_FILE_NAME, INDEX_SUBDIR};
|
||
|
|
|
||
|
|
// ── obipipeline data types ─────────────────────────────────────────────────
|
||
|
|
|
||
|
|
/// A batch of this layer's distinct k-mers (local MPHF slot + k-mer), the
|
||
|
|
/// pipeline's source item — batched, not one k-mer per item, so that
|
||
|
|
/// pipeline messages and their synchronisation cost stay amortised over
|
||
|
|
/// thousands of lookups (see `build_layer_sibling_annex`'s comment on
|
||
|
|
/// `BATCH_SIZE`). Carries the throttle permit for the whole batch, moved
|
||
|
|
/// (not cloned) into the corresponding `VariantBatch` — a 1-to-1 transform,
|
||
|
|
/// unlike the fan-out `Flat` stage this replaced, needs no `Arc` sharing.
|
||
|
|
struct SourceBatch {
|
||
|
|
items: Vec<(usize, CanonicalKmer)>,
|
||
|
|
_permit: ThrottleGuard,
|
||
|
|
}
|
||
|
|
|
||
|
|
/// One batch's worth of central-substitution variants (up to 3 per source
|
||
|
|
/// 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
|
||
|
|
/// hit. `(dest_partition, variant, source_slot, base)` per entry.
|
||
|
|
struct VariantBatch {
|
||
|
|
items: Vec<(usize, CanonicalKmer, usize, u8)>,
|
||
|
|
_permit: ThrottleGuard,
|
||
|
|
}
|
||
|
|
|
||
|
|
enum SibData {
|
||
|
|
Batch(SourceBatch),
|
||
|
|
Variants(VariantBatch),
|
||
|
|
}
|
||
|
|
|
||
|
|
impl KmerIndex {
|
||
|
|
/// Build the sibling-count/minorant annex for every layer of every
|
||
|
|
/// partition of this (already built) index, writing one annex file per
|
||
|
|
/// layer alongside its existing index files. Safe to call again later
|
||
|
|
/// (e.g. after a fresh `merge`) — each run simply overwrites the annex
|
||
|
|
/// files of the index it is called on.
|
||
|
|
///
|
||
|
|
/// Construction only — no statistics gathered here on purpose: this is
|
||
|
|
/// meant to run routinely (it is the artefact the SNP-family distances
|
||
|
|
/// will consume), while the sibling-count distribution
|
||
|
|
/// ([`sibling_annex_stats`](Self::sibling_annex_stats)) is a separate,
|
||
|
|
/// occasional diagnostic pass over the result, not run every time.
|
||
|
|
///
|
||
|
|
/// Cross-partition/cross-layer lookups are required (a k-mer's siblings
|
||
|
|
/// can live in any partition), but the layer loop itself — and thus the
|
||
|
|
/// annex file this produces — stays local to one layer at a time.
|
||
|
|
pub fn build_sibling_annex(&self) -> OKIResult<()> {
|
||
|
|
let n_parts = self.n_partitions();
|
||
|
|
let n_bits = n_parts.trailing_zeros() as usize;
|
||
|
|
|
||
|
|
let partition = KmerPartition::open_with_config(
|
||
|
|
&self.root_path,
|
||
|
|
self.kmer_size(),
|
||
|
|
self.minimizer_size(),
|
||
|
|
n_bits,
|
||
|
|
)
|
||
|
|
.map_err(OKIError::Partition)?;
|
||
|
|
|
||
|
|
tracing::info!("opening {n_parts} partition(s) for the sibling-annex sweep");
|
||
|
|
let cache = Arc::new(PartitionCache::build(&partition, n_parts, self.meta.config.with_counts)?);
|
||
|
|
|
||
|
|
let pb = progress_bar("sibling_annex", n_parts as u64, "partitions");
|
||
|
|
let mut total_slots: u64 = 0;
|
||
|
|
for part in 0..n_parts {
|
||
|
|
let index_dir = self.partition().part_dir(part).join(INDEX_SUBDIR);
|
||
|
|
if !index_dir.exists() {
|
||
|
|
pb.inc(1);
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?;
|
||
|
|
|
||
|
|
let mut part_slots: u64 = 0;
|
||
|
|
for l in 0..meta.n_layers {
|
||
|
|
let layer_dir = index_dir.join(format!("layer_{l}"));
|
||
|
|
part_slots += self.build_layer_sibling_annex(&layer_dir, n_parts, &cache)?;
|
||
|
|
}
|
||
|
|
total_slots += part_slots;
|
||
|
|
pb.inc(1);
|
||
|
|
pb.set_message(format!("partition {part}: {part_slots} kmers ({total_slots} total)"));
|
||
|
|
}
|
||
|
|
pb.finish_and_clear();
|
||
|
|
tracing::info!("sibling annex built — {total_slots} kmers across {n_parts} partitions");
|
||
|
|
|
||
|
|
Ok(())
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Returns the number of distinct k-mers (annex slots) processed, for
|
||
|
|
/// progress reporting.
|
||
|
|
fn build_layer_sibling_annex(
|
||
|
|
&self,
|
||
|
|
layer_dir: &Path,
|
||
|
|
n_parts: usize,
|
||
|
|
cache: &Arc<PartitionCache>,
|
||
|
|
) -> OKIResult<u64> {
|
||
|
|
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 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 ────────────
|
||
|
|
let mut slot_kmer: Vec<Option<CanonicalKmer>> = vec![None; n_slots];
|
||
|
|
let reader = UnitigFileReader::open_sequential(&layer_dir.join("unitigs.bin"))
|
||
|
|
.map_err(OKIError::Partition)?;
|
||
|
|
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
|
||
|
|
if let Some(slot) = mphf.find(kmer) {
|
||
|
|
slot_kmer[slot] = Some(kmer);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
let k = self.kmer_size();
|
||
|
|
|
||
|
|
// ── Reconciliation state, initialised with each slot's own base —
|
||
|
|
// that member is trivially present, no lookup needed. Built before
|
||
|
|
// the pipeline runs, from the same enumeration, since `sources`
|
||
|
|
// below is consumed as a throttled iterator, not collected.
|
||
|
|
// `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 slots — a
|
||
|
|
// lock-free `fetch_or` avoids needing any synchronisation beyond
|
||
|
|
// that. ─────────────────────────────────────────────────────────
|
||
|
|
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))) {
|
||
|
|
mask[slot].fetch_or(1 << central_base(kmer, k), Ordering::Relaxed);
|
||
|
|
}
|
||
|
|
|
||
|
|
// ── obipipeline: a *batch* transform, not a per-k-mer `Flat` one —
|
||
|
|
// the actual cross-partition lookup reuses
|
||
|
|
// `KmerPartition::query_partition_with` (the same batching mechanism
|
||
|
|
// `obikmer query` already uses: open a partition's files once,
|
||
|
|
// answer a whole batch of queries against it) instead of one lookup
|
||
|
|
// per pipeline item. A per-item lookup (tried first) reopened/
|
||
|
|
// re-mmap'd every target partition's files on every single variant
|
||
|
|
// — fine at toy scale, but ~90% system time against a real index,
|
||
|
|
// observed in practice. A *later* attempt still pushed one pipeline
|
||
|
|
// message per generated variant (a `Flat` stage, `SourceItem` =>
|
||
|
|
// `VariantQuery`, one k-mer in => up to 3 variants out as separate
|
||
|
|
// messages) — cheaper than reopening files, but sampling a real run
|
||
|
|
// showed most wall-clock time going into per-message channel
|
||
|
|
// send/notify syscalls instead of the lookup itself: the pipeline's
|
||
|
|
// whole point is amortising synchronisation over a batch, and a
|
||
|
|
// single k-mer's ≤3 variants is far too fine a granularity for
|
||
|
|
// that. Batching `BATCH_SIZE` source k-mers into one pipeline item
|
||
|
|
// — a plain 1-to-1 (`|`, not `||`) transform, batch in, batch of
|
||
|
|
// variants out, one message either way — keeps the per-message
|
||
|
|
// synchronisation cost amortised over thousands of lookups instead
|
||
|
|
// of one to three. ──────────────────────────────────────────────
|
||
|
|
const BATCH_SIZE: usize = 4096;
|
||
|
|
let n_workers = obisys::effective_parallelism();
|
||
|
|
let capacity = 256;
|
||
|
|
|
||
|
|
// Throttling limits how many *batches* are in flight at once — the
|
||
|
|
// permit is acquired per batch (not per k-mer) in the source
|
||
|
|
// thread, and released once its `VariantBatch` has been read out of
|
||
|
|
// the pipeline by the accumulation loop below. See
|
||
|
|
// `obipipeline::throttle`'s docs for why this is required, not
|
||
|
|
// optional, once a `Flat`-style stage sits in the pipeline.
|
||
|
|
let sources: Vec<(usize, CanonicalKmer)> = slot_kmer
|
||
|
|
.iter()
|
||
|
|
.enumerate()
|
||
|
|
.filter_map(|(slot, maybe_kmer)| maybe_kmer.map(|kmer| (slot, kmer)))
|
||
|
|
.collect();
|
||
|
|
let batches: Vec<Vec<(usize, CanonicalKmer)>> = sources
|
||
|
|
.chunks(BATCH_SIZE)
|
||
|
|
.map(|chunk| chunk.to_vec())
|
||
|
|
.collect();
|
||
|
|
let throttled = obipipeline::throttle(batches.into_iter(), n_workers).map(|t| SourceBatch {
|
||
|
|
items: t.item,
|
||
|
|
_permit: t.guard,
|
||
|
|
});
|
||
|
|
|
||
|
|
let pipe = obipipeline::make_pipe! {
|
||
|
|
SibData : SourceBatch => VariantBatch,
|
||
|
|
| {
|
||
|
|
move |batch: SourceBatch| -> VariantBatch {
|
||
|
|
let mut items = Vec::with_capacity(batch.items.len() * 3);
|
||
|
|
for (slot, kmer) in batch.items {
|
||
|
|
for variant in kmer.central_canonical_neighbors() {
|
||
|
|
if variant == kmer {
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
items.push((
|
||
|
|
partition_of(variant, n_parts),
|
||
|
|
variant,
|
||
|
|
slot,
|
||
|
|
central_base(variant, k),
|
||
|
|
));
|
||
|
|
}
|
||
|
|
}
|
||
|
|
VariantBatch { items, _permit: batch._permit }
|
||
|
|
}
|
||
|
|
} : Batch => Variants,
|
||
|
|
};
|
||
|
|
|
||
|
|
// ── Group generated variants by destination partition. `cache`
|
||
|
|
// holds every partition already mmap'd (no more `open()` cost), but
|
||
|
|
// `mmap` pages are still loaded on demand and can be evicted — a
|
||
|
|
// lookup is not free just because the file isn't reopened. Grouping
|
||
|
|
// keeps one partition's pages hot while its whole batch is resolved,
|
||
|
|
// instead of faulting pages in and out as lookups jump between
|
||
|
|
// partitions in whatever order the pipeline happens to produce
|
||
|
|
// 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();
|
||
|
|
for vb in pipe.apply(throttled, n_workers, capacity) {
|
||
|
|
for (dest_partition, variant, source_slot, base) in vb.items {
|
||
|
|
outgoing[dest_partition].push((variant, source_slot, base));
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// ── Resolve each partition's batch against the cache in one
|
||
|
|
// contiguous pass; parallelised across partitions (independent,
|
||
|
|
// read-only) so this keeps using multiple cores without giving up
|
||
|
|
// the per-partition locality above. ─────────────────────────────
|
||
|
|
outgoing.par_iter().enumerate().filter(|(_, q)| !q.is_empty()).for_each(|(dest, queries)| {
|
||
|
|
for &(variant, source_slot, base) in queries {
|
||
|
|
if cache.find(dest, variant) {
|
||
|
|
mask[source_slot].fetch_or(1 << base, Ordering::Relaxed);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
});
|
||
|
|
|
||
|
|
// ── Write the layer's annex file ─────────────────────────────────────
|
||
|
|
let annex_path = layer_dir.join(ANNEX_FILE_NAME);
|
||
|
|
let mut builder = SiblingAnnexBuilder::new(n_slots, &annex_path)?;
|
||
|
|
for (slot, m) in mask.iter().enumerate() {
|
||
|
|
if slot_kmer[slot].is_none() {
|
||
|
|
continue; // unused MPHF slot, if any — leave at the sentinel
|
||
|
|
}
|
||
|
|
builder.set(slot, FamilyMask::from_bits(m.load(Ordering::Relaxed)));
|
||
|
|
}
|
||
|
|
builder.close()?;
|
||
|
|
|
||
|
|
Ok(n_slots as u64)
|
||
|
|
}
|
||
|
|
}
|