OBIKmers/obikmer
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 34 Wiki Activity

rename KmerPartition to KmerPartitions and update Mat enum

Browse Source 
Rename the KmerPartition type to KmerPartitions across obikindex, obikpartitionner, and obikphylo/siblings to reflect an updated data model. Update the Mat enum in siblings/cache.rs to add a SparsePresence variant and simplify opening logic by delegating sparse versus dense detection to PersistentBitMatrix. Apply consistent code formatting, import reordering, and multi-line refactoring throughout the affected modules.
This commit is contained in:
Eric Coissac
2026-08-20 15:57:33 +02:00
parent 76cbd3a886
commit f9ef6b8391
26 changed files with 1290 additions and 733 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/obikphylo/src/siblings/build.rs
+225 -195
View File
@@ -1,22 +1,22 @@
use std::path::Path;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::sync::atomic::Ordering;
use rayon::prelude::*;
use obikpartitionner::KmerPartition;
use obipipeline::ThrottleGuard;
use obikpartitionner::KmerPartitions;
use obikseq::CanonicalKmer;
use obilayeredmap::MphfLayer;
use obilayeredmap::meta::IndexMode;
use obipipeline::ThrottleGuard;
use obisys::progress_bar;
use obikindex::{OKIError, OKIResult};
use obikindex::KmerIndex;
use obikindex::{OKIError, OKIResult};
use super::cache::PartitionCache;
use super::helpers::{central_base, is_minorant};
use super::{olm_to_ok, FamilyMask, SiblingAnnexBuilder, ANNEX_FILE_NAME};
use super::{ANNEX_FILE_NAME, FamilyMask, SiblingAnnexBuilder, olm_to_ok};
// ── obipipeline data types ─────────────────────────────────────────────────
@@ -78,7 +78,7 @@ impl SiblingAnnexBuildExt for KmerIndex {
let n_parts = self.n_partitions();
let n_bits = n_parts.trailing_zeros() as usize;
let partition = KmerPartition::open_with_config(
let partition = KmerPartitions::open_with_config(
self.root_path(),
self.kmer_size(),
self.minimizer_size(),
@@ -87,7 +87,11 @@ impl SiblingAnnexBuildExt for KmerIndex {
.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 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;
@@ -102,12 +106,19 @@ impl SiblingAnnexBuildExt for KmerIndex {
let mut part_slots: u64 = 0;
for l in 0..meta.n_layers {
part_slots += build_layer_sibling_annex(
self, &self.partition().layer_dir(part, l), &meta.mode, n_parts, l, &cache,
self,
&self.partition().layer_dir(part, l),
&meta.mode,
n_parts,
l,
&cache,
)?;
}
total_slots += part_slots;
pb.inc(1);
pb.set_message(format!("partition {part}: {part_slots} kmers ({total_slots} total)"));
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");
@@ -139,213 +150,232 @@ fn build_layer_sibling_annex(
// agree; see `PartitionCache::fast_mode`'s docs.
let fast_mode = cache.fast_mode();
// ── The annex file itself is the reconciliation state, indexed by
// this layer's k-mer iteration order (the physical layout of
// `unitigs.bin`), never by MPHF slot — this layer's own k-mers are
// known members by construction, so no evidence check, no MPHF
// slot, and no slot -> k-mer reconstruction is needed or legitimate
// here (see `DevDocMD/architecture/siblings.md`: the MPHF is not
// invertible, and evidence answers membership, not identity). No
// separate in-memory accumulator: every concurrent write goes
// straight through `SiblingAnnexBuilder::atomic_slot` into the
// mmap — a layer can hold billions of k-mers, so a `Vec` shadowing
// the whole file in RAM just to copy it out again at the end (the
// previous design) doubles memory for no benefit once the builder
// itself can be written concurrently.
//
// `Arc<SiblingAnnexBuilder>`, not a bare `SiblingAnnexBuilder` —
// `Pipe::apply` requires 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. `atomic_slot`,
// not `set`, for every concurrent write below: the gather phase
// parallelises across destination partitions (independent
// `cache.find` calls, safe to run concurrently) and their hits can
// land on arbitrary, possibly-shared source entries — a lock-free
// `fetch_or` avoids needing any synchronisation beyond that.
// Reclaimed as a plain owned value (`Arc::try_unwrap`) once every
// concurrent phase below is done, for the sequential finalisation
// pass. ─────────────────────────────────────────────────────────
let annex_path = layer_dir.join(ANNEX_FILE_NAME);
let builder = Arc::new(SiblingAnnexBuilder::new(n, &annex_path)?);
// ── The annex file itself is the reconciliation state, indexed by
// this layer's k-mer iteration order (the physical layout of
// `unitigs.bin`), never by MPHF slot — this layer's own k-mers are
// known members by construction, so no evidence check, no MPHF
// slot, and no slot -> k-mer reconstruction is needed or legitimate
// here (see `DevDocMD/architecture/siblings.md`: the MPHF is not
// invertible, and evidence answers membership, not identity). No
// separate in-memory accumulator: every concurrent write goes
// straight through `SiblingAnnexBuilder::atomic_slot` into the
// mmap — a layer can hold billions of k-mers, so a `Vec` shadowing
// the whole file in RAM just to copy it out again at the end (the
// previous design) doubles memory for no benefit once the builder
// itself can be written concurrently.
//
// `Arc<SiblingAnnexBuilder>`, not a bare `SiblingAnnexBuilder` —
// `Pipe::apply` requires 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. `atomic_slot`,
// not `set`, for every concurrent write below: the gather phase
// parallelises across destination partitions (independent
// `cache.find` calls, safe to run concurrently) and their hits can
// land on arbitrary, possibly-shared source entries — a lock-free
// `fetch_or` avoids needing any synchronisation beyond that.
// Reclaimed as a plain owned value (`Arc::try_unwrap`) once every
// concurrent phase below is done, for the sequential finalisation
// pass. ─────────────────────────────────────────────────────────
let annex_path = layer_dir.join(ANNEX_FILE_NAME);
let builder = Arc::new(SiblingAnnexBuilder::new(n, &annex_path)?);
// ── 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. `BATCH_SIZE` was originally tuned back when the
// per-k-mer cost inside this stage (minimiser via `RollingStat`)
// was ~3x higher than it is now (see `CanonicalKmerOf::minimizer`,
// a direct O(k) bit-arithmetic replacement) — at the old per-k-mer
// cost, dispatch overhead (the scheduler's single-threaded `Select`
// loop: one channel round-trip per batch) was negligible next to
// the work each batch represented; now that the work itself is
// cheaper, that fixed per-batch overhead is proportionally larger,
// so a bigger batch amortises it over more k-mers again. ─────────
const BATCH_SIZE: usize = 32768;
let n_workers = obisys::effective_parallelism();
let capacity = 256;
// ── 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. `BATCH_SIZE` was originally tuned back when the
// per-k-mer cost inside this stage (minimiser via `RollingStat`)
// was ~3x higher than it is now (see `CanonicalKmerOf::minimizer`,
// a direct O(k) bit-arithmetic replacement) — at the old per-k-mer
// cost, dispatch overhead (the scheduler's single-threaded `Select`
// loop: one channel round-trip per batch) was negligible next to
// the work each batch represented; now that the work itself is
// cheaper, that fixed per-batch overhead is proportionally larger,
// so a bigger batch amortises it over more k-mers again. ─────────
const BATCH_SIZE: usize = 32768;
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.
//
// 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 straight into the annex in
// this same pass, since this is exactly the iteration order the
// annex is keyed on — no separate seeding pass needed.
let seed_builder = Arc::clone(&builder);
let batches = mphf.enumerate_kmers_batch(BATCH_SIZE).map(move |(start, kmers)| {
// 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.
//
// 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 straight into the annex in
// this same pass, since this is exactly the iteration order the
// annex is keyed on — no separate seeding pass needed.
let seed_builder = Arc::clone(&builder);
let batches = mphf
.enumerate_kmers_batch(BATCH_SIZE)
.map(move |(start, kmers)| {
kmers
.into_iter()
.enumerate()
.map(|(i, kmer)| {
let base = central_base(kmer, k);
let bits = if fast_mode { FamilyMask::layer_bits(base, l) } else { FamilyMask::presence_bits(base) };
seed_builder.atomic_slot(start + i).fetch_or(bits, Ordering::Relaxed);
let bits = if fast_mode {
FamilyMask::layer_bits(base, l)
} else {
FamilyMask::presence_bits(base)
};
seed_builder
.atomic_slot(start + i)
.fetch_or(bits, Ordering::Relaxed);
(start + i, kmer)
})
.collect::<Vec<_>>()
});
// Diagnostic: count still-empty annex slots after seeding + cross-partition
// resolution. A non-zero count here means some k-mer's own base was never
// written (should be impossible after the batch_offset fix above).
let check_empty = |label: &str| {
let empty = (0..n).filter(|&slot| builder.get(slot).is_none()).count();
if empty > 0 {
tracing::warn!("{label}: {empty} mask slots still empty after resolution (layer {layer_dir:?})");
}
};
let throttled = obipipeline::throttle(batches, n_workers).map(|t| SourceBatch {
items: t.item,
_permit: t.guard,
});
// Diagnostic: count still-empty annex slots after seeding + cross-partition
// resolution. A non-zero count here means some k-mer's own base was never
// written (should be impossible after the batch_offset fix above).
let check_empty = |label: &str| {
let empty = (0..n).filter(|&slot| builder.get(slot).is_none()).count();
if empty > 0 {
tracing::warn!(
"{label}: {empty} mask slots still empty after resolution (layer {layer_dir:?})"
);
}
};
let throttled = obipipeline::throttle(batches, 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 (order, kmer) in batch.items {
for variant in kmer.central_canonical_neighbors() {
if variant == kmer {
continue;
}
items.push((
variant.partition(n_parts),
variant,
order,
central_base(variant, k),
));
let pipe = obipipeline::make_pipe! {
SibData : SourceBatch => VariantBatch,
| {
move |batch: SourceBatch| -> VariantBatch {
let mut items = Vec::with_capacity(batch.items.len() * 3);
for (order, kmer) in batch.items {
for variant in kmer.central_canonical_neighbors() {
if variant == kmer {
continue;
}
items.push((
variant.partition(n_parts),
variant,
order,
central_base(variant, k),
));
}
VariantBatch { items, _permit: batch._permit }
}
} : Batch => Variants,
};
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_order, base) in vb.items {
outgoing[dest_partition].push((variant, source_order, base));
// ── 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_order, base) in vb.items {
outgoing[dest_partition].push((variant, source_order, base));
}
}
check_empty("after_seeding");
// ── Resolve each partition's batch against the cache, parallelised
// over roughly equal-sized *chunks*, not over partitions — a plain
// `outgoing.par_iter()` over `n_parts` buckets gives one thread the
// whole of one partition's bucket, however large, and a lot of
// buckets are far from equal: a central-base substitution changes
// the minimiser (and thus routes to a different partition) only
// when the winning minimiser window overlaps the central position.
// For k=31/m=11 that is ~11 of the 21 possible windows — so *most*
// of the remaining ~10/21 send the variant right back to the
// partition already being built. That self-partition bucket ends up
// far larger than any other, so the naive per-partition split
// leaves one thread grinding through it alone long after every
// other partition's (tiny) bucket is done — confirmed by sampling a
// real run: one thread solely in `MphfLayer::find` while the rest
// of the pool sits idle. Splitting each non-empty bucket into
// `chunk_size`-sized pieces first keeps this pass's per-partition
// locality (each chunk is still contiguous within one partition,
// still resolved via `cache.find` grouped by that partition) while
// letting Rayon spread a single oversized bucket across several
// threads instead of pinning it to one. ──────────────────────────
let chunk_size = ((outgoing.iter().map(Vec::len).sum::<usize>() / n_workers).max(1)).min(4096);
let work: Vec<(usize, &[(CanonicalKmer, usize, u8)])> = outgoing
.iter()
.enumerate()
.filter(|(_, q)| !q.is_empty())
.flat_map(|(dest, q)| q.chunks(chunk_size).map(move |c| (dest, c)))
.collect();
work.par_iter().for_each(|&(dest, chunk)| {
for &(variant, source_order, base) in chunk {
if let Some(layer) = cache.find(dest, variant) {
let bits = if fast_mode {
FamilyMask::layer_bits(base, layer)
} else {
FamilyMask::presence_bits(base)
};
builder
.atomic_slot(source_order)
.fetch_or(bits, Ordering::Relaxed);
}
}
});
check_empty("after_seeding");
check_empty("after_resolution");
// ── Resolve each partition's batch against the cache, parallelised
// over roughly equal-sized *chunks*, not over partitions — a plain
// `outgoing.par_iter()` over `n_parts` buckets gives one thread the
// whole of one partition's bucket, however large, and a lot of
// buckets are far from equal: a central-base substitution changes
// the minimiser (and thus routes to a different partition) only
// when the winning minimiser window overlaps the central position.
// For k=31/m=11 that is ~11 of the 21 possible windows — so *most*
// of the remaining ~10/21 send the variant right back to the
// partition already being built. That self-partition bucket ends up
// far larger than any other, so the naive per-partition split
// leaves one thread grinding through it alone long after every
// other partition's (tiny) bucket is done — confirmed by sampling a
// real run: one thread solely in `MphfLayer::find` while the rest
// of the pool sits idle. Splitting each non-empty bucket into
// `chunk_size`-sized pieces first keeps this pass's per-partition
// locality (each chunk is still contiguous within one partition,
// still resolved via `cache.find` grouped by that partition) while
// letting Rayon spread a single oversized bucket across several
// threads instead of pinning it to one. ──────────────────────────
let chunk_size = ((outgoing.iter().map(Vec::len).sum::<usize>() / n_workers).max(1)).min(4096);
let work: Vec<(usize, &[(CanonicalKmer, usize, u8)])> = outgoing
.iter()
.enumerate()
.filter(|(_, q)| !q.is_empty())
.flat_map(|(dest, q)| q.chunks(chunk_size).map(move |c| (dest, c)))
.collect();
work.par_iter().for_each(|&(dest, chunk)| {
for &(variant, source_order, base) in chunk {
if let Some(layer) = cache.find(dest, variant) {
let bits = if fast_mode { FamilyMask::layer_bits(base, layer) } else { FamilyMask::presence_bits(base) };
builder.atomic_slot(source_order).fetch_or(bits, Ordering::Relaxed);
}
}
});
check_empty("after_resolution");
// ── Write the layer's annex file, indexed by iteration order — a
// second streamed pass over `unitigs.bin` (via `enumerate_kmers`),
// now that every entry's mask is final; never a `Vec` hold of the
// whole layer. The minorant flag is computed here, not by every
// later reader: this is the one place the whole family's *final*
// mask and this entry's own k-mer (already in hand, no extra
// lookup) are both available together. Every consumer that only
// needs to know "is this k-mer the family's minorant" — the common
// case, since a family is tallied once, at its minorant — reads
// 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).
// Every concurrent phase above is done and its `Arc` clone (the
// seeding pass's `seed_builder`) has already been dropped along
// with the now-fully-drained `batches` iterator — this is the only
// remaining strong reference, so reclaiming plain ownership for the
// sequential `set` calls below is guaranteed to succeed.
let mut builder = Arc::try_unwrap(builder).unwrap_or_else(|_| panic!("sibling-annex builder still shared after every concurrent phase completed"));
for (order, kmer) in mphf.enumerate_kmers() {
let final_mask = builder.get(order).expect("seeded by construction");
let minorant = is_minorant(kmer, final_mask, k);
builder.set(order, final_mask.with_minorant(minorant));
}
builder.close()?;
// ── Write the layer's annex file, indexed by iteration order — a
// second streamed pass over `unitigs.bin` (via `enumerate_kmers`),
// now that every entry's mask is final; never a `Vec` hold of the
// whole layer. The minorant flag is computed here, not by every
// later reader: this is the one place the whole family's *final*
// mask and this entry's own k-mer (already in hand, no extra
// lookup) are both available together. Every consumer that only
// needs to know "is this k-mer the family's minorant" — the common
// case, since a family is tallied once, at its minorant — reads
// 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).
// Every concurrent phase above is done and its `Arc` clone (the
// seeding pass's `seed_builder`) has already been dropped along
// with the now-fully-drained `batches` iterator — this is the only
// remaining strong reference, so reclaiming plain ownership for the
// sequential `set` calls below is guaranteed to succeed.
let mut builder = Arc::try_unwrap(builder).unwrap_or_else(|_| {
panic!("sibling-annex builder still shared after every concurrent phase completed")
});
for (order, kmer) in mphf.enumerate_kmers() {
let final_mask = builder.get(order).expect("seeded by construction");
let minorant = is_minorant(kmer, final_mask, k);
builder.set(order, final_mask.with_minorant(minorant));
}
builder.close()?;
Ok(n as u64)
}