refactor: merge KmerPartitions into KmerIndex and rename obikpartition

Consolidates partition logic, metadata storage, and layer management directly into KmerIndex. Renames obikpartitionner to obikpartition, retaining only PartitionRouter for superkmer routing. Removes intermediate .partition() accessors in favor of direct methods on the index and updates PartitionCache::build to accept &KmerIndex directly. Derives n_partitions from config.n_bits and consolidates k-mer/minimizer sizes into IndexMeta.config. Fixes a regression where PartitionRouter::open incorrectly defaulted to closed.
This commit is contained in:
Eric Coissac
2026-08-20 20:23:58 +02:00
parent b5ec0122d0
commit 6c860f120f
49 changed files with 357 additions and 427 deletions
+8
View File
@@ -11,6 +11,14 @@ obiskio = { path = "../obiskio" }
obisys = { path = "../obisys" }
obicompactvec = { path = "../obicompactvec" }
obilayeredmap = { path = "../obilayeredmap" }
obidebruinj = { path = "../obidebruinj" }
obipipeline = { path = "../obipipeline" }
obikentropy = { path = "../obikentropy" }
cacheline-ef = "1.1"
epserde = "0.8"
ptr_hash = "1.1"
niffler = "3.0.0"
memmap2 = "0.9.10"
ndarray = "0.17"
rayon = "1"
crossbeam-channel = "0.5"
+4 -4
View File
@@ -25,16 +25,16 @@ fn main() -> anyhow::Result<()> {
let mut first_mismatch = None;
for part in 0..n_parts {
let index_dir_sparse = sparse.partition().index_dir(part);
let index_dir_dense = dense.partition().index_dir(part);
let index_dir_sparse = sparse.index_dir(part);
let index_dir_dense = dense.index_dir(part);
if !index_dir_sparse.exists() || !index_dir_dense.exists() {
continue;
}
for layer in 0..n_layers {
let layer_dir_sparse = sparse.partition().layer_dir(part, layer);
let layer_dir_dense = dense.partition().layer_dir(part, layer);
let layer_dir_sparse = sparse.layer_dir(part, layer);
let layer_dir_dense = dense.layer_dir(part, layer);
if !layer_dir_sparse.exists() || !layer_dir_dense.exists() {
continue;
+65
View File
@@ -0,0 +1,65 @@
use std::path::Path;
use obicompactvec::{PersistentBitVecBuilder, PersistentCompactIntVecBuilder};
use obilayeredmap::meta::PartitionMeta;
use obilayeredmap::{layer_dir, IndexMode, OLMError};
use obiskio::{SKError, SKResult};
// ── olm_to_sk ────────────────────────────────────────────────────────────────
pub(crate) fn olm_to_sk(e: OLMError, context: &'static str) -> SKError {
match e {
OLMError::Io(e) => SKError::Io(e),
other => SKError::InvalidData {
context,
detail: other.to_string(),
},
}
}
// ── load_meta ────────────────────────────────────────────────────────────────
/// Load PartitionMeta, or recover it by probing layer directories.
/// Indexes built before meta.json was introduced lack the file.
pub(crate) fn load_meta(dir: &Path, context: &'static str) -> SKResult<PartitionMeta> {
match PartitionMeta::load(dir) {
Ok(m) => Ok(m),
Err(e) if matches!(e, OLMError::Io(ref io_e) if io_e.kind() == std::io::ErrorKind::NotFound) =>
{
let mut n = 0usize;
while layer_dir(dir, n).exists() {
n += 1;
}
let m = PartitionMeta {
n_layers: n,
mode: IndexMode::default(),
};
m.save(dir).map_err(|e| olm_to_sk(e, context))?;
Ok(m)
}
Err(e) => Err(olm_to_sk(e, context)),
}
}
// ── ColBuilder ────────────────────────────────────────────────────────────────
pub(crate) enum ColBuilder {
Bit(PersistentBitVecBuilder),
Int(PersistentCompactIntVecBuilder),
}
impl ColBuilder {
pub(crate) fn set_val(&mut self, slot: usize, value: u32) {
match self {
ColBuilder::Bit(b) => b.set(slot, value > 0),
ColBuilder::Int(b) => b.set(slot, value),
}
}
pub(crate) fn close(self) -> SKResult<()> {
match self {
ColBuilder::Bit(b) => b.close().map_err(SKError::Io),
ColBuilder::Int(b) => b.close().map_err(SKError::Io),
}
}
}
+2 -2
View File
@@ -74,7 +74,7 @@ impl KmerIndex {
if use_counts {
let stores: Vec<_> = (0..n_parts)
.into_par_iter()
.map(|i| self.partition.count_store(i).map_err(OKIError::Partition))
.map(|i| self.count_store(i).map_err(OKIError::Partition))
.collect::<OKIResult<_>>()?;
let global = LayeredStore::new(stores);
@@ -105,7 +105,7 @@ impl KmerIndex {
} else {
let stores: Vec<_> = (0..n_parts)
.into_par_iter()
.map(|i| self.partition.presence_store(i).map_err(OKIError::Partition))
.map(|i| self.presence_store(i).map_err(OKIError::Partition))
.collect::<OKIResult<_>>()?;
let global = LayeredStore::new(stores);
+5 -5
View File
@@ -5,7 +5,7 @@ use rayon::prelude::*;
use crate::error::{OKIError, OKIResult};
use crate::index::KmerIndex;
use obikpartitionner::KmerFilter;
use crate::KmerFilter;
impl KmerIndex {
/// Write a CSV table of all indexed kmers to `out`.
@@ -69,14 +69,14 @@ impl KmerIndex {
}
};
if debug {
self.partition
self
.iter_partition_kmers_located(i, use_counts, n_genomes, filters, |part, layer, kmer, row| {
let seq = String::from_utf8(kmer.to_ascii()).unwrap_or_else(|_| "?".repeat(kmer_size));
try_write(&mut buf, &row, &format!("{part},{layer},{seq}"))
})
.map_err(OKIError::Partition)?;
} else {
self.partition
self
.iter_partition_kmers(i, use_counts, n_genomes, filters, |kmer, row| {
let seq = String::from_utf8(kmer.to_ascii()).unwrap_or_else(|_| "?".repeat(kmer_size));
try_write(&mut buf, &row, &seq)
@@ -91,7 +91,7 @@ impl KmerIndex {
(0..n).into_par_iter().map(|i| {
let mut buf = Vec::<u8>::new();
if debug {
self.partition
self
.iter_partition_kmers_located(i, use_counts, n_genomes, filters, |part, layer, kmer, row| {
let seq = String::from_utf8(kmer.to_ascii()).unwrap_or_else(|_| "?".repeat(kmer_size));
write_row(&mut buf, &row, &format!("{part},{layer},{seq}"));
@@ -99,7 +99,7 @@ impl KmerIndex {
})
.map_err(OKIError::Partition)?;
} else {
self.partition
self
.iter_partition_kmers(i, use_counts, n_genomes, filters, |kmer, row| {
let seq = String::from_utf8(kmer.to_ascii()).unwrap_or_else(|_| "?".repeat(kmer_size));
write_row(&mut buf, &row, &seq);
+205
View File
@@ -0,0 +1,205 @@
use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix};
use obikseq::CanonicalKmer;
use obilayeredmap::{IndexMode, MphfLayer, OLMError};
use obiskio::{SKError, SKResult, UnitigFileReader};
use crate::filter::{KmerFilter, passes_all};
use crate::index::KmerIndex;
fn olm_to_sk(e: OLMError) -> SKError {
match e {
OLMError::Io(e) => SKError::Io(e),
other => SKError::InvalidData {
context: "dump",
detail: other.to_string(),
},
}
}
impl KmerIndex {
/// Iterate all indexed kmers in partition `part`, calling `cb(kmer, row)` for each
/// kmer that passes every filter in `filters`.
///
/// `use_counts = true` → reads count columns (u32 values per genome).
/// `use_counts = false` → reads presence columns, converted to 0/1 u32.
///
/// If no data matrix exists for a layer (pure set-membership, single genome),
/// a row of `n_genomes` ones is emitted for every kmer in that layer — unless
/// the filter rejects it, in which case the whole layer is skipped.
/// Like [`iter_partition_kmers`] but the callback returns `false` to stop early.
/// Returns `Ok(true)` if all kmers were visited, `Ok(false)` if the callback halted.
pub fn iter_partition_kmers(
&self,
part: usize,
use_counts: bool,
n_genomes: usize,
filters: &[Box<dyn KmerFilter>],
mut cb: impl FnMut(CanonicalKmer, Box<[u32]>) -> bool,
) -> SKResult<bool> {
let index_dir = self.index_dir(part);
if !index_dir.exists() {
return Ok(true);
}
let index_mode = self.partition_mode(part).unwrap_or(IndexMode::Exact);
let mut l = 0;
loop {
let layer_dir = self.layer_dir(part, l);
if !layer_dir.exists() {
break;
}
l += 1;
let mphf = MphfLayer::open(&layer_dir, &index_mode).map_err(olm_to_sk)?;
let reader = UnitigFileReader::open_sequential(&layer_dir.join("unitigs.bin"))?;
let counts_dir = layer_dir.join("counts");
let presence_dir = layer_dir.join("presence");
let cont = if use_counts && counts_dir.exists() {
let mat = PersistentCompactIntMatrix::open(&layer_dir).map_err(SKError::Io)?;
let mut cont = true;
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
if let Some(slot) = mphf.find(kmer) {
let row = mat.row(slot);
if passes_all(filters, kmer, &row, n_genomes) {
cont = cb(kmer, row);
if !cont {
break;
}
}
}
}
cont
} else if !use_counts && presence_dir.exists() {
let mat = PersistentBitMatrix::open(&layer_dir).map_err(SKError::Io)?;
let mut cont = true;
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
if let Some(slot) = mphf.find(kmer) {
let row: Box<[u32]> = mat.row(slot).iter().map(|&b| b as u32).collect();
if passes_all(filters, kmer, &row, n_genomes) {
cont = cb(kmer, row);
if !cont {
break;
}
}
}
}
cont
} else {
// No data matrix: implicit presence — all values are 1. `row`
// is identical for every kmer, but a filter can still depend
// on the kmer's own sequence (e.g. MinComplexity), so this
// cannot be evaluated once for the whole layer — filters must
// still be tested per kmer.
let all_present: Box<[u32]> = vec![1u32; n_genomes].into();
let mut cont = true;
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
if mphf.find(kmer).is_some()
&& passes_all(filters, kmer, &all_present, n_genomes)
{
cont = cb(kmer, all_present.clone());
if !cont {
break;
}
}
}
cont
};
if !cont {
return Ok(false);
}
}
Ok(true)
}
/// Like [`iter_partition_kmers`] but the callback also receives `(partition, layer)`
/// indices, enabling debug output that identifies where each kmer was stored.
/// Returns `Ok(true)` if all kmers were visited, `Ok(false)` if the callback halted.
pub fn iter_partition_kmers_located(
&self,
part: usize,
use_counts: bool,
n_genomes: usize,
filters: &[Box<dyn KmerFilter>],
mut cb: impl FnMut(usize, usize, CanonicalKmer, Box<[u32]>) -> bool,
) -> SKResult<bool> {
let index_dir = self.index_dir(part);
if !index_dir.exists() {
return Ok(true);
}
let index_mode = self.partition_mode(part).unwrap_or(IndexMode::Exact);
let mut layer = 0;
loop {
let layer_dir = self.layer_dir(part, layer);
if !layer_dir.exists() {
break;
}
let mphf = MphfLayer::open(&layer_dir, &index_mode).map_err(olm_to_sk)?;
let reader = UnitigFileReader::open_sequential(&layer_dir.join("unitigs.bin"))?;
let counts_dir = layer_dir.join("counts");
let presence_dir = layer_dir.join("presence");
let cont = if use_counts && counts_dir.exists() {
let mat = PersistentCompactIntMatrix::open(&layer_dir).map_err(SKError::Io)?;
let mut cont = true;
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
if let Some(slot) = mphf.find(kmer) {
let row = mat.row(slot);
if passes_all(filters, kmer, &row, n_genomes) {
cont = cb(part, layer, kmer, row);
if !cont {
break;
}
}
}
}
cont
} else if !use_counts && presence_dir.exists() {
let mat = PersistentBitMatrix::open(&layer_dir).map_err(SKError::Io)?;
let mut cont = true;
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
if let Some(slot) = mphf.find(kmer) {
let row: Box<[u32]> = mat.row(slot).iter().map(|&b| b as u32).collect();
if passes_all(filters, kmer, &row, n_genomes) {
cont = cb(part, layer, kmer, row);
if !cont {
break;
}
}
}
}
cont
} else {
// Same as iter_partition_kmers: row is constant but a filter
// may still depend on the kmer's own sequence, so this must
// be tested per kmer, not once for the whole layer.
let all_present: Box<[u32]> = vec![1u32; n_genomes].into();
let mut cont = true;
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
if mphf.find(kmer).is_some()
&& passes_all(filters, kmer, &all_present, n_genomes)
{
cont = cb(part, layer, kmer, all_present.clone());
if !cont {
break;
}
}
}
cont
};
if !cont {
return Ok(false);
}
layer += 1;
}
Ok(true)
}
}
+298
View File
@@ -0,0 +1,298 @@
use obicompactvec::FilterMask;
use obikseq::CanonicalKmer;
/// Trait for kmer filters.
///
/// `kmer` is the k-mer's own canonical sequence, reconstructed from the
/// source index's `unitigs.bin` (always present — see `rebuild_layer.rs`);
/// `row` contains raw per-genome counts (or 0/1 for presence/absence data).
/// `n_genomes` equals `row.len()`. Most filters only need `row` — `kmer` is
/// there for filters that reason about the k-mer's sequence itself (e.g.
/// [`MinComplexity`]).
pub trait KmerFilter: Send + Sync {
fn passes(&self, kmer: CanonicalKmer, row: &[u32], n_genomes: usize) -> bool;
/// Express this filter as a [`FilterMask`] column-operation expression.
///
/// Returns `Some(expr)` if the filter can be evaluated solely from matrix
/// column aggregates (no per-kmer row scan needed). Returns `None` if the
/// filter requires row-level inspection — always the case for a filter
/// that needs the k-mer's sequence, since a `FilterMask` only expresses
/// per-genome column aggregates, never per-slot sequence data.
///
/// `threshold` semantics in the returned mask use `>= threshold`, matching
/// [`obicompactvec::MatrixGroupOps`]. Implementations must add 1 to any
/// row-level threshold that uses strict `>` comparison.
fn column_mask_expr(&self, _n_genomes: usize) -> Option<FilterMask> {
None
}
}
/// True when `row` passes every filter in `filters`.
/// Returns `true` if `filters` is empty.
pub fn passes_all(
filters: &[Box<dyn KmerFilter>],
kmer: CanonicalKmer,
row: &[u32],
n_genomes: usize,
) -> bool {
filters.iter().all(|f| f.passes(kmer, row, n_genomes))
}
// ── Quorum filters ─────────────────────────────────────────────────────────────
fn present_count(row: &[u32], threshold: u32) -> usize {
row.iter().filter(|&&v| v > threshold).count()
}
/// At least `frac` fraction of genomes contain this kmer (count > `threshold`).
pub struct MinGenomeFraction {
pub frac: f64,
pub threshold: u32,
}
impl KmerFilter for MinGenomeFraction {
fn passes(&self, _kmer: CanonicalKmer, row: &[u32], n_genomes: usize) -> bool {
let p = present_count(row, self.threshold);
p as f64 / n_genomes as f64 >= self.frac
}
fn column_mask_expr(&self, n_genomes: usize) -> Option<FilterMask> {
let t = self.threshold.checked_add(1)?;
let min_count = (self.frac * n_genomes as f64).ceil() as usize;
Some(FilterMask::PresenceGeq {
indices: (0..n_genomes).collect(),
threshold: t,
min_count,
})
}
}
/// At most `frac` fraction of genomes contain this kmer (count > `threshold`).
pub struct MaxGenomeFraction {
pub frac: f64,
pub threshold: u32,
}
impl KmerFilter for MaxGenomeFraction {
fn passes(&self, _kmer: CanonicalKmer, row: &[u32], n_genomes: usize) -> bool {
let p = present_count(row, self.threshold);
p as f64 / n_genomes as f64 <= self.frac
}
fn column_mask_expr(&self, n_genomes: usize) -> Option<FilterMask> {
let t = self.threshold.checked_add(1)?;
let max_count = (self.frac * n_genomes as f64).floor() as usize;
Some(FilterMask::PresenceLeq {
indices: (0..n_genomes).collect(),
threshold: t,
max_count,
})
}
}
/// At least `count` genomes contain this kmer (count > `threshold`).
pub struct MinGenomeCount {
pub count: usize,
pub threshold: u32,
}
impl KmerFilter for MinGenomeCount {
fn passes(&self, _kmer: CanonicalKmer, row: &[u32], _n_genomes: usize) -> bool {
present_count(row, self.threshold) >= self.count
}
fn column_mask_expr(&self, n_genomes: usize) -> Option<FilterMask> {
let t = self.threshold.checked_add(1)?;
Some(FilterMask::PresenceGeq {
indices: (0..n_genomes).collect(),
threshold: t,
min_count: self.count,
})
}
}
/// At most `count` genomes contain this kmer (count > `threshold`).
pub struct MaxGenomeCount {
pub count: usize,
pub threshold: u32,
}
impl KmerFilter for MaxGenomeCount {
fn passes(&self, _kmer: CanonicalKmer, row: &[u32], _n_genomes: usize) -> bool {
present_count(row, self.threshold) <= self.count
}
fn column_mask_expr(&self, n_genomes: usize) -> Option<FilterMask> {
let t = self.threshold.checked_add(1)?;
Some(FilterMask::PresenceLeq {
indices: (0..n_genomes).collect(),
threshold: t,
max_count: self.count,
})
}
}
// ── Total-count filters (count indexes only) ───────────────────────────────────
/// Sum of counts across all genomes >= `total`.
pub struct MinTotalCount {
pub total: u32,
}
impl KmerFilter for MinTotalCount {
fn passes(&self, _kmer: CanonicalKmer, row: &[u32], _n_genomes: usize) -> bool {
row.iter().sum::<u32>() >= self.total
}
fn column_mask_expr(&self, n_genomes: usize) -> Option<FilterMask> {
Some(FilterMask::SumGeq {
indices: (0..n_genomes).collect(),
min_sum: self.total,
})
}
}
/// Sum of counts across all genomes <= `total`.
pub struct MaxTotalCount {
pub total: u32,
}
impl KmerFilter for MaxTotalCount {
fn passes(&self, _kmer: CanonicalKmer, row: &[u32], _n_genomes: usize) -> bool {
row.iter().sum::<u32>() <= self.total
}
fn column_mask_expr(&self, n_genomes: usize) -> Option<FilterMask> {
Some(FilterMask::SumLeq {
indices: (0..n_genomes).collect(),
max_sum: self.total,
})
}
}
// ── Group-based quorum filter ─────────────────────────────────────────────────
/// Quorum filter operating on pre-classified genome groups.
///
/// `ingroup_idx` / `outgroup_idx` are column indices into the per-genome row.
/// When `ingroup_idx` is empty, no ingroup quorum is checked.
/// When `outgroup_idx` is empty, no outgroup quorum is checked.
pub struct GroupQuorumFilter {
pub ingroup_idx: Vec<usize>,
pub outgroup_idx: Vec<usize>,
pub threshold: u32,
pub min_count: usize,
pub max_count: usize,
pub min_frac: f64,
pub max_frac: f64,
pub min_outgroup_count: usize,
pub max_outgroup_count: usize,
pub min_outgroup_frac: f64,
pub max_outgroup_frac: f64,
}
impl GroupQuorumFilter {
// Build PresenceGeq/PresenceLeq constraints for one group (ingroup or outgroup).
fn group_mask_parts(
indices: &[usize],
threshold: u32,
min_count: usize,
max_count: usize,
min_frac: f64,
max_frac: f64,
parts: &mut Vec<FilterMask>,
) {
let n = indices.len();
let geq = min_count.max((min_frac * n as f64).ceil() as usize);
if geq > 0 {
parts.push(FilterMask::PresenceGeq {
indices: indices.to_vec(),
threshold,
min_count: geq,
});
}
let leq = max_count.min((max_frac * n as f64).floor() as usize);
if leq < n {
parts.push(FilterMask::PresenceLeq {
indices: indices.to_vec(),
threshold,
max_count: leq,
});
}
}
}
impl KmerFilter for GroupQuorumFilter {
fn passes(&self, _kmer: CanonicalKmer, row: &[u32], _n_genomes: usize) -> bool {
if !self.ingroup_idx.is_empty() {
let n = self.ingroup_idx.iter()
.filter(|&&i| row.get(i).copied().unwrap_or(0) > self.threshold)
.count();
let denom = self.ingroup_idx.len();
if n < self.min_count { return false; }
if n > self.max_count { return false; }
let frac = n as f64 / denom as f64;
if frac < self.min_frac { return false; }
if frac > self.max_frac { return false; }
}
if !self.outgroup_idx.is_empty() {
let n = self.outgroup_idx.iter()
.filter(|&&i| row.get(i).copied().unwrap_or(0) > self.threshold)
.count();
let denom = self.outgroup_idx.len();
if n < self.min_outgroup_count { return false; }
if n > self.max_outgroup_count { return false; }
let frac = n as f64 / denom as f64;
if frac < self.min_outgroup_frac { return false; }
if frac > self.max_outgroup_frac { return false; }
}
true
}
fn column_mask_expr(&self, _n_genomes: usize) -> Option<FilterMask> {
let t = self.threshold.checked_add(1)?;
let mut parts: Vec<FilterMask> = Vec::new();
if !self.ingroup_idx.is_empty() {
Self::group_mask_parts(
&self.ingroup_idx, t,
self.min_count, self.max_count,
self.min_frac, self.max_frac,
&mut parts,
);
}
if !self.outgroup_idx.is_empty() {
Self::group_mask_parts(
&self.outgroup_idx, t,
self.min_outgroup_count, self.max_outgroup_count,
self.min_outgroup_frac, self.max_outgroup_frac,
&mut parts,
);
}
Some(FilterMask::And(parts))
}
}
// ── Complexity filter (post-hoc, sequence-based) ──────────────────────────────
/// Reject k-mers with normalized entropy below `theta` — the same complexity
/// metric `obikmer index`'s `--theta`/`--level-max` apply *during* superkmer
/// construction (see [`obikentropy::KmerEntropy`]), applied here after the
/// fact, to k-mers already committed to a built index.
///
/// Unlike every other filter in this module, this one needs the k-mer's own
/// sequence, not its per-genome row — `column_mask_expr` is never overridden
/// (stays `None`), so this filter always forces the row-level scan path in
/// `rebuild_layer.rs` (which reconstructs the sequence from `unitigs.bin`
/// regardless, so no extra I/O beyond what filtering already requires).
pub struct MinComplexity {
pub level_max: usize,
pub theta: f64,
}
impl KmerFilter for MinComplexity {
fn passes(&self, kmer: CanonicalKmer, _row: &[u32], _n_genomes: usize) -> bool {
use obikentropy::KmerEntropy;
kmer.entropy(self.level_max) >= self.theta
}
}
+152
View File
@@ -0,0 +1,152 @@
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use tracing::debug;
use obipipeline::{
Pipeline, PipelineError, PipelineSender, SharedFlatFn, Stage, WorkerPool,
make_sink, make_source, make_transform,
throttle,
};
use obidebruinj::GraphDeBruijn;
use obikseq::CanonicalKmer;
use obilayeredmap::{IndexMode, Layer};
use obiskio::{SKError, SKResult};
use crate::common::olm_to_sk;
// ── KmerGraphData ─────────────────────────────────────────────────────────────
enum KmerGraphData {
File(obipipeline::Throttled<PathBuf>),
RawBatch(Vec<CanonicalKmer>),
FilteredBatch(Vec<CanonicalKmer>),
}
// ── build_graph ───────────────────────────────────────────────────────────────
/// Phase 1: pipeline that reads files, filters kmers, pushes into a GraphDeBruijn.
///
/// `flat_fn(path, emit)`: opens path, iterates kmers, calls `emit(batch)` for each batch.
/// `filter(kmer) -> bool`: secondary filter applied in the Transform stage.
pub(crate) fn build_graph<I, F, G>(
file_source: I,
flat_fn: F,
filter: G,
n_workers: usize,
max_open: usize,
) -> SKResult<GraphDeBruijn>
where
I: Iterator<Item = PathBuf> + Send + 'static,
F: Fn(PathBuf, &mut dyn FnMut(Vec<CanonicalKmer>)) -> SKResult<()> + Send + Sync + 'static,
G: Fn(CanonicalKmer) -> bool + Send + Sync + 'static,
{
let capacity = 2;
let flat_fn = Arc::new(flat_fn);
let filter = Arc::new(filter);
let g_shared = Arc::new(Mutex::new(GraphDeBruijn::new()));
let g_sink = Arc::clone(&g_shared);
let err_cap: Arc<Mutex<Option<SKError>>> = Arc::new(Mutex::new(None));
let err_flat = Arc::clone(&err_cap);
let throttled = throttle(file_source, max_open);
let pipeline = Pipeline::new(
make_source!(KmerGraphData, throttled, File),
vec![
Stage::Flat(Arc::new(
move |data: KmerGraphData,
push: &PipelineSender<Result<KmerGraphData, PipelineError>>,
delta: &PipelineSender<isize>|
{
if let KmerGraphData::File(t) = data {
let path = t.item;
let _guard = t.guard; // released at end of block
let mut count: isize = 0;
let push_clone = push.clone();
let result = flat_fn(path, &mut |batch: Vec<CanonicalKmer>| {
push_clone.send(Ok(KmerGraphData::RawBatch(batch))).ok();
count += 1;
});
match result {
Ok(()) => {
delta.send(count - 1).ok();
}
Err(e) => {
*err_flat.lock().unwrap() = Some(e);
delta.send(-1).ok();
}
}
}
},
) as SharedFlatFn<KmerGraphData>),
make_transform!(KmerGraphData, {
let filter = Arc::clone(&filter);
move |batch: Vec<CanonicalKmer>| -> Vec<CanonicalKmer> {
batch.into_iter().filter(|k| filter(*k)).collect()
}
}, RawBatch, FilteredBatch),
],
make_sink!(KmerGraphData, {
move |batch: Vec<CanonicalKmer>| {
let mut g = g_sink.lock().unwrap();
for kmer in batch {
g.push(kmer);
}
}
}, FilteredBatch),
);
WorkerPool::new(pipeline, n_workers, capacity).run();
if let Some(e) = Arc::try_unwrap(err_cap)
.unwrap_or_else(|_| panic!("build_graph: err_cap not uniquely owned after pipeline"))
.into_inner()
.unwrap_or_else(|e| e.into_inner())
{
return Err(e);
}
let g = Arc::try_unwrap(g_shared)
.unwrap_or_else(|_| panic!("build_graph: g_shared not uniquely owned after pipeline"))
.into_inner()
.unwrap_or_else(|e| e.into_inner());
Ok(g)
}
// ── write_graph_as_unitigs ────────────────────────────────────────────────────
/// Phase 2 (write unitigs only): compute degrees, write unitigs to `layer_dir`, drop graph.
///
/// Returns n_kmers. Does NOT build the MPHF — caller does it.
pub(crate) fn write_graph_as_unitigs(g: GraphDeBruijn, layer_dir: &Path) -> SKResult<usize> {
let n_kmers = g.len();
g.compute_degrees_and_mark_starts();
std::fs::create_dir_all(layer_dir)?;
let mut uw = Layer::<()>::unitig_writer(layer_dir).map_err(|e| olm_to_sk(e, "graph pipeline"))?;
g.try_for_each_unitig(|unitig| uw.write(unitig))?;
uw.close()?;
drop(g);
Ok(n_kmers)
}
// ── materialize_layer ─────────────────────────────────────────────────────────
/// Phase 2 (full): write_graph_as_unitigs + `Layer::<()>::build`.
///
/// Returns n_kmers.
pub(crate) fn materialize_layer(
g: GraphDeBruijn,
layer_dir: &Path,
block_bits: u8,
evidence: &IndexMode,
) -> SKResult<usize> {
let n = write_graph_as_unitigs(g, layer_dir)?;
debug!("materialize_layer: unitigs written ({n} kmers), building MPHF");
Layer::<()>::build(layer_dir, block_bits, evidence)
.map_err(|e| olm_to_sk(e, "graph pipeline"))?;
debug!("materialize_layer: MPHF build done");
Ok(n)
}
+76 -65
View File
@@ -2,13 +2,15 @@ use std::collections::BTreeMap;
use std::fs;
use std::path::{Path, PathBuf};
use obikpartitionner::{KmerPartitions, KmerSpectrum, PARTITIONS_SUBDIR};
use obikpartitionner::{KmerSpectrum, PartitionRouter};
use obilayeredmap::meta::PartitionMeta;
use obisys::{Reporter, Stage, progress_bar};
use rayon::prelude::*;
use tracing::info;
use obikseq::{set_k, set_m};
use crate::common::load_meta;
use crate::error::{OKIError, OKIResult};
use crate::meta::{GenomeInfo, IndexConfig, IndexMeta};
use crate::state::{IndexState, SENTINEL_COUNTED, SENTINEL_INDEXED, SENTINEL_SCATTERED};
@@ -16,7 +18,6 @@ use crate::state::{IndexState, SENTINEL_COUNTED, SENTINEL_INDEXED, SENTINEL_SCAT
pub struct KmerIndex {
pub(crate) root_path: PathBuf,
pub(crate) meta: IndexMeta,
pub(crate) partition: KmerPartitions,
}
impl KmerIndex {
@@ -31,13 +32,7 @@ impl KmerIndex {
force: bool,
) -> OKIResult<Self> {
let root_path = path.as_ref().to_owned();
let partition = KmerPartitions::create(
&root_path,
config.n_bits,
config.kmer_size,
config.minimizer_size,
force,
)?;
PartitionRouter::create(&root_path, config.n_bits, force)?;
set_k(config.kmer_size);
set_m(config.minimizer_size);
let mut meta = IndexMeta::new(config);
@@ -45,11 +40,7 @@ impl KmerIndex {
meta.genomes.push(info);
}
meta.write(&root_path)?;
Ok(Self {
root_path,
meta,
partition,
})
Ok(Self { root_path, meta })
}
pub fn open<P: AsRef<Path>>(path: P) -> OKIResult<Self> {
@@ -57,17 +48,7 @@ impl KmerIndex {
let meta = IndexMeta::read(&root_path).map_err(OKIError::Io)?;
set_k(meta.config.kmer_size);
set_m(meta.config.minimizer_size);
let partition = KmerPartitions::open_with_config(
&root_path,
meta.config.kmer_size,
meta.config.minimizer_size,
meta.config.n_bits,
)?;
Ok(Self {
root_path,
meta,
partition,
})
Ok(Self { root_path, meta })
}
/// Return `true` if `path` contains an `index.meta` file.
@@ -96,25 +77,17 @@ impl KmerIndex {
}
/// Lay out a fresh index skeleton at `output`: the root directory,
/// `index.meta` (from `meta`), and an opened, empty partition set.
/// `index.meta` (from `meta`), and an empty partition layout.
///
/// For construction paths that build partitions from scratch (`select`,
/// `rebuild`). `merge` bootstraps by copying a source index instead, so
/// it does not use this.
pub(crate) fn create_skeleton<P: AsRef<Path>>(
output: P,
meta: &IndexMeta,
) -> OKIResult<KmerPartitions> {
pub(crate) fn create_skeleton<P: AsRef<Path>>(output: P, meta: &IndexMeta) -> OKIResult<KmerIndex> {
let output = output.as_ref();
fs::create_dir_all(output).map_err(OKIError::Io)?;
meta.write(output).map_err(OKIError::Io)?;
fs::create_dir_all(output.join(PARTITIONS_SUBDIR)).map_err(OKIError::Io)?;
Ok(KmerPartitions::open_with_config(
output,
meta.config.kmer_size,
meta.config.minimizer_size,
meta.config.n_bits,
)?)
PartitionRouter::create(output, meta.config.n_bits, false)?;
Ok(KmerIndex { root_path: output.to_owned(), meta: meta.clone() })
}
/// Mark `output` as fully indexed, pack its column matrices, and reopen it.
@@ -139,9 +112,7 @@ impl KmerIndex {
IndexState::detect(&self.root_path).unwrap_or(IndexState::Empty)
}
/// The index's root directory — needed by out-of-crate extension code
/// (e.g. `obikphylo`) that opens its own `KmerPartition` handle onto
/// the same on-disk index.
/// The index's root directory.
pub fn root_path(&self) -> &Path {
&self.root_path
}
@@ -185,7 +156,48 @@ impl KmerIndex {
&self.meta.genomes
}
pub fn n_partitions(&self) -> usize {
self.partition.n_partitions()
1usize << self.meta.config.n_bits
}
/// Path of partition `i`'s raw directory (`partitions/part_{i:05}`) —
/// the on-disk naming convention `obikpartitionner::PartitionRouter`
/// also writes to (raw/dereplicated superkmer files, `mphf1.bin`,
/// `counts1.bin`), the single point of agreement between the two.
pub fn partition_dir(&self, i: usize) -> PathBuf {
obikpartitionner::partition_dir(&self.root_path, i)
}
/// Path of partition `i`'s layered-index directory (`<partition>/index`).
pub fn index_dir(&self, i: usize) -> PathBuf {
self.partition_dir(i).join("index")
}
/// Path of layer `l` within partition `i`'s layered index.
pub fn layer_dir(&self, i: usize, l: usize) -> PathBuf {
obilayeredmap::layer_dir(&self.index_dir(i), l)
}
/// Partition `i`'s metadata (layer count, evidence mode). Returns
/// `obiskio::SKResult`, not `OKIResult` — matches the error convention
/// of the partition/layer-construction code below (moved here from
/// `obikpartitionner`, which predates `OKIError`); `?` still converts
/// it to `OKIResult` at any call site that needs one (`OKIError: From<SKError>`).
pub fn partition_meta(&self, i: usize) -> obiskio::SKResult<PartitionMeta> {
load_meta(&self.index_dir(i), "partition_meta")
}
/// Number of layers in partition `i` — see [`partition_meta`](Self::partition_meta).
pub fn n_layers(&self, i: usize) -> obiskio::SKResult<usize> {
Ok(self.partition_meta(i)?.n_layers)
}
/// Evidence mode partition `i` was actually built with — see
/// [`partition_meta`](Self::partition_meta). Distinct from
/// [`evidence_mode`](Self::evidence_mode): that one is the index-level
/// config, this one is per-partition ground truth (the two agree in
/// practice, but this is what `Layer::open` needs).
pub fn partition_mode(&self, i: usize) -> obiskio::SKResult<obilayeredmap::IndexMode> {
Ok(self.partition_meta(i)?.mode)
}
/// Number of layers per partition.
@@ -194,13 +206,7 @@ impl KmerIndex {
/// homogeneous across all partitions — reading it off partition 0
/// is enough, no need to scan every partition.
pub fn n_layers_per_partition(&self) -> OKIResult<usize> {
Ok(self.partition.n_layers(0)?)
}
/// Expose the inner partition so the caller can run scatter into it.
/// Call `mark_scattered` once scatter is complete.
pub fn partition_mut(&mut self) -> &mut KmerPartitions {
&mut self.partition
Ok(self.n_layers(0)?)
}
/// Mark scatter as complete and write `scatter.done`.
@@ -217,6 +223,14 @@ impl KmerIndex {
Ok(())
}
/// Open a fresh [`PartitionRouter`] onto this index's partition layout
/// — the write-side handle for `scatter`, or for `dereplicate_and_count`
/// below. Transient: no state is kept in `KmerIndex` itself between
/// calls, only on disk.
pub fn partition_router(&self) -> OKIResult<PartitionRouter> {
Ok(PartitionRouter::open(&self.root_path, self.meta.config.n_bits)?)
}
/// Dereplicate all partitions then compute kmer counts.
///
/// Writes `spectrums/{label}.json` and touches `count.done` upon completion.
@@ -226,12 +240,14 @@ impl KmerIndex {
keep_intermediate: bool,
rep: &mut Reporter,
) -> OKIResult<()> {
let router = self.partition_router()?;
let t = Stage::start("dereplicate");
self.partition.dereplicate()?;
router.dereplicate()?;
rep.push(t.stop());
let t = Stage::start("count_kmer");
let spectrum = self.partition.count_kmer(keep_intermediate)?;
let spectrum = router.count_kmer(keep_intermediate)?;
rep.push(t.stop());
self.write_spectrum(&spectrum)?;
@@ -273,7 +289,7 @@ impl KmerIndex {
keep_intermediate: bool,
rep: &mut Reporter,
) -> OKIResult<()> {
let n = self.partition.n_partitions();
let n = self.n_partitions();
let t = Stage::start("index");
let with_counts = self.meta.config.with_counts;
let evidence = self.meta.config.evidence.clone();
@@ -287,7 +303,7 @@ impl KmerIndex {
.run(
&order,
|i| {
self.partition.build_index_layer(
self.build_index_layer(
i,
min_ab,
max_ab,
@@ -311,7 +327,7 @@ impl KmerIndex {
if !keep_intermediate {
for i in 0..n {
self.partition.remove_build_artifacts(i);
self.remove_build_artifacts(i);
}
}
@@ -320,14 +336,9 @@ impl KmerIndex {
Ok(())
}
/// Borrow the inner partition for direct superkmer-level queries.
pub fn partition(&self) -> &KmerPartitions {
&self.partition
}
/// Path to the unitigs file for partition `part`, layer `layer`.
pub fn layer_unitigs_path(&self, part: usize, layer: usize) -> PathBuf {
self.partition.layer_dir(part, layer).join("unitigs.bin")
self.layer_dir(part, layer).join("unitigs.bin")
}
/// Pack all partition matrices into single-file format (presence → .pbmx, counts → .pcmx).
@@ -349,13 +360,13 @@ impl KmerIndex {
crate::numa::PartitionRunner::new().run(
&order,
|i| -> OKIResult<()> {
let index_dir = self.partition.index_dir(i);
let index_dir = self.index_dir(i);
if !index_dir.exists() {
return Ok(());
}
let n_layers = self.partition.n_layers(i)?;
let n_layers = self.n_layers(i)?;
for l in 0..n_layers {
let layer_dir = self.partition.layer_dir(i, l);
let layer_dir = self.layer_dir(i, l);
let presence_dir = layer_dir.join("presence");
let counts_dir = layer_dir.join("counts");
if presence_dir.exists() {
@@ -391,11 +402,11 @@ impl KmerIndex {
let errors: Vec<_> = (0..n)
.into_par_iter()
.filter_map(|i| {
let index_dir = self.partition.index_dir(i);
let index_dir = self.index_dir(i);
if !index_dir.exists() {
return None;
}
let n_layers = match self.partition.n_layers(i) {
let n_layers = match self.n_layers(i) {
Ok(n) => n,
Err(e) => {
return Some(OKIError::Io(std::io::Error::new(
@@ -405,7 +416,7 @@ impl KmerIndex {
}
};
for l in 0..n_layers {
let layer_dir = self.partition.layer_dir(i, l);
let layer_dir = self.layer_dir(i, l);
let meta_path = layer_dir.join(LayerMeta::FILENAME);
if meta_path.exists() {
continue;
+133
View File
@@ -0,0 +1,133 @@
use std::fs;
use std::io;
use cacheline_ef::{CachelineEf, CachelineEfVec};
use epserde::prelude::*;
use obicompactvec::{PersistentCompactIntMatrix, PersistentCompactIntVec};
use obidebruinj::GraphDeBruijn;
use obilayeredmap::meta::PartitionMeta;
use obilayeredmap::{IndexMode, layer::Layer};
use obiskio::{SKError, SKFileMeta, SKFileReader};
use ptr_hash::{PtrHash, bucket_fn::CubicEps, hash::Xx64};
use crate::common::olm_to_sk;
use crate::graph_pipeline::{materialize_layer, write_graph_as_unitigs};
use crate::index::KmerIndex;
type Mphf = PtrHash<u64, CubicEps, CachelineEfVec<Vec<CachelineEf>>, Xx64, Vec<u8>>;
fn remove_if_exists(path: &std::path::Path) {
if let Err(e) = fs::remove_file(path) {
if e.kind() != io::ErrorKind::NotFound {
eprintln!("warning: could not remove {}: {e}", path.display());
}
}
}
impl KmerIndex {
/// Build the layered MPHF index for partition `i`.
///
/// Returns the number of canonical k-mers indexed, or 0 if the partition
/// has no data or its layer was already built (resume-safe).
///
/// Abundance filtering is applied when `min_ab > 1` or `max_ab.is_some()`,
/// using `mphf1.bin` + `counts1.bin` if they exist.
/// Count payload is stored iff `with_counts` is true.
pub fn build_index_layer(
&self,
i: usize,
min_ab: u32,
max_ab: Option<u32>,
with_counts: bool,
mode: &IndexMode,
block_bits: u8,
) -> Result<usize, SKError> {
let partition_dir = self.partition_dir(i);
let dedup_path = partition_dir.join("dereplicated.skmer.zst");
if !dedup_path.exists() {
return Ok(0);
}
let layer_dir = self.layer_dir(i, 0);
if layer_dir.join("mphf.bin").exists() {
return Ok(0);
}
let filter_active = min_ab > 1 || max_ab.is_some();
let need_counts = filter_active || with_counts;
let mphf1_opt: Option<Mphf> = if need_counts {
let p = partition_dir.join("mphf1.bin");
p.exists().then(|| Mphf::load_full(&p).ok()).flatten()
} else {
None
};
let counts1_opt: Option<PersistentCompactIntVec> = if need_counts {
let p = partition_dir.join("counts1.bin");
p.exists()
.then(|| PersistentCompactIntVec::open(&p).ok())
.flatten()
} else {
None
};
let mut g = GraphDeBruijn::new();
let mut reader = SKFileReader::open(&dedup_path)?;
for sk in reader.iter() {
for kmer in sk.iter_canonical_kmers() {
let accept = if filter_active {
match (&mphf1_opt, &counts1_opt) {
(Some(mphf), Some(counts)) => {
let ab = counts.get(mphf.index(&kmer.raw()));
ab >= min_ab && max_ab.map_or(true, |max| ab <= max)
}
_ => true,
}
} else {
true
};
if accept {
g.push(kmer);
}
}
}
let n_kmers =
if with_counts {
let n = write_graph_as_unitigs(g, &layer_dir)?;
Layer::<PersistentCompactIntMatrix>::build(&layer_dir, block_bits, mode, |kmer| {
match (&mphf1_opt, &counts1_opt) {
(Some(mphf), Some(counts)) => counts.get(mphf.index(&kmer.raw())),
_ => 1,
}
})
.map_err(|e| olm_to_sk(e, "layer build"))?;
n
} else {
materialize_layer(g, &layer_dir, block_bits, mode)?
};
let index_dir = layer_dir.parent().expect("layer_dir has a parent");
PartitionMeta {
n_layers: 1,
mode: mode.clone(),
}
.save(index_dir)
.map_err(|e| olm_to_sk(e, "layer build"))?;
Ok(n_kmers)
}
/// Remove intermediate build artifacts for partition `i`.
///
/// Deletes `dereplicated.skmer.zst` (+ sidecar), `mphf1.bin`, `counts1.bin`.
pub fn remove_build_artifacts(&self, i: usize) {
let partition_dir = self.partition_dir(i);
let dedup = partition_dir.join("dereplicated.skmer.zst");
remove_if_exists(&SKFileMeta::sidecar_path(&dedup));
remove_if_exists(&dedup);
remove_if_exists(&partition_dir.join("mphf1.bin"));
remove_if_exists(&partition_dir.join("counts1.bin"));
}
}
+14 -1
View File
@@ -2,22 +2,35 @@ pub mod error;
pub mod meta;
pub mod predicate;
pub mod state;
mod common;
mod distance;
mod dump;
mod dump_layer;
pub mod filter;
mod graph_pipeline;
mod index;
mod index_layer;
mod matrix_store;
mod merge;
mod merge_layer;
mod numa;
mod query_layer;
mod rebuild;
mod rebuild_layer;
mod reindex;
mod select;
mod select_layer;
mod stats;
pub use error::{OKIError, OKIResult};
pub use distance::{DistanceMetric, DistanceOutput};
pub use filter::{GroupQuorumFilter, KmerFilter, passes_all};
pub use index::KmerIndex;
pub use merge::MergeMode;
pub use merge_layer::MergeMode;
pub use meta::{validate_label, GenomeInfo, IndexConfig, IndexMeta, META_FILENAME};
pub use predicate::{GroupFilterParams, MetaPred};
pub use query_layer::{KmerDesc, QueryHit, QueryStats};
pub use select_layer::{AggOp, OutputCol};
pub use state::{IndexState, SENTINEL_COUNTED, SENTINEL_INDEXED, SENTINEL_SCATTERED};
pub use stats::IndexBitsPerKmer;
pub use numa::PartitionRunner;
+46
View File
@@ -0,0 +1,46 @@
use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix};
use obilayeredmap::{LayeredStore, open_data};
use obiskio::SKResult;
use crate::common::{load_meta, olm_to_sk};
use crate::index::KmerIndex;
impl KmerIndex {
/// Open all count matrices for partition `part`, one per layer.
/// Layers without a `counts/` directory are skipped.
pub fn count_store(&self, part: usize) -> SKResult<LayeredStore<PersistentCompactIntMatrix>> {
let index_dir = self.index_dir(part);
if !index_dir.exists() {
return Ok(LayeredStore::new(vec![]));
}
let n_layers = load_meta(&index_dir, "distance")?.n_layers;
let matrices = (0..n_layers)
.filter_map(|l| {
self.layer_dir(part, l)
.join("counts")
.exists()
.then(|| open_data(&index_dir, l).map_err(|e| olm_to_sk(e, "distance")))
})
.collect::<SKResult<Vec<_>>>()?;
Ok(LayeredStore::new(matrices))
}
/// Open all presence matrices for partition `part`, one per layer.
/// Layers without a `presence/` directory are skipped.
pub fn presence_store(&self, part: usize) -> SKResult<LayeredStore<PersistentBitMatrix>> {
let index_dir = self.index_dir(part);
if !index_dir.exists() {
return Ok(LayeredStore::new(vec![]));
}
let n_layers = load_meta(&index_dir, "distance")?.n_layers;
let matrices = (0..n_layers)
.filter_map(|l| {
self.layer_dir(part, l)
.join("presence")
.exists()
.then(|| open_data(&index_dir, l).map_err(|e| olm_to_sk(e, "distance")))
})
.collect::<SKResult<Vec<_>>>()?;
Ok(LayeredStore::new(matrices))
}
}
+5 -6
View File
@@ -13,7 +13,7 @@ use crate::index::KmerIndex;
use crate::meta::{GenomeInfo, IndexMeta};
use crate::state::{IndexState, SENTINEL_INDEXED};
pub use obikpartitionner::MergeMode;
pub use crate::merge_layer::MergeMode;
// ── per-partition diagnostic record ──────────────────────────────────────────
@@ -189,13 +189,12 @@ impl KmerIndex {
let t = Stage::start("merge_partitions");
let pb = progress_bar("merge", n_partitions as u64, "partitions");
let dst_partition = &dst.partition;
let block_bits = dst.meta.config.block_bits;
// Pre-build source list once (avoid rebuilding per partition)
let srcs: Vec<(&obikpartitionner::KmerPartitions, usize)> = remaining_sources
let srcs: Vec<(&KmerIndex, usize)> = remaining_sources
.iter()
.map(|s| (&s.partition, s.meta.genomes.len()))
.map(|s| (*s, s.meta.genomes.len()))
.collect();
// Per-partition unitig byte sizes across remaining sources (stat() only)
@@ -225,7 +224,7 @@ impl KmerIndex {
.run(
&order,
|i| {
dst_partition.merge_partition(
dst.merge_partition(
i,
srcs,
mode,
@@ -418,7 +417,7 @@ fn is_trivial(src: &KmerIndex, mode: MergeMode) -> bool {
}
fn index_unitig_size(src: &KmerIndex) -> u64 {
let n = src.partition.n_partitions();
let n = src.n_partitions();
(0..n).map(|i| partition_unitig_bytes(src, i)).sum()
}
+598
View File
@@ -0,0 +1,598 @@
//! Merging a source partition's new layer into a destination partition:
//! de Bruijn graph union (pass 1) then column fill (pass 2).
//!
//! Submodules: [`src_layer`] (`SrcLayerData`, the opened-source-matrix
//! lookup used by pass 2 here and by `rebuild_layer`). The `merge_partition`
//! orchestration itself stays in this file — its ~400-line body is one
//! tightly threaded pipeline (shared `Arc`/`Mutex` state across pass 1,
//! builder setup, and pass 2), not a set of independently callable steps.
use std::fs;
use std::io;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use obipipeline::{
Pipeline, PipelineError, PipelineSender, SharedFlatFn, Stage, ThrottleGuard, WorkerPool,
make_sink, make_source, make_transform, throttle,
};
use tracing::debug;
use obicompactvec::{PersistentBitMatrixBuilder, PersistentCompactIntMatrixBuilder};
use obikseq::CanonicalKmer;
use obilayeredmap::{IndexMode, Layer, LayeredMap, MphfOnly, layer_dir};
use obiskio::{SKError, SKResult, UnitigFileReader};
use crate::common::{ColBuilder, load_meta, olm_to_sk};
use crate::graph_pipeline::{build_graph, materialize_layer};
use crate::index::KmerIndex;
mod src_layer;
pub(crate) use src_layer::SrcLayerData;
// ── MergeMode ─────────────────────────────────────────────────────────────────
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MergeMode {
Presence,
Count,
}
// ── MatrixBuilder ─────────────────────────────────────────────────────────────
//
// Wraps whichever matrix builder `mode` calls for, so the merge pipeline never
// has to know the on-disk column naming (`col_NNNNNN.pbiv`/`.pciv`) or the
// matrix `meta.json` schema itself — both stay private to obicompactvec.
// `resume` reopens a matrix directory already closed by a previous builder
// session (an existing destination layer), continuing from its current
// `n_cols` instead of starting a fresh matrix at 0.
enum MatrixBuilder {
Bit(PersistentBitMatrixBuilder),
Int(PersistentCompactIntMatrixBuilder),
}
impl MatrixBuilder {
fn new(mode: MergeMode, n: usize, dir: &Path) -> io::Result<Self> {
Ok(match mode {
MergeMode::Presence => MatrixBuilder::Bit(PersistentBitMatrixBuilder::new(n, dir)?),
MergeMode::Count => MatrixBuilder::Int(PersistentCompactIntMatrixBuilder::new(n, dir)?),
})
}
fn resume(mode: MergeMode, dir: &Path) -> io::Result<Self> {
Ok(match mode {
MergeMode::Presence => MatrixBuilder::Bit(PersistentBitMatrixBuilder::resume(dir)?),
MergeMode::Count => MatrixBuilder::Int(PersistentCompactIntMatrixBuilder::resume(dir)?),
})
}
/// Add a column with no data written (all-zero/false) — for genome
/// columns absent from this source (e.g. dst genomes in a new layer).
fn add_absent_col(&mut self) -> io::Result<()> {
match self {
MatrixBuilder::Bit(b) => b.add_col()?.close(),
MatrixBuilder::Int(b) => b.add_col()?.close(),
}
}
fn add_col(&mut self) -> io::Result<ColBuilder> {
Ok(match self {
MatrixBuilder::Bit(b) => ColBuilder::Bit(b.add_col()?),
MatrixBuilder::Int(b) => ColBuilder::Int(b.add_col()?),
})
}
fn close(self) -> io::Result<()> {
match self {
MatrixBuilder::Bit(b) => b.close(),
MatrixBuilder::Int(b) => b.close(),
}
}
}
#[cfg(test)]
mod matrix_builder_tests {
use tempfile::tempdir;
use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix};
use super::{ColBuilder, MatrixBuilder, MergeMode};
/// Mirrors `merge_partition`'s "new layer" setup: absent (dst-genome)
/// columns get no data, then source columns are filled — all through one
/// continuous `MatrixBuilder` session, closed once at the end.
#[test]
fn new_layer_absent_then_source_columns_presence() {
let dir = tempdir().unwrap();
let data_dir = dir.path().join("presence");
let mut mb = MatrixBuilder::new(MergeMode::Presence, 3, &data_dir).unwrap();
// Two absent (dst-genome) columns.
mb.add_absent_col().unwrap();
mb.add_absent_col().unwrap();
// One source column, filled like pass 2 would.
let mut col = mb.add_col().unwrap();
match &mut col {
ColBuilder::Bit(b) => {
b.set(0, true);
b.set(1, false);
b.set(2, true);
}
ColBuilder::Int(_) => unreachable!(),
}
col.close().unwrap();
mb.close().unwrap();
let m = PersistentBitMatrix::open(dir.path()).unwrap();
assert_eq!(m.n_cols(), 3);
assert_eq!(&*m.row(0), &[false, false, true]);
assert_eq!(&*m.row(1), &[false, false, false]);
assert_eq!(&*m.row(2), &[false, false, true]);
}
#[test]
fn new_layer_absent_then_source_columns_count() {
let dir = tempdir().unwrap();
let data_dir = dir.path().join("counts");
let mut mb = MatrixBuilder::new(MergeMode::Count, 2, &data_dir).unwrap();
mb.add_absent_col().unwrap();
let mut col = mb.add_col().unwrap();
match &mut col {
ColBuilder::Int(b) => {
b.set(0, 7);
b.set(1, 42);
}
ColBuilder::Bit(_) => unreachable!(),
}
col.close().unwrap();
mb.close().unwrap();
let m = PersistentCompactIntMatrix::open(dir.path()).unwrap();
assert_eq!(m.n_cols(), 2);
assert_eq!(&*m.row(0), &[0u32, 7]);
assert_eq!(&*m.row(1), &[0u32, 42]);
}
/// Mirrors `merge_partition`'s "existing dst layer" setup: an
/// already-closed matrix (from a previous merge) gets more columns
/// appended via `resume`, without callers ever seeing `col_path`/
/// `MatrixMeta` — `n` itself is read back from the matrix directory.
#[test]
fn resume_appends_source_columns_to_existing_layer() {
let dir = tempdir().unwrap();
let data_dir = dir.path().join("presence");
// Previous merge: one dst-genome column already on disk.
let mut mb0 = MatrixBuilder::new(MergeMode::Presence, 3, &data_dir).unwrap();
mb0.add_absent_col().unwrap();
mb0.close().unwrap();
// This merge: resume and append two more source columns.
let mut mb = MatrixBuilder::resume(MergeMode::Presence, &data_dir).unwrap();
for vals in [[true, false, true], [false, true, false]] {
let mut col = mb.add_col().unwrap();
match &mut col {
ColBuilder::Bit(b) => {
for (slot, v) in vals.into_iter().enumerate() {
b.set(slot, v);
}
}
ColBuilder::Int(_) => unreachable!(),
}
col.close().unwrap();
}
mb.close().unwrap();
let m = PersistentBitMatrix::open(dir.path()).unwrap();
assert_eq!(m.n_cols(), 3);
assert_eq!(&*m.row(0), &[false, true, false]);
assert_eq!(&*m.row(1), &[false, false, true]);
assert_eq!(&*m.row(2), &[false, true, false]);
}
}
// ── KmerPartition::merge_partition ────────────────────────────────────────────
impl KmerIndex {
/// Merge `sources` into destination partition `i`.
///
/// Each entry in `sources` is `(partition, n_genomes)` where `n_genomes` is
/// the number of genome columns that source contributes. A merged index
/// contributes more than one. The total new columns added to the destination
/// is `sum(n_genomes)`.
///
/// `n_dst_genomes` is the number of genome columns already in the destination
/// matrices (copied from source_0 before this call).
pub fn merge_partition(
&self,
i: usize,
sources: &[(&KmerIndex, usize)],
mode: MergeMode,
n_dst_genomes: usize,
block_bits: u8,
evidence: &IndexMode,
) -> SKResult<usize> {
let dst_index_dir = self.index_dir(i);
if !dst_index_dir.exists() {
return Ok(0);
}
load_meta(&dst_index_dir, "merge")?; // ensure meta.json exists before LayeredMap::open
let dst_map =
Arc::new(LayeredMap::<()>::open(&dst_index_dir).map_err(|e| olm_to_sk(e, "merge"))?);
let n_dst_layers = dst_map.n_layers();
let n_src_total: usize = sources.iter().map(|(_, n)| *n).sum();
// First merge in presence mode: init presence matrices on existing layers
// (all slots true — every kmer in those layers belongs to genome_0).
if n_dst_genomes == 1 && mode == MergeMode::Presence {
for l in 0..n_dst_layers {
Layer::<()>::init_presence_matrix(
&layer_dir(&dst_index_dir, l),
dst_map.layer(l).n(),
)
.map_err(|e| olm_to_sk(e, "merge"))?;
}
}
// ── Pass 1: pipeline — parallel file read + dst_map filter + graph fill ─
//
// Source : list of unitigs.bin paths (one per source × layer)
// Flat : open file, emit Vec<CanonicalKmer> batches (BeeGFS parallel I/O)
// Transform: filter via dst_map.query() — thread-safe, LayeredMap<()>: Sync
// Sink : push new kmers into GraphDeBruijn (single thread, no locks needed)
// Collect file paths (propagates load_meta errors before the pipeline starts)
let mut unitig_paths: Vec<PathBuf> = Vec::new();
for (src, _) in sources.iter() {
let src_index_dir = src.index_dir(i);
if !src_index_dir.exists() {
continue;
}
let src_meta = load_meta(&src_index_dir, "merge")?;
for l in 0..src_meta.n_layers {
let p = layer_dir(&src_index_dir, l).join("unitigs.bin");
if p.exists() {
unitig_paths.push(p);
}
}
}
let n_src_layers = unitig_paths.len();
debug!("partition {i}: de Bruijn graph build start — {n_src_layers} source layer(s)");
const BATCH: usize = 4096;
let n_workers = rayon::current_num_threads().min(16).max(4);
// At most 2 files open simultaneously: keeps n_workers-2 workers free
// for the Transform stage. Each open file monopolises one worker for the
// full duration of its read, so this must stay well below n_workers.
let max_open = 2;
let dst_filter = Arc::clone(&dst_map);
let g = build_graph(
unitig_paths.into_iter(),
move |path: PathBuf, emit: &mut dyn FnMut(Vec<CanonicalKmer>)| -> SKResult<()> {
let reader = UnitigFileReader::open_sequential(&path)?;
let mut batch: Vec<CanonicalKmer> = Vec::with_capacity(BATCH);
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
batch.push(kmer);
if batch.len() == BATCH {
emit(std::mem::replace(&mut batch, Vec::with_capacity(BATCH)));
}
}
if !batch.is_empty() {
emit(batch);
}
Ok(())
},
move |kmer| dst_filter.query(kmer).is_none(),
n_workers,
max_open,
)?;
let any_new = g.len() > 0;
debug!(
"partition {i}: de Bruijn graph done — {} new kmers",
g.len()
);
// Build new layer from de Bruijn graph if there are new kmers.
let new_layer_idx = n_dst_layers;
let new_layer_dir = layer_dir(&dst_index_dir, new_layer_idx);
let n_new = if any_new {
debug!("partition {i}: unitig traversal start — {} nodes", g.len());
let n_nodes = materialize_layer(g, &new_layer_dir, block_bits, evidence)?;
debug!("partition {i}: MPHF build done");
n_nodes
} else {
drop(g);
0
};
let t_open = std::time::Instant::now();
let new_mphf: Option<Arc<MphfOnly>> = if any_new {
Some(Arc::new(
MphfOnly::open(&new_layer_dir).map_err(|e| olm_to_sk(e, "merge"))?,
))
} else {
None
};
debug!(
"partition {i}: MPHF open in {:.3}s",
t_open.elapsed().as_secs_f64()
);
// ── Prepare matrix directories for the new layer ──────────────────────
// Absent columns (dst genomes) get an all-zero/false column. Source-genome
// columns are created as mutable builders for pass 2. `new_mb` is kept
// open (not `close`d) until pass 2 has filled every source column, so its
// `n_cols` bookkeeping stays in sync with the columns actually created.
let (new_src_builders, new_mb): (Vec<ColBuilder>, Option<MatrixBuilder>) = if any_new {
let data_dir = match mode {
MergeMode::Presence => new_layer_dir.join("presence"),
MergeMode::Count => new_layer_dir.join("counts"),
};
fs::create_dir_all(&data_dir)?;
let mut mb = MatrixBuilder::new(mode, n_new, &data_dir).map_err(SKError::Io)?;
for _ in 0..n_dst_genomes {
mb.add_absent_col().map_err(SKError::Io)?;
}
let cols = (0..n_src_total)
.map(|_| mb.add_col().map_err(SKError::Io))
.collect::<SKResult<Vec<_>>>()?;
(cols, Some(mb))
} else {
(vec![], None)
};
let t_builders = std::time::Instant::now();
// Builders for existing layers: n_src_total per layer, resumed from
// each layer's own matrix directory (already holding n_dst_genomes
// columns from a previous merge). Columns land at
// n_dst_genomes .. n_dst_genomes + n_src_total - 1.
let mut exist_mbs: Vec<MatrixBuilder> = Vec::with_capacity(n_dst_layers);
let mut exist_builders: Vec<Vec<ColBuilder>> = Vec::with_capacity(n_dst_layers);
for l in 0..n_dst_layers {
let layer_dir = layer_dir(&dst_index_dir, l);
let data_dir = match mode {
MergeMode::Presence => layer_dir.join("presence"),
MergeMode::Count => layer_dir.join("counts"),
};
let mut mb = MatrixBuilder::resume(mode, &data_dir).map_err(SKError::Io)?;
let cols = (0..n_src_total)
.map(|_| mb.add_col().map_err(SKError::Io))
.collect::<SKResult<Vec<_>>>()?;
exist_mbs.push(mb);
exist_builders.push(cols);
}
debug!(
"partition {i}: builders ready in {:.3}s",
t_builders.elapsed().as_secs_f64()
);
// ── Pass 2: fill builders (pipeline) ─────────────────────────────────
let t_pass2 = std::time::Instant::now();
// Collect source items before the pipeline so load_meta errors propagate
// via ? before any worker thread is spawned.
let mut pass2_items: Vec<(usize, usize, PathBuf)> = Vec::new();
{
let mut col_offset = 0usize;
for (src, src_n) in sources.iter() {
let src_index_dir = src.index_dir(i);
if !src_index_dir.exists() {
col_offset += src_n;
continue;
}
let src_meta = load_meta(&src_index_dir, "merge")?;
for l in 0..src_meta.n_layers {
let src_layer_dir = layer_dir(&src_index_dir, l);
if src_layer_dir.join("unitigs.bin").exists() {
pass2_items.push((col_offset, *src_n, src_layer_dir));
}
}
col_offset += src_n;
}
}
enum Pass2Data {
SrcLayer((usize, usize, PathBuf, ThrottleGuard)),
RawBatch((usize, usize, Arc<SrcLayerData>, Vec<CanonicalKmer>)),
WriteBatch(Vec<(Option<usize>, usize, usize, u32)>),
}
let exist_locked: Vec<Vec<Arc<Mutex<ColBuilder>>>> = exist_builders
.into_iter()
.map(|layer| layer.into_iter().map(|b| Arc::new(Mutex::new(b))).collect())
.collect();
let new_locked: Vec<Arc<Mutex<ColBuilder>>> = new_src_builders
.into_iter()
.map(|b| Arc::new(Mutex::new(b)))
.collect();
let exist_sink: Vec<Vec<Arc<Mutex<ColBuilder>>>> = exist_locked
.iter()
.map(|layer| layer.iter().map(Arc::clone).collect())
.collect();
let new_sink: Vec<Arc<Mutex<ColBuilder>>> = new_locked.iter().map(Arc::clone).collect();
let dst_map_t2 = Arc::clone(&dst_map);
let new_mphf_t2 = new_mphf.clone();
let pass2_err: Arc<Mutex<Option<String>>> = Arc::new(Mutex::new(None));
let err_cap2 = Arc::clone(&pass2_err);
let capacity = 2;
let throttled_pass2 = throttle(pass2_items.into_iter(), max_open);
let pipeline2 = Pipeline::new(
make_source!(
Pass2Data,
throttled_pass2.map(|t| {
let (col_offset, src_n, src_layer_dir) = t.item;
(col_offset, src_n, src_layer_dir, t.guard)
}),
SrcLayer
),
vec![
Stage::Flat(Arc::new(
move |data: Pass2Data,
push: &PipelineSender<Result<Pass2Data, PipelineError>>,
delta: &PipelineSender<isize>| {
if let Pass2Data::SrcLayer((col_offset, src_n, src_layer_dir, _guard)) =
data
{
// _guard dropped at end of block, releasing the slot.
let reader = match UnitigFileReader::open_sequential(
&src_layer_dir.join("unitigs.bin"),
) {
Ok(r) => r,
Err(e) => {
*err_cap2.lock().unwrap() = Some(e.to_string());
delta.send(-1).ok();
return;
}
};
let src_data = match SrcLayerData::open(&src_layer_dir, mode) {
Ok(d) => Arc::new(d),
Err(e) => {
*err_cap2.lock().unwrap() = Some(e.to_string());
delta.send(-1).ok();
return;
}
};
const BATCH: usize = 4096;
let mut batch: Vec<CanonicalKmer> = Vec::with_capacity(BATCH);
let mut count: isize = 0;
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
batch.push(kmer);
if batch.len() == BATCH {
let b =
std::mem::replace(&mut batch, Vec::with_capacity(BATCH));
push.send(Ok(Pass2Data::RawBatch((
col_offset,
src_n,
Arc::clone(&src_data),
b,
))))
.ok();
count += 1;
}
}
if !batch.is_empty() {
push.send(Ok(Pass2Data::RawBatch((
col_offset, src_n, src_data, batch,
))))
.ok();
count += 1;
}
delta.send(count - 1).ok();
}
},
) as SharedFlatFn<Pass2Data>),
make_transform!(
Pass2Data,
{
move |(col_offset, src_n, src_data, kmers): (
usize,
usize,
Arc<SrcLayerData>,
Vec<CanonicalKmer>,
)|
-> Vec<(Option<usize>, usize, usize, u32)> {
let mut ops: Vec<(Option<usize>, usize, usize, u32)> = Vec::new();
for kmer in kmers {
let values = src_data.lookup(kmer, src_n);
if let Some((dst_layer, hit)) = dst_map_t2.query(kmer) {
for (g, val) in values.into_iter().enumerate() {
ops.push((Some(dst_layer), col_offset + g, hit.slot, val));
}
} else if let Some(ref mphf) = new_mphf_t2 {
let slot = mphf.index(kmer);
for (g, val) in values.into_iter().enumerate() {
ops.push((None, col_offset + g, slot, val));
}
}
}
ops
}
},
RawBatch,
WriteBatch
),
],
make_sink!(
Pass2Data,
{
move |ops: Vec<(Option<usize>, usize, usize, u32)>| {
for (layer_opt, col, slot, val) in ops {
match layer_opt {
Some(l) => exist_sink[l][col].lock().unwrap().set_val(slot, val),
None => new_sink[col].lock().unwrap().set_val(slot, val),
}
}
}
},
WriteBatch
),
);
WorkerPool::new(pipeline2, n_workers, capacity).run();
debug!(
"partition {i}: pass2 pipeline done in {:.3}s",
t_pass2.elapsed().as_secs_f64()
);
if let Some(msg) = Arc::try_unwrap(pass2_err)
.unwrap_or_else(|_| panic!("pass2: pass2_err not uniquely owned"))
.into_inner()
.unwrap_or_else(|e| e.into_inner())
{
return Err(SKError::InvalidData {
context: "merge pass2",
detail: msg,
});
}
let t_close = std::time::Instant::now();
// ── Close builders and update metadata ────────────────────────────────
for (mb, builders) in exist_mbs.into_iter().zip(exist_locked.into_iter()) {
for b in builders {
Arc::try_unwrap(b)
.unwrap_or_else(|_| panic!("pass2: exist_builder not uniquely owned"))
.into_inner()
.unwrap_or_else(|e| e.into_inner())
.close()?;
}
mb.close().map_err(SKError::Io)?;
}
for b in new_locked {
Arc::try_unwrap(b)
.unwrap_or_else(|_| panic!("pass2: new_builder not uniquely owned"))
.into_inner()
.unwrap_or_else(|e| e.into_inner())
.close()?;
}
if let Some(mb) = new_mb {
mb.close().map_err(SKError::Io)?;
let mut part_meta = self.partition_meta(i)?;
part_meta.n_layers = new_layer_idx + 1;
part_meta
.save(&dst_index_dir)
.map_err(|e| olm_to_sk(e, "merge"))?;
}
debug!(
"partition {i}: builders closed in {:.3}s",
t_close.elapsed().as_secs_f64()
);
Ok(n_new)
}
}
@@ -0,0 +1,95 @@
use std::path::Path;
use obicompactvec::{MatrixGroupOps, PersistentBitMatrix, PersistentCompactIntMatrix};
use obikseq::CanonicalKmer;
use obilayeredmap::MphfOnly;
use obiskio::{SKError, SKResult};
use crate::common::olm_to_sk;
use super::MergeMode;
// ── SrcLayerData — opened source matrix for pass-2 lookup ─────────────────────
pub(crate) enum SrcLayerData {
Presence(MphfOnly, PersistentBitMatrix),
Count(MphfOnly, PersistentCompactIntMatrix),
}
impl SrcLayerData {
pub(crate) fn open(layer_dir: &Path, merge_mode: MergeMode) -> SKResult<Self> {
let counts_dir = layer_dir.join("counts");
match merge_mode {
MergeMode::Presence => {
if counts_dir.exists() && !layer_dir.join("presence").exists() {
let mphf = MphfOnly::open(layer_dir).map_err(|e| olm_to_sk(e, "merge"))?;
let mat = PersistentCompactIntMatrix::open(layer_dir).map_err(SKError::Io)?;
Ok(SrcLayerData::Count(mphf, mat))
} else {
// presence dir exists, or neither exists → Implicit handled by open()
let mphf = MphfOnly::open(layer_dir).map_err(|e| olm_to_sk(e, "merge"))?;
let mat = PersistentBitMatrix::open(layer_dir).map_err(SKError::Io)?;
Ok(SrcLayerData::Presence(mphf, mat))
}
}
MergeMode::Count => {
let mphf = MphfOnly::open(layer_dir).map_err(|e| olm_to_sk(e, "merge"))?;
if counts_dir.exists() {
let mat = PersistentCompactIntMatrix::open(layer_dir).map_err(SKError::Io)?;
Ok(SrcLayerData::Count(mphf, mat))
} else {
// No counts → treat as implicit presence (all 1s)
let mat = PersistentBitMatrix::open(layer_dir).map_err(SKError::Io)?;
Ok(SrcLayerData::Presence(mphf, mat))
}
}
}
}
/// Return one value per source genome for `kmer`.
/// The caller guarantees `kmer` is in the source MPHF domain.
#[inline]
pub(crate) fn lookup(&self, kmer: CanonicalKmer, n_genomes: usize) -> Vec<u32> {
let mut buf = vec![0u32; n_genomes];
match self {
SrcLayerData::Presence(mphf, mat) => mat.fill_row(mphf.index(kmer), &mut buf),
SrcLayerData::Count(mphf, mat) => mat.fill_row(mphf.index(kmer), &mut buf),
}
buf
}
pub(crate) fn n_slots(&self) -> usize {
match self {
SrcLayerData::Presence(_, mat) => mat.n(),
SrcLayerData::Count(_, mat) => mat.n(),
}
}
/// MPHF lookup: returns the slot index for `kmer` (kmer must be in the domain).
#[inline]
pub(crate) fn slot(&self, kmer: CanonicalKmer) -> usize {
match self {
SrcLayerData::Presence(mphf, _) => mphf.index(kmer),
SrcLayerData::Count(mphf, _) => mphf.index(kmer),
}
}
/// Row lookup by slot index, bypassing the MPHF.
#[inline]
pub(crate) fn fill_row_by_slot(&self, slot: usize, n_genomes: usize) -> Vec<u32> {
let mut buf = vec![0u32; n_genomes];
match self {
SrcLayerData::Presence(_, mat) => mat.fill_row(slot, &mut buf),
SrcLayerData::Count(_, mat) => mat.fill_row(slot, &mut buf),
}
buf
}
/// Call `f` with a reference to the underlying matrix as `&dyn MatrixGroupOps`.
pub(crate) fn with_matrix<R>(&self, f: impl FnOnce(&dyn MatrixGroupOps) -> R) -> R {
match self {
SrcLayerData::Presence(_, mat) => f(mat),
SrcLayerData::Count(_, mat) => f(mat),
}
}
}
+1 -1
View File
@@ -1,6 +1,6 @@
use std::collections::HashMap;
use obikpartitionner::GroupQuorumFilter;
use crate::GroupQuorumFilter;
use obitaxonomy::{TaxPath, TaxPattern};
use crate::meta::{GenomeInfo, IndexMeta};
+240
View File
@@ -0,0 +1,240 @@
use std::collections::HashMap;
use std::path::Path;
use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix};
use obikseq::CanonicalKmer;
use obilayeredmap::{IndexMode, MphfLayer, OLMError};
use obiskio::{SKError, SKResult};
use crate::index::KmerIndex;
fn olm_to_sk(e: OLMError) -> SKError {
match e {
OLMError::Io(io_err) => SKError::Io(io_err),
other => SKError::InvalidData {
context: "query",
detail: other.to_string(),
},
}
}
// ── per-layer query handle ────────────────────────────────────────────────────
enum QueryLayer {
Presence(MphfLayer, PersistentBitMatrix),
Count(MphfLayer, PersistentCompactIntMatrix),
}
impl QueryLayer {
fn open(layer_dir: &Path, with_counts: bool, mode: &IndexMode) -> SKResult<Self> {
let mphf = MphfLayer::open(layer_dir, mode).map_err(olm_to_sk)?;
let counts_dir = layer_dir.join("counts");
let presence_dir = layer_dir.join("presence");
if with_counts && counts_dir.exists() {
let mat = PersistentCompactIntMatrix::open(layer_dir).map_err(SKError::Io)?;
Ok(QueryLayer::Count(mphf, mat))
} else if presence_dir.exists() || !counts_dir.exists() {
// presence mode, or no matrix at all → Implicit handled inside open()
let mat = PersistentBitMatrix::open(layer_dir).map_err(SKError::Io)?;
Ok(QueryLayer::Presence(mphf, mat))
} else {
// counts exist but not presence — count layer, no presence requested
let mat = PersistentCompactIntMatrix::open(layer_dir).map_err(SKError::Io)?;
Ok(QueryLayer::Count(mphf, mat))
}
}
/// MPHF lookup only — no matrix access. `Some(slot)` on hit.
fn find_slot(&self, kmer: CanonicalKmer) -> Option<usize> {
match self {
QueryLayer::Presence(mphf, _) | QueryLayer::Count(mphf, _) => mphf.find(kmer),
}
}
/// Number of genome columns this layer's matrix actually has. Bounds
/// column-major iteration — usually equal to the index's `n_genomes`, but
/// `PersistentBitMatrix::Implicit` (the documented mono-genome fast path)
/// always reports exactly `1`, regardless of the index's real genome
/// count, so callers must use this rather than assuming `n_genomes`.
fn n_cols(&self) -> usize {
match self {
QueryLayer::Presence(_, mat) => mat.n_cols(),
QueryLayer::Count(_, mat) => mat.n_cols(),
}
}
/// Every nonzero `(idx into slots, col, value)` triple among `slots`.
/// Format-agnostic: each matrix picks its own natural traversal
/// (`PersistentBitMatrix::nonzero_iter` dispatches to a genuinely
/// row-major decode on `Sparse`, not a column-major point-probe loop —
/// see `DevDocMD/architecture/siblings.md`, "`query` never benefits
/// from sparse row-major access"). Replaces the old per-`(genome,
/// slot)` `col_value` point lookup, which this layer's `Sparse`
/// presence matrices paid for badly: each such lookup rebuilt the
/// entire row just to return one cell.
fn nonzero_iter<'a>(
&'a self,
slots: &'a [usize],
) -> Box<dyn Iterator<Item = (usize, usize, u32)> + 'a> {
match self {
QueryLayer::Presence(_, mat) => mat.nonzero_iter(slots),
QueryLayer::Count(_, mat) => Box::new(mat.nonzero_iter(slots)),
}
}
}
// ── KmerDesc — one occurrence of a k-mer in the query batch ──────────────────
/// Describes one occurrence of a (deduplicated) k-mer in the query batch:
/// which sequence it came from, and its absolute s-mer position within it.
#[derive(Debug, Clone, Copy)]
pub struct KmerDesc {
pub seq_idx: u32,
pub pos: u32,
}
/// Aggregate counters for one `query_partition_with` call — feeds the
/// dedup-ratio and column-scan logging in `obikmer::cmd::query` (occurrences
/// vs. unique k-mers is the whole justification for k-mer-level
/// dereplication; columns scanned / `get()` calls quantify the column-major
/// fetch's locality claim).
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
pub struct QueryStats {
/// Distinct canonical k-mers queried in this partition.
pub n_unique_kmers: usize,
/// Total `MphfLayer::find` calls issued (a k-mer tried against more than
/// one layer before a hit, or against all layers on a miss, counts once
/// per layer attempted).
pub n_mphf_calls: usize,
/// Distinct canonical k-mers that matched some layer.
pub n_hits: usize,
/// Total genome columns scanned across all hit layers (sum of
/// `layer.n_cols()` over layers with at least one hit).
pub n_columns_scanned: usize,
/// Total `col_value` calls issued during the column-major fetch pass
/// (`n_columns_scanned` × hits-per-layer, summed over layers).
pub n_col_get_calls: usize,
}
impl std::ops::AddAssign for QueryStats {
fn add_assign(&mut self, other: Self) {
self.n_unique_kmers += other.n_unique_kmers;
self.n_mphf_calls += other.n_mphf_calls;
self.n_hits += other.n_hits;
self.n_columns_scanned += other.n_columns_scanned;
self.n_col_get_calls += other.n_col_get_calls;
}
}
// ── QueryHit — one event delivered to query_partition_with's callback ───────
/// One event from [`KmerPartition::query_partition_with`]'s two-stage query:
/// a `Found` event once per hit k-mer (stage 1, MPHF-only — mark the k-mer as
/// indexed regardless of any genome's value), then a `Value` event per
/// `(hit k-mer, genome)` pair with a nonzero matrix value (stage 2,
/// column-major fetch). Carried as one enum, not two separate callbacks, so
/// the caller only needs one `FnMut` closure — passing two closures that each
/// need to mutably borrow the same accumulator does not borrow-check.
pub enum QueryHit<'a> {
Found(&'a [KmerDesc]),
Value(&'a [KmerDesc], usize, u32),
}
// ── KmerPartition::query_partition_with ──────────────────────────────────────
impl KmerIndex {
/// Query a single partition for a pre-deduplicated map of canonical
/// k-mers → their occurrences (`seq_idx`, `pos`) in the query batch.
///
/// Two stages:
/// 1. **MPHF-only pass**: for each unique k-mer, try each layer's MPHF in
/// turn (stopping at the first hit) and bucket confirmed hits by
/// `(layer, slot)`. Emits one `QueryHit::Found` per hit k-mer. This
/// stage's cost is independent of the index's genome count.
/// 2. **Column-major fetch**: for each layer with at least one hit, walk
/// its matrix **column by column** (genome by genome) — for each
/// genome, scan the slots bucketed in stage 1 and look up their value.
/// Emits one `QueryHit::Value` per nonzero `(k-mer, genome)` pair.
/// Total lookups are the same as a row-major pass (`n_hits × n_cols`
/// in the worst case); the win is memory locality — both persistent
/// matrix formats are column-oriented on disk (one `mmap`'d region per
/// genome), so scanning one column at a time touches far fewer
/// distinct mmap regions than fetching one full row per hit.
pub fn query_partition_with<F>(
&self,
part_idx: usize,
kmers: &HashMap<CanonicalKmer, Vec<KmerDesc>>,
n_genomes: usize,
with_counts: bool,
mut on_event: F,
) -> SKResult<QueryStats>
where
F: FnMut(QueryHit),
{
let mut stats = QueryStats::default();
if kmers.is_empty() {
return Ok(stats);
}
let index_dir = self.index_dir(part_idx);
if !index_dir.exists() {
return Ok(stats);
}
let meta = self.partition_meta(part_idx)?;
let layers: Vec<QueryLayer> = (0..meta.n_layers)
.map(|i| QueryLayer::open(&self.layer_dir(part_idx, i), with_counts, &meta.mode))
.collect::<SKResult<_>>()?;
// ── Stage 1: MPHF-only pass, bucket hits by (layer_idx, slot) ────────
let mut by_layer: Vec<HashMap<usize, &Vec<KmerDesc>>> =
(0..layers.len()).map(|_| HashMap::new()).collect();
for (kmer, descs) in kmers {
stats.n_unique_kmers += 1;
for (layer_idx, layer) in layers.iter().enumerate() {
stats.n_mphf_calls += 1;
if let Some(slot) = layer.find_slot(*kmer) {
by_layer[layer_idx].insert(slot, descs);
on_event(QueryHit::Found(descs));
stats.n_hits += 1;
break;
}
}
}
// ── Stage 2: nonzero-cell fetch, per layer ────────────────────────────
// Format-agnostic — see `QueryLayer::nonzero_iter`. `n_cols` still
// bounds accepted genome columns (Implicit reports fewer than
// `n_genomes`; see `n_cols`'s doc), cells beyond it are dropped
// rather than ever produced, since `nonzero_iter` only knows the
// matrix's own column count, not the caller's `n_genomes`.
for (layer_idx, slots) in by_layer.iter().enumerate() {
if slots.is_empty() {
continue;
}
let layer = &layers[layer_idx];
let n_cols = layer.n_cols().min(n_genomes);
stats.n_columns_scanned += n_cols;
let slot_list: Vec<usize> = slots.keys().copied().collect();
for (idx, g, v) in layer.nonzero_iter(&slot_list) {
if g >= n_cols {
continue;
}
stats.n_col_get_calls += 1;
debug_assert_ne!(v, 0, "nonzero_iter must not yield zero-valued cells");
let descs = slots[&slot_list[idx]];
on_event(QueryHit::Value(descs, g, v));
}
}
Ok(stats)
}
}
#[cfg(test)]
#[path = "tests/query_layer.rs"]
mod tests;
+3 -4
View File
@@ -1,6 +1,6 @@
use std::path::Path;
use obikpartitionner::{KmerFilter, MergeMode};
use crate::{KmerFilter, MergeMode};
use obisys::{Reporter, Stage, progress_bar};
use tracing::info;
@@ -46,7 +46,7 @@ impl KmerIndex {
meta.genomes = src.meta.genomes.clone();
let n_genomes = src.meta.genomes.len();
let n_partitions = src.partition.n_partitions();
let n_partitions = src.n_partitions();
// ── Create an empty destination KmerPartition ─────────────────────────
let dst_partition = KmerIndex::create_skeleton(output, &meta)?;
@@ -59,14 +59,13 @@ impl KmerIndex {
let t = Stage::start("rebuild");
let pb = progress_bar("rebuild", n_partitions as u64, "partitions");
let src_partition = &src.partition;
let block_bits = meta.config.block_bits;
let order: Vec<usize> = (0..n_partitions).collect();
let runner = crate::numa::PartitionRunner::new();
runner.run(
&order,
|i| dst_partition.rebuild_partition(src_partition, i, filters, mode, n_genomes, block_bits),
|i| dst_partition.rebuild_partition(src, i, filters, mode, n_genomes, block_bits),
|_, _, _| { pb.inc(1); },
).map_err(OKIError::Partition)?;
+270
View File
@@ -0,0 +1,270 @@
use std::path::Path;
use obicompactvec::{
FilterMask, PersistentBitMatrixBuilder, PersistentBitVecBuilder,
PersistentCompactIntMatrixBuilder, PersistentCompactIntVecBuilder, eval_filter_mask,
};
use obidebruinj::GraphDeBruijn;
use obikseq::CanonicalKmer;
use obilayeredmap::meta::PartitionMeta;
use obilayeredmap::{IndexMode, MphfLayer, layer_dir};
use obiskio::{SKError, SKResult, UnitigFileReader};
use crate::common::{load_meta, olm_to_sk};
use crate::filter::KmerFilter;
use crate::graph_pipeline::materialize_layer;
use crate::merge_layer::{MergeMode, SrcLayerData};
use crate::index::KmerIndex;
// ── Builders — pair matrix builder + column builders for one mode ─────────────
enum Builders {
Presence(PersistentBitMatrixBuilder, Vec<PersistentBitVecBuilder>),
Count(
PersistentCompactIntMatrixBuilder,
Vec<PersistentCompactIntVecBuilder>,
),
}
impl Builders {
fn new(mode: MergeMode, n: usize, dir: &Path, n_genomes: usize) -> SKResult<Self> {
match mode {
MergeMode::Presence => {
let mut mat = PersistentBitMatrixBuilder::new(n, dir).map_err(SKError::Io)?;
let mut cols = Vec::with_capacity(n_genomes);
for _ in 0..n_genomes {
cols.push(mat.add_col().map_err(SKError::Io)?);
}
Ok(Builders::Presence(mat, cols))
}
MergeMode::Count => {
let mut mat =
PersistentCompactIntMatrixBuilder::new(n, dir).map_err(SKError::Io)?;
let mut cols = Vec::with_capacity(n_genomes);
for _ in 0..n_genomes {
cols.push(mat.add_col().map_err(SKError::Io)?);
}
Ok(Builders::Count(mat, cols))
}
}
}
fn set_val(&mut self, col: usize, slot: usize, value: u32) {
match self {
Builders::Presence(_, cols) => cols[col].set(slot, value > 0),
Builders::Count(_, cols) => cols[col].set(slot, value),
}
}
fn close(self) -> SKResult<()> {
match self {
Builders::Presence(mat, cols) => {
for b in cols {
b.close().map_err(SKError::Io)?;
}
mat.close().map_err(SKError::Io)
}
Builders::Count(mat, cols) => {
for b in cols {
b.close().map_err(SKError::Io)?;
}
mat.close().map_err(SKError::Io)
}
}
}
}
// ── try_compute_combined_mask ─────────────────────────────────────────────────
/// Build a per-slot `TempBitVec` mask from `filters` using column operations
/// on the source matrix — no per-kmer MPHF lookup or row read needed.
///
/// Returns `Some(mask)` when every filter in `filters` can express itself as
/// a [`FilterMask`] expression. Returns `None` when any filter requires
/// row-level inspection (fall back to `passes_all`).
fn try_compute_combined_mask(
filters: &[Box<dyn KmerFilter>],
src_data: &SrcLayerData,
n_genomes: usize,
) -> SKResult<Option<obicompactvec::TempBitVec>> {
if filters.is_empty() {
return Ok(None);
}
let mut exprs: Vec<FilterMask> = Vec::with_capacity(filters.len());
for f in filters {
match f.column_mask_expr(n_genomes) {
Some(expr) => exprs.push(expr),
None => return Ok(None),
}
}
let combined = FilterMask::And(exprs);
let n = src_data.n_slots();
let mask = src_data
.with_matrix(|mat| eval_filter_mask(&combined, mat, n))
.map_err(SKError::Io)?;
Ok(Some(mask))
}
// ── iter_src_kmers_masked (pass 1) ────────────────────────────────────────────
/// Iterate all passing kmers in `src_index_dir`, yielding only the kmer value.
///
/// When all filters can be expressed as column operations, a per-slot mask is
/// computed once per layer and used for O(1) slot-check per kmer instead of a
/// full row read. Falls back to row-level `passes_all` otherwise.
fn iter_src_kmers_masked(
src_index_dir: &Path,
mode: MergeMode,
n_genomes: usize,
filters: &[Box<dyn KmerFilter>],
mut cb: impl FnMut(CanonicalKmer),
) -> SKResult<()> {
let src_meta = load_meta(src_index_dir, "rebuild")?;
for l in 0..src_meta.n_layers {
let src_layer_dir = layer_dir(src_index_dir, l);
let unitigs_path = src_layer_dir.join("unitigs.bin");
if !unitigs_path.exists() {
continue;
}
let src_data = SrcLayerData::open(&src_layer_dir, mode)?;
let mask = try_compute_combined_mask(filters, &src_data, n_genomes)?;
let reader = UnitigFileReader::open_sequential(&unitigs_path)?;
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
let slot = src_data.slot(kmer);
let passes = match &mask {
Some(m) => m.get(slot),
None => {
let row = src_data.fill_row_by_slot(slot, n_genomes);
filters.iter().all(|f| f.passes(kmer, &row, n_genomes))
}
};
if passes {
cb(kmer);
}
}
}
Ok(())
}
// ── iter_src_layers (pass 2) ──────────────────────────────────────────────────
/// Iterate all passing kmers in `src_index_dir`, yielding `(kmer, row)`.
///
/// When the slot mask is available, skips the row read for filtered-out slots.
fn iter_src_layers(
src_index_dir: &Path,
mode: MergeMode,
n_genomes: usize,
filters: &[Box<dyn KmerFilter>],
mut cb: impl FnMut(CanonicalKmer, Box<[u32]>),
) -> SKResult<()> {
let src_meta = load_meta(src_index_dir, "rebuild")?;
for l in 0..src_meta.n_layers {
let src_layer_dir = layer_dir(src_index_dir, l);
let unitigs_path = src_layer_dir.join("unitigs.bin");
if !unitigs_path.exists() {
continue;
}
let src_data = SrcLayerData::open(&src_layer_dir, mode)?;
let mask = try_compute_combined_mask(filters, &src_data, n_genomes)?;
let reader = UnitigFileReader::open_sequential(&unitigs_path)?;
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
let slot = src_data.slot(kmer);
if let Some(ref m) = mask {
if !m.get(slot) {
continue;
}
let row = src_data.fill_row_by_slot(slot, n_genomes);
cb(kmer, row.into_boxed_slice());
} else {
let row = src_data.fill_row_by_slot(slot, n_genomes);
if filters.iter().all(|f| f.passes(kmer, &row, n_genomes)) {
cb(kmer, row.into_boxed_slice());
}
}
}
}
Ok(())
}
// ── KmerPartition::rebuild_partition ─────────────────────────────────────────
impl KmerIndex {
/// Rebuild partition `i` from `src` into `self` (an empty destination partition).
///
/// Only k-mers whose per-genome row passes all `filters` are written.
/// The output is a single-layer index — regardless of how many layers the
/// source has.
///
/// `n_genomes` is the number of genome columns in the source (and destination).
pub fn rebuild_partition(
&self,
src: &KmerIndex,
i: usize,
filters: &[Box<dyn KmerFilter>],
mode: MergeMode,
n_genomes: usize,
block_bits: u8,
) -> SKResult<()> {
let src_index_dir = src.index_dir(i);
if !src_index_dir.exists() {
return Ok(());
}
let src_meta = load_meta(&src_index_dir, "rebuild")?;
if src_meta.n_layers == 0 {
return Ok(());
}
// ── Pass 1: collect filtered kmers into de Bruijn graph ───────────────
let mut g = GraphDeBruijn::new();
iter_src_kmers_masked(&src_index_dir, mode, n_genomes, filters, |kmer| {
g.push(kmer);
})?;
if g.len() == 0 {
return Ok(());
}
// ── Build MPHF in dst layer_0 ─────────────────────────────────────────
let dst_index_dir = self.index_dir(i);
let dst_layer_dir = self.layer_dir(i, 0);
let n_new = materialize_layer(g, &dst_layer_dir, block_bits, &IndexMode::Exact)?;
let dst_mphf = MphfLayer::open(&dst_layer_dir, &IndexMode::Exact)
.map_err(|e| olm_to_sk(e, "rebuild"))?;
// ── Prepare matrix builders (one column per genome) ───────────────────
let data_dir = match mode {
MergeMode::Presence => dst_layer_dir.join("presence"),
MergeMode::Count => dst_layer_dir.join("counts"),
};
std::fs::create_dir_all(&data_dir)?;
let mut builders = Builders::new(mode, n_new, &data_dir, n_genomes)?;
// ── Pass 2: fill builders ─────────────────────────────────────────────
iter_src_layers(&src_index_dir, mode, n_genomes, filters, |kmer, row| {
if let Some(slot) = dst_mphf.find(kmer) {
for (col, &value) in row.iter().enumerate() {
builders.set_val(col, slot, value);
}
}
})?;
// ── Close builders and write metadata ─────────────────────────────────
builders.close()?;
PartitionMeta {
n_layers: 1,
mode: IndexMode::Exact,
}
.save(&dst_index_dir)
.map_err(|e| olm_to_sk(e, "rebuild"))?;
Ok(())
}
}
+6 -7
View File
@@ -1,4 +1,3 @@
use obikpartitionner::KmerPartitions;
use obilayeredmap::{IndexMode, layer::Layer};
use obisys::{Reporter, Stage, progress_bar};
use std::fs;
@@ -35,7 +34,7 @@ impl KmerIndex {
return Err(OKIError::NotIndexed(self.root_path.clone()));
}
let n = self.partition.n_partitions();
let n = self.n_partitions();
info!(
"reindex {} partition(s): {:?} → {:?}",
n, self.meta.config.evidence, target,
@@ -49,7 +48,7 @@ impl KmerIndex {
runner.run(
&order,
|i| {
reindex_partition(&self.partition, i, &target, block_bits)
reindex_partition(self, i, &target, block_bits)
.map_err(|e| OKIError::InvalidInput(format!("partition {i}: {e}")))
},
|_, _, _| {
@@ -71,19 +70,19 @@ impl KmerIndex {
/// Process all layers of one partition's index directory.
fn reindex_partition(
partition: &KmerPartitions,
index: &KmerIndex,
i: usize,
target: &IndexMode,
block_bits: u8,
) -> OKIResult<()> {
if !partition.index_dir(i).exists() {
if !index.index_dir(i).exists() {
return Ok(());
}
let n_layers = partition
let n_layers = index
.n_layers(i)
.map_err(|e| OKIError::InvalidInput(e.to_string()))?;
for layer_idx in 0..n_layers {
reindex_layer(&partition.layer_dir(i, layer_idx), target, block_bits)?;
reindex_layer(&index.layer_dir(i, layer_idx), target, block_bits)?;
}
Ok(())
}
+6 -15
View File
@@ -1,6 +1,6 @@
use std::path::Path;
use obikpartitionner::{KmerPartitions, OutputCol};
use crate::OutputCol;
use obisys::{Reporter, Stage, progress_bar};
use tracing::info;
@@ -41,7 +41,7 @@ impl KmerIndex {
.collect();
let n_src_genomes = src.meta.genomes.len();
let n_partitions = src.partition.n_partitions();
let n_partitions = src.n_partitions();
let dst_partition = KmerIndex::create_skeleton(output, &meta)?;
@@ -54,7 +54,6 @@ impl KmerIndex {
let t = Stage::start("select");
let pb = progress_bar("select", n_partitions as u64, "partitions");
let src_partition = &src.partition;
let order: Vec<usize> = (0..n_partitions).collect();
let runner = crate::numa::PartitionRunner::new();
@@ -63,7 +62,7 @@ impl KmerIndex {
&order,
|i| {
dst_partition.select_partition(
src_partition,
src,
i,
specs,
n_src_genomes,
@@ -99,14 +98,7 @@ impl KmerIndex {
}
let n_src_genomes = self.meta.genomes.len();
let n_partitions = self.partition.n_partitions();
let src_partition = KmerPartitions::open_with_config(
&self.root_path,
self.meta.config.kmer_size,
self.meta.config.minimizer_size,
self.meta.config.n_bits,
)?;
let n_partitions = self.n_partitions();
info!(
"select (in-place): {} partition(s), {} source genome(s) → {} output column(s)",
@@ -118,15 +110,14 @@ impl KmerIndex {
let t = Stage::start("select");
let pb = progress_bar("select", n_partitions as u64, "partitions");
let partition = &self.partition;
let order: Vec<usize> = (0..n_partitions).collect();
let runner = crate::numa::PartitionRunner::new();
runner
.run(
&order,
|i| {
partition.select_partition(
&src_partition,
self.select_partition(
self,
i,
specs,
n_src_genomes,
+309
View File
@@ -0,0 +1,309 @@
use std::fs;
use std::io;
use std::path::{Path, PathBuf};
use obicompactvec::{
ColGroup, MatrixGroupOps, PersistentBitMatrix, PersistentBitMatrixBuilder,
PersistentCompactIntMatrix, PersistentCompactIntMatrixBuilder,
};
use obilayeredmap::OLMError;
use obiskio::{SKError, SKResult};
use crate::index::KmerIndex;
// ── AggOp ─────────────────────────────────────────────────────────────────────
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AggOp {
Any,
All,
None,
Sum,
Min,
Max,
}
impl AggOp {
pub fn is_logical(self) -> bool {
matches!(self, AggOp::Any | AggOp::All | AggOp::None)
}
}
// ── OutputCol ─────────────────────────────────────────────────────────────────
pub struct OutputCol {
pub label: String,
pub indices: Vec<usize>,
pub op: AggOp,
}
// ── Helpers ───────────────────────────────────────────────────────────────────
fn olm_to_sk(e: OLMError) -> SKError {
match e {
OLMError::Io(e) => SKError::Io(e),
other => SKError::InvalidData {
context: "select",
detail: other.to_string(),
},
}
}
/// Copy all plain files (not subdirectories) from `src_dir` to `dst_dir`.
fn copy_layer_files(src_dir: &Path, dst_dir: &Path) -> io::Result<()> {
for entry in fs::read_dir(src_dir)? {
let entry = entry?;
let path = entry.path();
if path.is_file() {
fs::copy(&path, dst_dir.join(entry.file_name()))?;
}
}
Ok(())
}
// ── fill_builders ─────────────────────────────────────────────────────────────
fn fill_builders(
specs: &[OutputCol],
src_layer_dir: &Path,
src_is_count: bool,
threshold: u32,
output_presence: bool,
mut dst_bit: Option<&mut PersistentBitMatrixBuilder>,
mut dst_int: Option<&mut PersistentCompactIntMatrixBuilder>,
) -> SKResult<()> {
if src_is_count {
let mat = PersistentCompactIntMatrix::open(src_layer_dir).map_err(SKError::Io)?;
for spec in specs {
let g = ColGroup::new(&spec.label, spec.indices.clone());
if output_presence {
let b = dst_bit.as_deref_mut().unwrap();
match spec.op {
AggOp::Any => {
b.add_col_from(&mat.partial_group_any(&g, threshold).map_err(SKError::Io)?)
}
AggOp::All => {
b.add_col_from(&mat.partial_group_all(&g, threshold).map_err(SKError::Io)?)
}
AggOp::None => {
b.add_col_from(&mat.partial_group_none(&g, threshold).map_err(SKError::Io)?)
}
AggOp::Sum => {
b.add_col_from_int(&mat.partial_group_sum(&g).map_err(SKError::Io)?)
}
AggOp::Min => {
b.add_col_from_int(&mat.partial_group_min(&g).map_err(SKError::Io)?)
}
AggOp::Max => {
b.add_col_from_int(&mat.partial_group_max(&g).map_err(SKError::Io)?)
}
}
.map_err(SKError::Io)?;
} else {
let b = dst_int.as_deref_mut().unwrap();
match spec.op {
AggOp::Sum => b.add_col_from(&mat.partial_group_sum(&g).map_err(SKError::Io)?),
AggOp::Min => b.add_col_from(&mat.partial_group_min(&g).map_err(SKError::Io)?),
AggOp::Max => b.add_col_from(&mat.partial_group_max(&g).map_err(SKError::Io)?),
AggOp::Any => b.add_col_from_bit(
&mat.partial_group_any(&g, threshold).map_err(SKError::Io)?,
),
AggOp::All => b.add_col_from_bit(
&mat.partial_group_all(&g, threshold).map_err(SKError::Io)?,
),
AggOp::None => b.add_col_from_bit(
&mat.partial_group_none(&g, threshold).map_err(SKError::Io)?,
),
}
.map_err(SKError::Io)?;
}
}
} else {
let mat = PersistentBitMatrix::open(src_layer_dir).map_err(SKError::Io)?;
for spec in specs {
let g = ColGroup::new(&spec.label, spec.indices.clone());
if output_presence {
let b = dst_bit.as_deref_mut().unwrap();
match spec.op {
AggOp::Any => {
b.add_col_from(&mat.partial_group_any(&g, 1).map_err(SKError::Io)?)
}
AggOp::All => {
b.add_col_from(&mat.partial_group_all(&g, 1).map_err(SKError::Io)?)
}
AggOp::None => {
b.add_col_from(&mat.partial_group_none(&g, 1).map_err(SKError::Io)?)
}
AggOp::Sum => {
b.add_col_from_int(&mat.partial_group_sum(&g).map_err(SKError::Io)?)
}
AggOp::Min => {
b.add_col_from_int(&mat.partial_group_min(&g).map_err(SKError::Io)?)
}
AggOp::Max => {
b.add_col_from_int(&mat.partial_group_max(&g).map_err(SKError::Io)?)
}
}
.map_err(SKError::Io)?;
} else {
let b = dst_int.as_deref_mut().unwrap();
match spec.op {
AggOp::Sum => b.add_col_from(&mat.partial_group_sum(&g).map_err(SKError::Io)?),
AggOp::Min => b.add_col_from(&mat.partial_group_min(&g).map_err(SKError::Io)?),
AggOp::Max => b.add_col_from(&mat.partial_group_max(&g).map_err(SKError::Io)?),
AggOp::Any => {
b.add_col_from_bit(&mat.partial_group_any(&g, 1).map_err(SKError::Io)?)
}
AggOp::All => {
b.add_col_from_bit(&mat.partial_group_all(&g, 1).map_err(SKError::Io)?)
}
AggOp::None => {
b.add_col_from_bit(&mat.partial_group_none(&g, 1).map_err(SKError::Io)?)
}
}
.map_err(SKError::Io)?;
}
}
}
Ok(())
}
// ── KmerPartition::select_partition ──────────────────────────────────────────
impl KmerIndex {
/// Rewrite the data matrices of partition `i` in `src` into `self`.
///
/// `specs` defines the output columns (projection/aggregation).
/// `output_presence` — if true, all output builders use bit (0/1) format.
/// `in_place` — `self` and `src` share the same root; write to temp dirs then swap.
pub fn select_partition(
&self,
src: &KmerIndex,
i: usize,
specs: &[OutputCol],
_n_src_genomes: usize,
threshold: u32,
output_presence: bool,
in_place: bool,
) -> SKResult<()> {
let src_index_dir = src.index_dir(i);
if !src_index_dir.exists() {
return Ok(());
}
let n_src_layers = src.n_layers(i)?;
if n_src_layers == 0 {
return Ok(());
}
let dst_index_dir = self.index_dir(i);
if !in_place {
fs::create_dir_all(&dst_index_dir)?;
}
let data_subdir = if output_presence {
"presence"
} else {
"counts"
};
for l in 0..n_src_layers {
let src_layer_dir = src.layer_dir(i, l);
if !src_layer_dir.exists() {
continue;
}
let dst_layer_dir = self.layer_dir(i, l);
let counts_dir = src_layer_dir.join("counts");
let presence_dir = src_layer_dir.join("presence");
let src_is_count = counts_dir.exists() && !presence_dir.exists();
// Determine number of slots and detect implicit layers.
let n = if counts_dir.exists() {
PersistentCompactIntMatrix::open(&src_layer_dir)
.map_err(SKError::Io)?
.n()
} else if presence_dir.exists() {
PersistentBitMatrix::open(&src_layer_dir)
.map_err(SKError::Io)?
.n()
} else {
// Implicit single-genome layer: no data matrix needed in output either.
if !in_place {
fs::create_dir_all(&dst_layer_dir)?;
copy_layer_files(&src_layer_dir, &dst_layer_dir)?;
}
continue;
};
// Choose the output data directory (temp name for in-place).
let (dst_data_dir, final_data_dir): (PathBuf, PathBuf) = if in_place {
let tmp = dst_layer_dir.join(format!("{data_subdir}_new"));
let perm = dst_layer_dir.join(data_subdir);
(tmp, perm)
} else {
let perm = dst_layer_dir.join(data_subdir);
(perm.clone(), perm)
};
if !in_place {
fs::create_dir_all(&dst_layer_dir)?;
copy_layer_files(&src_layer_dir, &dst_layer_dir)?;
}
fs::create_dir_all(&dst_data_dir)?;
let (mut dst_bit, mut dst_int) = if output_presence {
(
Some(PersistentBitMatrixBuilder::new(n, &dst_data_dir).map_err(SKError::Io)?),
None,
)
} else {
(
None,
Some(
PersistentCompactIntMatrixBuilder::new(n, &dst_data_dir)
.map_err(SKError::Io)?,
),
)
};
fill_builders(
specs,
&src_layer_dir,
src_is_count,
threshold,
output_presence,
dst_bit.as_mut(),
dst_int.as_mut(),
)?;
if output_presence {
dst_bit.unwrap().close().map_err(SKError::Io)?;
} else {
dst_int.unwrap().close().map_err(SKError::Io)?;
}
// In-place: swap old data dir for new.
if in_place {
let old_data_dir = if src_is_count {
dst_layer_dir.join("counts")
} else {
dst_layer_dir.join("presence")
};
if old_data_dir.exists() {
fs::remove_dir_all(&old_data_dir)?;
}
fs::rename(&dst_data_dir, &final_data_dir)?;
}
}
if !in_place {
src.partition_meta(i)?
.save(&dst_index_dir)
.map_err(olm_to_sk)?;
}
Ok(())
}
}
+6 -6
View File
@@ -89,13 +89,13 @@ impl KmerIndex {
let (n_kmers, mphf_b, evidence_b, matrix_b) = (0..n)
.into_par_iter()
.map(|i| {
let index_dir = self.partition.index_dir(i);
let index_dir = self.index_dir(i);
if !index_dir.exists() { return (0usize, 0u64, 0u64, 0u64); }
let n_layers = self.partition.n_layers(i).unwrap_or(0);
let n_layers = self.n_layers(i).unwrap_or(0);
(0..n_layers).fold((0usize, 0u64, 0u64, 0u64), |acc, l| {
let lb = layer_bytes(&self.partition.layer_dir(i, l));
let lb = layer_bytes(&self.layer_dir(i, l));
(acc.0 + lb.n_kmers, acc.1 + lb.mphf, acc.2 + lb.evidence, acc.3 + lb.matrix)
})
})
@@ -138,13 +138,13 @@ impl KmerIndex {
let mut counts = vec![0u64; n_genomes];
let mut n_kmers = 0usize;
let index_dir = self.partition.index_dir(i);
let index_dir = self.index_dir(i);
if !index_dir.exists() { return (0, counts); }
let n_layers = self.partition.n_layers(i).unwrap_or(0);
let n_layers = self.n_layers(i).unwrap_or(0);
for l in 0..n_layers {
let this_layer_dir = self.partition.layer_dir(i, l);
let this_layer_dir = self.layer_dir(i, l);
if !this_layer_dir.exists() { continue; }
n_kmers += LayerMeta::load(&this_layer_dir).map(|m| m.n).unwrap_or(0);
+94
View File
@@ -0,0 +1,94 @@
use super::*;
use crate::meta::IndexConfig;
// ── QueryStats::AddAssign ───────────────────────────────────────────────────
#[test]
fn query_stats_add_assign_sums_fields() {
let mut total = QueryStats {
n_unique_kmers: 3,
n_mphf_calls: 5,
n_hits: 2,
n_columns_scanned: 1,
n_col_get_calls: 7,
};
total += QueryStats {
n_unique_kmers: 1,
n_mphf_calls: 4,
n_hits: 1,
n_columns_scanned: 2,
n_col_get_calls: 3,
};
assert_eq!(total.n_unique_kmers, 4);
assert_eq!(total.n_mphf_calls, 9);
assert_eq!(total.n_hits, 3);
assert_eq!(total.n_columns_scanned, 3);
assert_eq!(total.n_col_get_calls, 10);
}
#[test]
fn query_stats_default_is_zero() {
let s = QueryStats::default();
assert_eq!(s.n_unique_kmers, 0);
assert_eq!(s.n_mphf_calls, 0);
assert_eq!(s.n_hits, 0);
assert_eq!(s.n_columns_scanned, 0);
assert_eq!(s.n_col_get_calls, 0);
}
// ── query_partition_with on a not-yet-indexed partition ─────────────────────
/// A `KmerPartition` created but never taken through `build_layers` has no
/// `index/` subdirectory under any partition — `query_partition_with` must
/// recognise this and return default (all-zero) stats rather than erroring,
/// exactly like an empty `kmers` map.
#[test]
fn query_partition_with_missing_index_dir_returns_default_stats() {
let tmp = tempfile::tempdir().expect("tempdir");
let config = IndexConfig {
kmer_size: 21,
minimizer_size: 9,
n_bits: 2,
with_counts: false,
evidence: obilayeredmap::IndexMode::Exact,
block_bits: 0,
};
let index = KmerIndex::create(tmp.path().join("idx"), config, None, false).expect("create index");
let mut kmers: HashMap<CanonicalKmer, Vec<KmerDesc>> = HashMap::new();
// Any well-formed canonical k-mer works here — the call must return
// before ever attempting an MPHF lookup, since `index/` doesn't exist.
let kmer = CanonicalKmer::from_raw_unchecked(0u64);
kmers.insert(kmer, vec![KmerDesc { seq_idx: 0, pos: 0 }]);
let stats = index
.query_partition_with(0, &kmers, 1, false, |_event| {
panic!("on_event must not be called: no index was built");
})
.expect("query_partition_with should not error on a missing index dir");
assert_eq!(stats, QueryStats::default());
}
#[test]
fn query_partition_with_empty_kmers_is_a_noop() {
let tmp = tempfile::tempdir().expect("tempdir");
let config = IndexConfig {
kmer_size: 21,
minimizer_size: 9,
n_bits: 2,
with_counts: false,
evidence: obilayeredmap::IndexMode::Exact,
block_bits: 0,
};
let index = KmerIndex::create(tmp.path().join("idx"), config, None, false).expect("create index");
let kmers: HashMap<CanonicalKmer, Vec<KmerDesc>> = HashMap::new();
let stats = index
.query_partition_with(0, &kmers, 1, false, |_event| {
panic!("on_event must not be called on an empty kmer map");
})
.expect("query_partition_with on an empty map should not error");
assert_eq!(stats, QueryStats::default());
}