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

Extract index modules into specialized workspace subcrates

Browse Source 
This commit partitions the obikindex crate into multiple focused subcrates (obikfilter, obikmerge, obikquery, obikrebuild, obikselect, obikstats, obikdump, and obikidxcache) to reduce coupling and clarify module boundaries. It standardizes error handling across the workspace using OKIError and OKIResult, updates index APIs to support lazy, disk-backed partition access, and migrates NUMA system utilities to a new obisys crate. All modifications are structural, focusing on dependency graph expansion, import path updates, and API surface reorganization without altering core runtime behavior.
This commit is contained in:
Eric Coissac
2026-08-22 06:25:28 +02:00
parent c9d10d55c7
commit fc4464a0ef
81 changed files with 1640 additions and 1196 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/obikindex/Cargo.toml
-6
View File
@@ -19,13 +19,11 @@ niffler = "3.0.0"
memmap2 = "0.9.10"
ndarray = "0.17"
rayon = "1"
crossbeam-channel = "0.5"
serde = { version = "1", features = ["derive"] }
serde_json = "1"
indicatif = "0.18"
tracing = "0.1.44"
bitvec = "1"
hwlocality = { version = "1.0.0-alpha.11", features = ["vendored"], optional = true }
[dev-dependencies]
obiread = { path = "../obiread" }
@@ -33,7 +31,3 @@ obikseq = { path = "../obikseq", features = ["test-utils"] }
tempfile = "3"
tracing-subscriber = { version = "0.3", features = ["fmt", "env-filter"] }
anyhow = "1"
[features]
default = ["numa"]
numa = ["hwlocality"]
src/obikindex/src/index/builder.rs
+15 -4
View File
@@ -44,7 +44,11 @@ pub trait IndexBuilder: Sized {
/// For construction paths that build partitions from scratch (`select`,
/// `rebuild`). `merge` bootstraps by copying a source index instead, so
/// it does not use this.
fn create_skeleton<P: AsRef<Path>>(output: P, config: IndexConfig, genomes: Vec<GenomeInfo>) -> OKIResult<Self>;
fn create_skeleton<P: AsRef<Path>>(
output: P,
config: IndexConfig,
genomes: Vec<GenomeInfo>,
) -> OKIResult<Self>;
/// Mark `output` as fully indexed, pack its column matrices, and reopen it.
///
@@ -61,7 +65,7 @@ pub trait IndexBuilder: Sized {
impl IndexBuilder for KmerIndex {
fn clear_output_for_create<P: AsRef<Path>>(output: P, force: bool) -> OKIResult<()> {
let output = output.as_ref();
if Self::exists(output) {
if Self::is_an_index(output) {
if force {
fs::remove_dir_all(output).map_err(OKIError::Io)?;
} else {
@@ -74,11 +78,18 @@ impl IndexBuilder for KmerIndex {
Ok(())
}
fn create_skeleton<P: AsRef<Path>>(output: P, config: IndexConfig, genomes: Vec<GenomeInfo>) -> OKIResult<Self> {
fn create_skeleton<P: AsRef<Path>>(
output: P,
config: IndexConfig,
genomes: Vec<GenomeInfo>,
) -> OKIResult<Self> {
let output = output.as_ref();
fs::create_dir_all(output).map_err(OKIError::Io)?;
let meta = IndexMeta::create_at(output, config, genomes).map_err(OKIError::Io)?;
Ok(KmerIndex { root_path: output.to_owned(), meta: Arc::new(meta) })
Ok(KmerIndex {
root_path: output.to_owned(),
meta: Arc::new(meta),
})
}
fn finalize_indexed<P: AsRef<Path>>(output: P, rep: &mut Reporter) -> OKIResult<Self> {
src/obikindex/src/index/common.rs
+6 -46
View File
@@ -1,65 +1,25 @@
use std::path::Path;
use crate::index::error::{OKIError, OKIResult};
use obicompactvec::{PersistentBitVecBuilder, PersistentCompactIntVecBuilder};
use crate::layer::meta::PartitionMeta;
use crate::layer::{layer_dir, IndexMode, OLMError};
use obiskio::{SKError, SKResult};
// ── olm_to_sk ────────────────────────────────────────────────────────────────
pub 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 {
pub enum ColBuilder {
Bit(PersistentBitVecBuilder),
Int(PersistentCompactIntVecBuilder),
}
impl ColBuilder {
pub(crate) fn set_val(&mut self, slot: usize, value: u32) {
pub 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<()> {
pub fn close(self) -> OKIResult<()> {
match self {
ColBuilder::Bit(b) => b.close().map_err(SKError::Io),
ColBuilder::Int(b) => b.close().map_err(SKError::Io),
ColBuilder::Bit(b) => b.close().map_err(OKIError::Io),
ColBuilder::Int(b) => b.close().map_err(OKIError::Io),
}
}
}
src/obikindex/src/index/distance.rs
-135
View File
@@ -1,135 +0,0 @@
use ndarray::Array2;
use obicompactvec::traits::{BitPartials, CountPartials};
use crate::layer::LayeredStore;
use rayon::prelude::*;
use crate::index::error::{OKIError, OKIResult};
use crate::index::kmer_index::KmerIndex;
// ── Public API ────────────────────────────────────────────────────────────────
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DistanceMetric {
/// Jaccard distance on presence/absence data.
Jaccard,
/// Hamming distance (number of differing kmer positions) on presence/absence data.
Hamming,
/// Mash distance on presence/absence data (Jaccard-derived mutation-rate estimate).
Mash,
/// Bray-Curtis dissimilarity on raw counts.
BrayCurtis,
/// Bray-Curtis dissimilarity normalised by per-genome total counts.
RelfreqBrayCurtis,
/// Euclidean distance on raw counts.
Euclidean,
/// Euclidean distance on relative frequencies.
RelfreqEuclidean,
/// Hellinger distance on counts.
Hellinger,
/// Euclidean distance in the Hellinger (√relative-frequency) space (unnormalised variant).
HellingerEuclidean,
}
pub struct DistanceOutput {
/// n×n pairwise distance matrix (genomes in index order).
pub matrix: Array2<f64>,
/// n×n shared-kmer count matrix (intersection), if requested.
pub shared_kmers: Option<Array2<u64>>,
}
impl DistanceMetric {
pub fn requires_counts(self) -> bool {
matches!(
self,
DistanceMetric::BrayCurtis
| DistanceMetric::RelfreqBrayCurtis
| DistanceMetric::Euclidean
| DistanceMetric::RelfreqEuclidean
| DistanceMetric::Hellinger
| DistanceMetric::HellingerEuclidean
)
}
}
// ── KmerIndex::distance ───────────────────────────────────────────────────────
impl KmerIndex {
pub fn distance(&self, metric: DistanceMetric, shared_kmers: bool, presence_threshold: u32) -> OKIResult<DistanceOutput> {
let n_genomes = self.meta.genomes().map_err(OKIError::Io)?.len();
if n_genomes < 2 {
return Err(OKIError::InvalidInput(
"distance requires at least 2 genomes in the index".into(),
));
}
let use_counts = self.meta.config.with_counts;
if metric.requires_counts() && !use_counts {
return Err(OKIError::InvalidInput(format!(
"{metric:?} requires a count index (with_counts = true)"
)));
}
let n_parts = self.n_partitions();
if use_counts {
let stores: Vec<_> = (0..n_parts)
.into_par_iter()
.map(|i| self.count_store(i).map_err(OKIError::Partition))
.collect::<OKIResult<_>>()?;
let global = LayeredStore::new(stores);
let matrix = match metric {
DistanceMetric::BrayCurtis => CountPartials::bray_dist_matrix(&global),
DistanceMetric::RelfreqBrayCurtis => CountPartials::relfreq_bray_dist_matrix(&global),
DistanceMetric::Euclidean => CountPartials::euclidean_dist_matrix(&global),
DistanceMetric::RelfreqEuclidean => CountPartials::relfreq_euclidean_dist_matrix(&global),
DistanceMetric::Hellinger => CountPartials::hellinger_dist_matrix(&global),
DistanceMetric::HellingerEuclidean => CountPartials::hellinger_euclidean_dist_matrix(&global),
DistanceMetric::Jaccard => CountPartials::threshold_jaccard_dist_matrix(&global, presence_threshold),
DistanceMetric::Mash => CountPartials::threshold_mash_dist_matrix(&global, self.kmer_size(), presence_threshold),
DistanceMetric::Hamming => {
return Err(OKIError::InvalidInput(
"Hamming is only available for presence/absence indexes".into(),
));
}
};
let shared = if shared_kmers {
let (inter, _) = CountPartials::partial_threshold_jaccard(&global, presence_threshold);
Some(inter)
} else {
None
};
Ok(DistanceOutput { matrix, shared_kmers: shared })
} else {
let stores: Vec<_> = (0..n_parts)
.into_par_iter()
.map(|i| self.presence_store(i).map_err(OKIError::Partition))
.collect::<OKIResult<_>>()?;
let global = LayeredStore::new(stores);
let matrix = match metric {
DistanceMetric::Jaccard => BitPartials::jaccard_dist_matrix(&global),
DistanceMetric::Mash => BitPartials::mash_dist_matrix(&global, self.kmer_size()),
DistanceMetric::Hamming => {
BitPartials::hamming_dist_matrix(&global).mapv(|v| v as f64)
}
other => {
return Err(OKIError::InvalidInput(format!(
"{other:?} requires a count index; use --metric jaccard or --metric hamming"
)));
}
};
let shared = if shared_kmers {
let (inter, _) = BitPartials::partial_jaccard(&global);
Some(inter)
} else {
None
};
Ok(DistanceOutput { matrix, shared_kmers: shared })
}
}
}
src/obikindex/src/index/dump.rs
-123
View File
@@ -1,123 +0,0 @@
use std::io::Write;
use std::sync::atomic::{AtomicUsize, Ordering};
use rayon::prelude::*;
use crate::index::error::{OKIError, OKIResult};
use crate::index::kmer_index::KmerIndex;
use crate::index::KmerFilter;
impl KmerIndex {
/// Write a CSV table of all indexed kmers to `out`.
///
/// Columns: `kmer`, then one column per genome (in index order).
/// Values are counts (u32) when `use_counts = true`, otherwise 0/1.
///
/// `force_presence` overrides `with_counts`: even if the index stores counts,
/// the output uses 0/1 presence columns.
///
/// Partitions are scanned in parallel; each partition buffers its output locally
/// before the main thread writes the chunks in partition order.
///
/// The caller must have set the global kmer length (`obikseq::set_k`) before
/// calling this method.
pub fn dump<W: Write, F: Fn() + Send + Sync>(
&self,
out: &mut W,
force_presence: bool,
debug: bool,
head: Option<usize>,
filters: &[Box<dyn KmerFilter>],
on_partition: F,
) -> OKIResult<()> {
let genomes = self.meta.genomes().map_err(OKIError::Io)?;
let use_counts = self.meta.config.with_counts && !force_presence;
let n_genomes = genomes.len().max(1);
let kmer_size = self.kmer_size();
// ── Header ────────────────────────────────────────────────────────────
if debug {
write!(out, "partition,layer,")?;
}
write!(out, "kmer")?;
for g in &genomes {
write!(out, ",{}", g.label)?;
}
writeln!(out)?;
// ── Rows — parallel over partitions ───────────────────────────────────
let n = self.n_partitions();
let write_row = |buf: &mut Vec<u8>, row: &[u32], prefix: &str| {
let _ = buf.write_all(prefix.as_bytes());
for &v in row { let _ = write!(buf, ",{v}"); }
let _ = buf.write_all(b"\n");
};
let chunks: Vec<OKIResult<Vec<u8>>> = if let Some(limit) = head {
// ── Bounded: atomic counter, early exit when limit reached ────────
let remaining = AtomicUsize::new(limit);
(0..n).into_par_iter().map(|i| {
if remaining.load(Ordering::Relaxed) == 0 { return Ok(vec![]); }
let mut buf = Vec::<u8>::new();
let try_write = |buf: &mut Vec<u8>, row: &[u32], prefix: &str| -> bool {
match remaining.fetch_update(Ordering::SeqCst, Ordering::SeqCst, |cur| {
if cur > 0 { Some(cur - 1) } else { None }
}) {
Err(_) => false,
Ok(_) => { write_row(buf, row, prefix); true }
}
};
if debug {
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
.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)
})
.map_err(OKIError::Partition)?;
}
on_partition();
Ok(buf)
}).collect()
} else {
// ── Unbounded: no atomic, no contention ───────────────────────────
(0..n).into_par_iter().map(|i| {
let mut buf = Vec::<u8>::new();
if debug {
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}"));
true
})
.map_err(OKIError::Partition)?;
} else {
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);
true
})
.map_err(OKIError::Partition)?;
}
on_partition();
Ok(buf)
}).collect()
};
// ── Sequential write ──────────────────────────────────────────────────
for chunk in chunks {
out.write_all(&chunk?)?;
}
out.flush()?;
Ok(())
}
}
src/obikindex/src/index/dump_layer.rs
-201
View File
@@ -1,201 +0,0 @@
use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix};
use obikseq::CanonicalKmer;
use crate::layer::{MphfLayer, OLMError};
use obiskio::{SKError, SKResult, UnitigFileReader};
use crate::index::filter::{KmerFilter, passes_all};
use crate::index::kmer_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 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).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 mut layer = 0;
loop {
let layer_dir = self.layer_dir(part, layer);
if !layer_dir.exists() {
break;
}
let mphf = MphfLayer::open(&layer_dir).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)
}
}
src/obikindex/src/index/error.rs
+33 -14
View File
@@ -8,6 +8,7 @@ pub enum OKIError {
Io(io::Error),
Json(serde_json::Error),
Partition(SKError),
Layer(SKError),
/// Source index is not in `Indexed` state.
NotIndexed(std::path::PathBuf),
/// Source indexes have incompatible configurations (k, m, n_bits).
@@ -20,6 +21,9 @@ pub enum OKIError {
InvalidInput(String),
/// Sources mix exact and approximate evidence, or use incompatible approx parameters.
IncompatibleEvidence(String),
/// Malformed data encountered while reading/writing a layer file, tagged
/// with the operation that hit it.
InvalidData { context: &'static str, detail: String },
}
pub type OKIResult<T> = Result<T, OKIError>;
@@ -27,15 +31,24 @@ pub type OKIResult<T> = Result<T, OKIError>;
impl fmt::Display for OKIError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
OKIError::Io(e) => write!(f, "I/O error: {e}"),
OKIError::Json(e) => write!(f, "JSON error: {e}"),
OKIError::Partition(e) => write!(f, "partition error: {e}"),
OKIError::NotIndexed(p) => write!(f, "index not fully built: {}", p.display()),
OKIError::Io(e) => write!(f, "I/O error: {e}"),
OKIError::Json(e) => write!(f, "JSON error: {e}"),
OKIError::Partition(e) => write!(f, "partition error: {e}"),
OKIError::Layer(e) => write!(f, "layer error: {e}"),
OKIError::NotIndexed(p) => write!(f, "index not fully built: {}", p.display()),
OKIError::IncompatibleConfig => write!(f, "incompatible index configurations"),
OKIError::MismatchedMode => write!(f, "count mode requires all sources to have with_counts=true"),
OKIError::DuplicateGenomeLabel(l) => write!(f, "duplicate genome label across sources: {l}"),
OKIError::InvalidInput(m) => write!(f, "invalid input: {m}"),
OKIError::IncompatibleEvidence(m) => write!(f, "incompatible evidence: {m}"),
OKIError::MismatchedMode => write!(
f,
"count mode requires all sources to have with_counts=true"
),
OKIError::DuplicateGenomeLabel(l) => {
write!(f, "duplicate genome label across sources: {l}")
}
OKIError::InvalidInput(m) => write!(f, "invalid input: {m}"),
OKIError::IncompatibleEvidence(m) => write!(f, "incompatible evidence: {m}"),
OKIError::InvalidData { context, detail } => {
write!(f, "invalid data ({context}): {detail}")
}
}
}
}
@@ -43,22 +56,28 @@ impl fmt::Display for OKIError {
impl std::error::Error for OKIError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
OKIError::Io(e) => Some(e),
OKIError::Json(e) => Some(e),
OKIError::Io(e) => Some(e),
OKIError::Json(e) => Some(e),
OKIError::Partition(e) => Some(e),
_ => None, // IncompatibleConfig, MismatchedMode, DuplicateGenomeLabel
_ => None, // IncompatibleConfig, MismatchedMode, DuplicateGenomeLabel
}
}
}
impl From<io::Error> for OKIError {
fn from(e: io::Error) -> Self { OKIError::Io(e) }
fn from(e: io::Error) -> Self {
OKIError::Io(e)
}
}
impl From<serde_json::Error> for OKIError {
fn from(e: serde_json::Error) -> Self { OKIError::Json(e) }
fn from(e: serde_json::Error) -> Self {
OKIError::Json(e)
}
}
impl From<SKError> for OKIError {
fn from(e: SKError) -> Self { OKIError::Partition(e) }
fn from(e: SKError) -> Self {
OKIError::Partition(e)
}
}
src/obikindex/src/index/filter.rs
-298
View File
@@ -1,298 +0,0 @@
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
}
}
src/obikindex/src/index/graph_pipeline.rs
-152
View File
@@ -1,152 +0,0 @@
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 crate::layer::{IndexMode, TypedLayer};
use obiskio::{SKError, SKResult};
use crate::index::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 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 = TypedLayer::<()>::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 + `TypedLayer::<()>::build`.
///
/// Returns n_kmers.
pub 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");
TypedLayer::<()>::build(layer_dir, block_bits, evidence)
.map_err(|e| olm_to_sk(e, "graph pipeline"))?;
debug!("materialize_layer: MPHF build done");
Ok(n)
}
src/obikindex/src/index/kmer_index.rs
+76 -43
View File
@@ -2,13 +2,13 @@ use std::fs;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use crate::layer::meta::PartitionMeta;
use crate::layer::KmerLayer;
use crate::partition::KmerPartition;
use obisys::progress_bar;
use rayon::prelude::*;
use obikseq::{set_k, set_m};
use crate::index::common::load_meta;
use crate::index::error::{OKIError, OKIResult};
use crate::index::meta::{GenomeInfo, IndexConfig, IndexMeta};
@@ -28,12 +28,15 @@ impl KmerIndex {
genome_info: Option<GenomeInfo>,
) -> OKIResult<Self> {
let root_path = path.as_ref().to_owned();
fs::create_dir_all(&root_path).map_err(OKIError::Io)?;
fs::create_dir(&root_path).map_err(OKIError::Io)?;
set_k(config.kmer_size);
set_m(config.minimizer_size);
let genomes = genome_info.into_iter().collect();
let meta = IndexMeta::create_at(&root_path, config, genomes).map_err(OKIError::Io)?;
Ok(Self { root_path, meta: Arc::new(meta) })
Ok(Self {
root_path,
meta: Arc::new(meta),
})
}
pub fn open<P: AsRef<Path>>(path: P) -> OKIResult<Self> {
@@ -41,16 +44,30 @@ impl KmerIndex {
let meta = IndexMeta::open_at(&root_path).map_err(OKIError::Io)?;
set_k(meta.config.kmer_size);
set_m(meta.config.minimizer_size);
Ok(Self { root_path, meta: Arc::new(meta) })
Ok(Self {
root_path,
meta: Arc::new(meta),
})
}
/// Return `true` if `path` contains an `index.meta` file.
pub fn exists<P: AsRef<Path>>(path: P) -> bool {
/// Return `true` if `path` points to a valid obikmer index directory.
///
/// A directory is considered an index only when it contains an `index.meta`
/// file. The absence of that file means the directory is not an index, even
/// if it otherwise exists.
///
/// This is a static method (associated function), not an instance method.
/// Call it as `KmerIndex::is_an_index(path)`.
///
/// `path` is generic over any type implementing [`AsRef<Path>`]. This allows
/// callers to pass `&str`, `String`, `PathBuf`, `&Path`, etc. without an
/// explicit conversion; the value is coerced via `.as_ref()` inside.
pub fn is_an_index<P: AsRef<Path>>(path: P) -> bool {
IndexMeta::exists(path.as_ref())
}
/// The index's root directory.
pub fn root_path(&self) -> &Path {
pub fn dir(&self) -> &Path {
&self.root_path
}
pub fn meta(&self) -> Arc<IndexMeta> {
@@ -102,42 +119,48 @@ impl KmerIndex {
/// same directory only indirectly, through this method — `obikindexer`
/// depends on `KmerIndex`, not the other way around — see
/// `DevDocMD/implementation/partition_layer_cache.md`.
pub fn partition_dir(&self, i: usize) -> PathBuf {
crate::partition::partition_dir(&self.root_path, i)
pub(crate) fn partition_dir(&self, i: usize) -> OKIResult<PathBuf> {
Ok(self.partition(i)?.dir().to_path_buf())
}
pub fn partition(&self, i: usize) -> OKIResult<KmerPartition> {
let n = self.n_partitions();
if i >= n {
return Err(OKIError::InvalidData {
context: "New KmerPartition",
detail: format!("Partition {i} requested, Index has only {n} partitions"),
});
}
Ok(KmerPartition::new(self, i))
}
pub(crate) fn layer_dir(&self, i: usize, j: usize) -> OKIResult<PathBuf> {
Ok(self.layer(i, j)?.dir().to_path_buf())
}
pub fn layer(&self, i: usize, j: usize) -> OKIResult<KmerLayer> {
self.partition(i)?.layer(j)
}
/// Path of partition `i`'s layered-index directory (`<partition>/index`).
pub fn index_dir(&self, i: usize) -> PathBuf {
pub(crate) fn index_dir(&self, i: usize) -> PathBuf {
crate::partition::index_dir(&self.root_path, i)
}
/// Path of layer `l` within partition `i`'s layered index.
pub fn layer_dir(&self, i: usize, l: usize) -> PathBuf {
crate::layer::layer_dir(&self.index_dir(i), l)
/// Number of layers in partition `i` — a disk scan, not a metadata read
/// (see [`KmerPartition::n_layers`](crate::partition::KmerPartition::n_layers)).
pub fn n_layers(&self, i: usize) -> OKIResult<usize> {
Ok(self.partition(i)?.n_layers())
}
/// 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
/// the former `obikpartitionner` crate, 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<crate::layer::IndexMode> {
Ok(self.partition_meta(i)?.mode)
/// Evidence mode partition `i` was actually built with — detected from
/// layer 0's evidence files on disk (see
/// [`KmerPartition::mode`](crate::partition::KmerPartition::mode)).
/// 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) -> OKIResult<crate::layer::IndexMode> {
self.partition(i)?.mode()
}
/// Number of layers per partition.
@@ -150,8 +173,8 @@ impl KmerIndex {
}
/// Path to the unitigs file for partition `part`, layer `layer`.
pub fn layer_unitigs_path(&self, part: usize, layer: usize) -> PathBuf {
self.layer_dir(part, layer).join("unitigs.bin")
pub fn layer_unitigs_path(&self, part: usize, layer: usize) -> OKIResult<PathBuf> {
Ok(self.layer_dir(part, layer)?.join("unitigs.bin"))
}
/// Pack all partition matrices into single-file format (presence → .pbmx, counts → .pcmx).
@@ -164,13 +187,14 @@ impl KmerIndex {
/// on-disk format (see `DevDocMD/architecture/siblings.md`) — count
/// matrices are unaffected, sparse count matrices aren't implemented
/// (see the sparse-matrix design plan's "Explicitly deferred").
pub fn pack_matrices(&self, sparse: bool) -> OKIResult<()> {
/// TODO: To be moved somewhere else
pub(crate) fn pack_matrices(&self, sparse: bool) -> OKIResult<()> {
use obicompactvec::{pack_bit_matrix, pack_compact_int_matrix, pack_sparse_bit_matrix};
let n = self.n_partitions();
let order: Vec<usize> = (0..n).collect();
let pb = progress_bar("pack", n as u64, "partitions");
crate::index::numa::PartitionRunner::new().run(
obisys::PartitionRunner::new().run(
&order,
|i| -> OKIResult<()> {
let index_dir = self.index_dir(i);
@@ -179,7 +203,7 @@ impl KmerIndex {
}
let n_layers = self.n_layers(i)?;
for l in 0..n_layers {
let layer_dir = self.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() {
@@ -207,7 +231,8 @@ impl KmerIndex {
///
/// Old indexes were built before this file was required. The number of
/// kmers is recovered from `unitigs.bin`, which is always present.
pub fn upgrade_layer_meta(&self) -> OKIResult<()> {
/// TODO: Should not exist anymore
pub(crate) fn upgrade_layer_meta(&self) -> OKIResult<()> {
use obicompactvec::LayerMeta;
use obiskio::UnitigFileReader;
@@ -229,7 +254,15 @@ impl KmerIndex {
}
};
for l in 0..n_layers {
let layer_dir = self.layer_dir(i, l);
let layer_dir = match self.layer_dir(i, l) {
Ok(d) => d,
Err(e) => {
return Some(OKIError::Io(std::io::Error::new(
std::io::ErrorKind::Other,
e.to_string(),
)));
}
};
let meta_path = layer_dir.join(LayerMeta::FILENAME);
if meta_path.exists() {
continue;
src/obikindex/src/index/matrix_store.rs
-46
View File
@@ -1,46 +0,0 @@
use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix};
use crate::layer::{LayeredStore, open_data};
use obiskio::SKResult;
use crate::index::common::{load_meta, olm_to_sk};
use crate::index::kmer_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))
}
}
src/obikindex/src/index/merge.rs
-472
View File
@@ -1,472 +0,0 @@
use std::collections::HashMap;
use std::fs;
use std::io;
use std::path::Path;
use crate::index::builder::IndexBuilder;
use obisys::{Reporter, Stage, progress_bar, spinner};
use tracing::{debug, info};
use crate::layer::IndexMode;
use crate::index::error::{OKIError, OKIResult};
use crate::index::kmer_index::KmerIndex;
use crate::index::meta::{GenomeInfo, IndexMeta};
use crate::index::state::IndexState;
pub use crate::index::merge_layer::MergeMode;
// ── per-partition diagnostic record ──────────────────────────────────────────
#[derive(Debug)]
struct PartStat {
id: usize,
unitig_bytes: u64,
g_len: usize,
}
// ── main merge entry point ────────────────────────────────────────────────────
impl KmerIndex {
pub fn merge<P: AsRef<Path>>(
output: P,
sources: &[&KmerIndex],
mode: MergeMode,
force: bool,
rename_duplicates: bool,
budget_fraction: f64,
rep: &mut Reporter,
) -> OKIResult<Self> {
let output = output.as_ref();
if sources.is_empty() {
return Err(OKIError::Io(io::Error::new(
io::ErrorKind::InvalidInput,
"merge requires at least one source index",
)));
}
// ── Validate config compatibility ─────────────────────────────────────
let ref0 = sources[0];
for src in sources {
if src.state()? != IndexState::Indexed {
return Err(OKIError::NotIndexed(src.root_path.clone()));
}
if src.kmer_size() != ref0.kmer_size()
|| src.minimizer_size() != ref0.minimizer_size()
|| src.n_partitions() != ref0.n_partitions()
{
return Err(OKIError::IncompatibleConfig);
}
if mode == MergeMode::Count && !src.meta.config.with_counts {
return Err(OKIError::MismatchedMode);
}
}
// Read each source's genome list once — `IndexMeta::genomes` is a
// fresh disk read every call, so cache it rather than re-reading it
// repeatedly through the rest of this function.
let src_genomes: Vec<Vec<GenomeInfo>> =
sources.iter().map(|s| s.genomes()).collect::<OKIResult<_>>()?;
// ── Log source characteristics and choose base ────────────────────────
let mode_str = if mode == MergeMode::Presence {
"presence"
} else {
"count"
};
info!(
"merge: {} source(s), smer-size={}, mode={}",
sources.len(),
sources[0].kmer_size(),
mode_str,
);
for (i, src) in sources.iter().enumerate() {
let genome_str = if src_genomes[i].len() == 1 {
"mono-genome".to_string()
} else {
format!("{} genomes", src_genomes[i].len())
};
let trivial_str = if is_trivial(&src_genomes[i], mode) {
" [trivial: no data approximation]"
} else {
""
};
info!(
" [{}] {} — {}, {}, {}{}",
i,
src.root_path.display(),
format_evidence(&src.meta.config.evidence),
genome_str,
mode_str,
trivial_str,
);
}
let base_idx = choose_base(sources, &src_genomes, mode);
let needs_approx = sources.iter().enumerate().any(|(i, src)| {
!is_trivial(&src_genomes[i], mode)
&& matches!(
src.meta.config.evidence,
IndexMode::Approx { .. } | IndexMode::Hybrid { .. }
)
});
info!(
"output evidence: {} ({}base: [{}] {})",
format_evidence(&sources[base_idx].meta.config.evidence),
if needs_approx {
"forced approx — "
} else {
""
},
base_idx,
sources[base_idx].root_path.display(),
);
let mut ordered: Vec<&KmerIndex> = Vec::with_capacity(sources.len());
let mut ordered_genomes: Vec<Vec<GenomeInfo>> = Vec::with_capacity(sources.len());
ordered.push(sources[base_idx]);
ordered_genomes.push(src_genomes[base_idx].clone());
for (i, &src) in sources.iter().enumerate() {
if i != base_idx {
ordered.push(src);
ordered_genomes.push(src_genomes[i].clone());
}
}
let sources: &[&KmerIndex] = &ordered;
let src_genomes = ordered_genomes;
let evidence = sources[0].meta.config.evidence.clone();
// ── Compute final genome labels ────────────────────────────────────────
let (source_labels, all_genomes) = compute_labels(&src_genomes, rename_duplicates)?;
// ── Prepare output directory ──────────────────────────────────────────
KmerIndex::clear_output_for_create(output, force)?;
// ── Bootstrap: copy first source to output ────────────────────────────
info!(
"bootstrap: copying {} → {} ({} genome(s))",
sources[0].root_path.display(),
output.display(),
src_genomes[0].len(),
);
let t = Stage::start("bootstrap");
let pb = spinner("bootstrap");
pb.set_message("copying index …");
copy_dir_all(&sources[0].root_path, output)?;
let dst_meta = IndexMeta::open_at(output).map_err(OKIError::Io)?;
let mut config = dst_meta.config.clone();
config.with_counts = mode == MergeMode::Count;
config.evidence = evidence.clone();
dst_meta.rewrite_config(config, all_genomes).map_err(OKIError::Io)?;
if mode == MergeMode::Presence {
remove_dirs_named(output, "counts")?;
}
pb.finish_and_clear();
rep.push(t.stop());
// ── Rebuild spectrums ─────────────────────────────────────────────────
info!("rebuilding spectrums for {} source(s)", sources.len());
let t = Stage::start("spectrums");
let pb = spinner("spectrums");
pb.set_message("copying …");
let spectrums_dir = output.join("spectrums");
if spectrums_dir.exists() {
fs::remove_dir_all(&spectrums_dir)?;
}
for ((src, new_labels), genomes) in sources.iter().zip(&source_labels).zip(&src_genomes) {
let old_labels: Vec<String> = genomes.iter().map(|g| g.label.clone()).collect();
copy_spectrums(&src.root_path, output, &old_labels, new_labels)?;
}
pb.finish_and_clear();
rep.push(t.stop());
// ── Open destination ──────────────────────────────────────────────────
let dst = KmerIndex::open(output)?;
let n_partitions = dst.n_partitions();
let n_dst_genomes = src_genomes[0].len();
// ── Merge partitions ──────────────────────────────────────────────────
let remaining_sources: Vec<&KmerIndex> = sources[1..].to_vec();
if !remaining_sources.is_empty() {
let n_src_genomes: usize = src_genomes[1..].iter().map(|g| g.len()).sum();
info!(
"merging {} partition(s) × {} additional source genome(s) into {} destination genome(s)",
n_partitions, n_src_genomes, n_dst_genomes,
);
let t = Stage::start("merge_partitions");
let pb = progress_bar("merge", n_partitions as u64, "partitions");
let block_bits = dst.meta.config.block_bits;
// Pre-build source list once (avoid rebuilding per partition)
let srcs: Vec<(&KmerIndex, usize)> = remaining_sources
.iter()
.zip(&src_genomes[1..])
.map(|(s, g)| (*s, g.len()))
.collect();
// Per-partition unitig byte sizes across remaining sources (stat() only)
let partition_sizes: Vec<u64> = (0..n_partitions)
.map(|i| {
remaining_sources
.iter()
.map(|s| partition_unitig_bytes(s, i))
.sum()
})
.collect();
// LFD sort: largest partition first
let mut order: Vec<usize> = (0..n_partitions).collect();
order.sort_unstable_by_key(|&i| std::cmp::Reverse(partition_sizes[i]));
let _ = budget_fraction; // kept in signature for CLI compatibility
// Shadow as references so closures can capture them by copy.
let srcs = &srcs;
let evidence = &evidence;
let runner = crate::index::numa::PartitionRunner::new();
let mut part_stats: Vec<PartStat> = Vec::with_capacity(n_partitions);
runner
.run(
&order,
|i| {
dst.merge_partition(
i,
srcs,
mode,
n_dst_genomes,
block_bits,
evidence,
)
},
|i, g_len, dur| {
pb.inc(1);
debug!(
"partition {i}: done in {:.1}s — {} new kmers",
dur.as_secs_f64(),
g_len,
);
part_stats.push(PartStat {
id: i,
unitig_bytes: partition_sizes[i],
g_len,
});
},
)
.map_err(OKIError::Partition)?;
pb.finish_and_clear();
// ── Diagnostic report ─────────────────────────────────────────────
print_merge_partition_report(&part_stats, runner.max_workers());
rep.push(t.stop());
}
// ── Pack matrices after merge ─────────────────────────────────────────
{
let t = Stage::start("pack");
let pb = spinner("pack");
pb.set_message("consolidating column files …");
let dst2 = KmerIndex::open(output)?;
dst2.pack_matrices(false)?;
dst2.meta.mark_indexed().map_err(OKIError::Io)?;
pb.finish_and_clear();
rep.push(t.stop());
}
KmerIndex::open(output)
}
}
// ── Diagnostic report ─────────────────────────────────────────────────────────
fn print_merge_partition_report(stats: &[PartStat], max_workers: usize) {
let total_new: usize = stats.iter().map(|s| s.g_len).sum();
let non_empty = stats.iter().filter(|s| s.unitig_bytes > 0).count();
if non_empty == 0 {
info!("merge_partitions report: no data (all partitions empty)");
return;
}
info!("─── merge_partitions report ───");
info!(
" {} partition(s) processed, {} total new kmers",
non_empty, total_new,
);
info!(" max workers: {max_workers}");
// Top 8 partitions by new-kmer count
let mut by_new: Vec<&PartStat> = stats.iter().filter(|s| s.g_len > 0).collect();
by_new.sort_by_key(|s| std::cmp::Reverse(s.g_len));
if !by_new.is_empty() {
info!(" top partitions by new kmers:");
for s in by_new.iter().take(8) {
info!(
" partition {:4} : {}M new kmers ({} unitig bytes)",
s.id,
s.g_len / 1_000_000,
fmt_bytes(s.unitig_bytes),
);
}
}
info!("───────────────────────────────");
}
// ── helpers ───────────────────────────────────────────────────────────────────
fn fmt_bytes(b: u64) -> String {
if b >= 1 << 30 {
format!("{:.1} GB", b as f64 / (1u64 << 30) as f64)
} else if b >= 1 << 20 {
format!("{:.1} MB", b as f64 / (1u64 << 20) as f64)
} else if b >= 1 << 10 {
format!("{:.1} KB", b as f64 / (1u64 << 10) as f64)
} else {
format!("{b} B")
}
}
/// Sum of all unitigs.bin sizes across all layers of partition `i` in `src`.
fn partition_unitig_bytes(src: &KmerIndex, i: usize) -> u64 {
let mut total = 0u64;
for l in 0.. {
let p = src.layer_unitigs_path(i, l);
if !p.exists() {
break;
}
if let Ok(m) = std::fs::metadata(&p) {
total += m.len();
}
}
total
}
fn compute_labels(
src_genomes: &[Vec<GenomeInfo>],
rename_duplicates: bool,
) -> OKIResult<(Vec<Vec<String>>, Vec<GenomeInfo>)> {
let mut seen: HashMap<String, usize> = HashMap::new();
let mut source_labels: Vec<Vec<String>> = Vec::with_capacity(src_genomes.len());
let mut all_genomes: Vec<GenomeInfo> = Vec::new();
for genomes in src_genomes {
let mut labels = Vec::with_capacity(genomes.len());
for genome in genomes {
let label = &genome.label;
let count = seen.entry(label.clone()).or_insert(0);
let new_label = if *count == 0 {
label.clone()
} else if rename_duplicates {
format!("{label}.{count}")
} else {
return Err(OKIError::DuplicateGenomeLabel(label.clone()));
};
*count += 1;
labels.push(new_label.clone());
all_genomes.push(GenomeInfo {
label: new_label,
meta: genome.meta.clone(),
});
}
source_labels.push(labels);
}
Ok((source_labels, all_genomes))
}
fn copy_spectrums(
src_root: &Path,
dst_root: &Path,
old_labels: &[String],
new_labels: &[String],
) -> io::Result<()> {
let src_dir = src_root.join("spectrums");
let dst_dir = dst_root.join("spectrums");
fs::create_dir_all(&dst_dir)?;
for (old, new) in old_labels.iter().zip(new_labels.iter()) {
let src_file = src_dir.join(format!("{old}.json"));
if src_file.exists() {
fs::copy(&src_file, dst_dir.join(format!("{new}.json")))?;
}
}
Ok(())
}
fn remove_dirs_named(root: &Path, name: &str) -> io::Result<()> {
for entry in fs::read_dir(root)? {
let entry = entry?;
let path = entry.path();
if path.is_dir() {
if path.file_name().and_then(|n| n.to_str()) == Some(name) {
fs::remove_dir_all(&path)?;
} else {
remove_dirs_named(&path, name)?;
}
}
}
Ok(())
}
fn format_evidence(ev: &IndexMode) -> String {
match ev {
IndexMode::Exact => "exact".to_string(),
IndexMode::Approx { b, z } => format!("approx (b={b}, z={z})"),
IndexMode::Hybrid { b, z } => format!("hybrid (b={b}, z={z})"),
}
}
fn is_trivial(genomes: &[GenomeInfo], mode: MergeMode) -> bool {
genomes.len() == 1 && mode == MergeMode::Presence
}
fn index_unitig_size(src: &KmerIndex) -> u64 {
let n = src.n_partitions();
(0..n).map(|i| partition_unitig_bytes(src, i)).sum()
}
fn choose_base(sources: &[&KmerIndex], src_genomes: &[Vec<GenomeInfo>], mode: MergeMode) -> usize {
let needs_approx = sources.iter().enumerate().any(|(i, src)| {
!is_trivial(&src_genomes[i], mode)
&& matches!(
src.meta.config.evidence,
IndexMode::Approx { .. } | IndexMode::Hybrid { .. }
)
});
sources
.iter()
.enumerate()
.filter(|(_, src)| {
!needs_approx
|| matches!(
src.meta.config.evidence,
IndexMode::Approx { .. } | IndexMode::Hybrid { .. }
)
})
.max_by_key(|(_, src)| index_unitig_size(src))
.map(|(i, _)| i)
.unwrap()
}
fn copy_dir_all(src: &Path, dst: &Path) -> io::Result<()> {
fs::create_dir_all(dst)?;
for entry in fs::read_dir(src)? {
let entry = entry?;
let src_path = entry.path();
let dst_path = dst.join(entry.file_name());
if src_path.is_dir() {
copy_dir_all(&src_path, &dst_path)?;
} else {
fs::copy(&src_path, &dst_path)?;
}
}
Ok(())
}
src/obikindex/src/index/merge_layer/mod.rs
-598
View File
@@ -1,598 +0,0 @@
//! 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 crate::layer::{IndexMode, TypedLayer, LayeredMap, MphfOnly, layer_dir};
use obiskio::{SKError, SKResult, UnitigFileReader};
use crate::index::common::{ColBuilder, load_meta, olm_to_sk};
use crate::index::graph_pipeline::{build_graph, materialize_layer};
use crate::index::kmer_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 {
TypedLayer::<()>::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)
}
}
src/obikindex/src/index/merge_layer/src_layer.rs
-95
View File
@@ -1,95 +0,0 @@
use std::path::Path;
use obicompactvec::{MatrixGroupOps, PersistentBitMatrix, PersistentCompactIntMatrix};
use obikseq::CanonicalKmer;
use crate::layer::MphfOnly;
use obiskio::{SKError, SKResult};
use crate::index::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),
}
}
}
src/obikindex/src/index/meta.rs
+71 -35
View File
@@ -21,8 +21,44 @@ pub struct GenomeInfo {
}
impl GenomeInfo {
/// Panics if `label` fails [`validate_label`](Self::validate_label) —
/// an invalid genome label is a programmer error to surface loudly,
/// not a recoverable condition for this constructor to paper over.
/// Callers that need a graceful error (e.g. a CLI parsing user input)
/// should call [`validate_label`](Self::validate_label) themselves
/// first and report it their own way.
pub fn new(label: impl Into<String>) -> Self {
Self { label: label.into(), meta: HashMap::new() }
let label = label.into();
Self::validate_label(&label).expect("invalid genome label");
Self {
label,
meta: HashMap::new(),
}
}
/// Validate a user-supplied genome label.
///
/// Forbidden: `/` (filesystem separator), `=` (--new-label parser separator),
/// `\0` (null byte), `\n`, `\r`, `\t` (break CSV output).
/// Empty labels are also rejected.
pub fn validate_label(label: &str) -> Result<(), String> {
if label.is_empty() {
return Err("genome label must not be empty".into());
}
const FORBIDDEN: &[char] = &['/', '=', '\0', '\n', '\r', '\t'];
if let Some(c) = label.chars().find(|c| FORBIDDEN.contains(c)) {
let display = match c {
'\0' => "\\0 (null)".to_string(),
'\n' => "\\n (newline)".to_string(),
'\r' => "\\r (carriage return)".to_string(),
'\t' => "\\t (tab)".to_string(),
c => format!("'{c}'"),
};
return Err(format!(
"genome label contains forbidden character {display}"
));
}
Ok(())
}
}
@@ -47,7 +83,7 @@ pub struct IndexConfig {
/// On-disk shape of `index.meta` — the whole file, read/written atomically
/// as one JSON blob by every `IndexMeta` operation below.
#[derive(Debug, Clone, Serialize, Deserialize)]
struct OnDiskMeta {
struct IndexMetadata {
version: u32,
config: IndexConfig,
#[serde(default)]
@@ -81,21 +117,29 @@ impl IndexMeta {
/// Create a brand-new `index.meta` for `index`, from fixed config and
/// an initial genome list (often empty — genomes are usually added
/// later via [`push_genome`](Self::push_genome)).
pub fn create(index: &KmerIndex, config: IndexConfig, genomes: Vec<GenomeInfo>) -> io::Result<Self> {
Self::create_at(index.root_path(), config, genomes)
pub fn create(
index: &KmerIndex,
config: IndexConfig,
genomes: Vec<GenomeInfo>,
) -> io::Result<Self> {
Self::create_at(index.dir(), config, genomes)
}
/// Same as [`create`](Self::create), taking a raw path — for
/// `KmerIndex::create` itself, which doesn't have a complete
/// `KmerIndex` yet to hand in (it's what's being built).
pub(crate) fn create_at(root_path: &Path, config: IndexConfig, genomes: Vec<GenomeInfo>) -> io::Result<Self> {
pub(crate) fn create_at(
root_path: &Path,
config: IndexConfig,
genomes: Vec<GenomeInfo>,
) -> io::Result<Self> {
let meta = Self {
root_path: root_path.to_owned(),
version: META_VERSION,
config,
lock: RwLock::new(()),
};
meta.write_full(&OnDiskMeta {
meta.write_full(&IndexMetadata {
version: meta.version,
config: meta.config.clone(),
genomes,
@@ -106,7 +150,7 @@ impl IndexMeta {
/// Reopen the `index.meta` already on disk for `index`.
pub fn open(index: &KmerIndex) -> io::Result<Self> {
Self::open_at(index.root_path())
Self::open_at(index.dir())
}
pub(crate) fn open_at(root_path: &Path) -> io::Result<Self> {
@@ -119,7 +163,12 @@ impl IndexMeta {
})
}
pub fn exists(root: &Path) -> bool {
/// Returns whether the metadata file exists in `root`.
///
/// This is a static method (associated function), not an instance method.
/// It only checks for the presence of the `index.meta` file inside the
/// given directory; it does not validate its contents.
pub(crate) fn exists(root: &Path) -> bool {
root.join(META_FILENAME).exists()
}
@@ -177,10 +226,20 @@ impl IndexMeta {
/// (e.g. `self.meta = Arc::new(IndexMeta::open(self)?)`) — this method
/// only updates the file, it has no way to reach back into whatever
/// `KmerIndex` holds it.
pub fn rewrite_config(&self, config: IndexConfig, genomes: Vec<GenomeInfo>) -> io::Result<()> {
/// TODO: This methode is strange
pub(crate) fn rewrite_config(
&self,
config: IndexConfig,
genomes: Vec<GenomeInfo>,
) -> io::Result<()> {
let _guard = self.lock.write().unwrap();
let state = Self::read_full(&self.root_path)?.state;
self.write_full(&OnDiskMeta { version: self.version, config, genomes, state })
self.write_full(&IndexMetadata {
version: self.version,
config,
genomes,
state,
})
}
pub fn set_state(&self, state: IndexState) -> io::Result<()> {
@@ -221,12 +280,12 @@ impl IndexMeta {
// ── private ──────────────────────────────────────────────────────────────
fn read_full(root: &Path) -> io::Result<OnDiskMeta> {
fn read_full(root: &Path) -> io::Result<IndexMetadata> {
let file = fs::File::open(root.join(META_FILENAME))?;
serde_json::from_reader(file).map_err(io::Error::other)
}
fn write_full(&self, on_disk: &OnDiskMeta) -> io::Result<()> {
fn write_full(&self, on_disk: &IndexMetadata) -> io::Result<()> {
let file = fs::File::create(self.root_path.join(META_FILENAME))?;
serde_json::to_writer_pretty(file, on_disk).map_err(io::Error::other)
}
@@ -248,26 +307,3 @@ fn label_from_path(path: &Path) -> String {
s
}
}
/// Validate a user-supplied genome label.
///
/// Forbidden: `/` (filesystem separator), `=` (--new-label parser separator),
/// `\0` (null byte), `\n`, `\r`, `\t` (break CSV output).
/// Empty labels are also rejected.
pub fn validate_label(label: &str) -> Result<(), String> {
if label.is_empty() {
return Err("genome label must not be empty".into());
}
const FORBIDDEN: &[char] = &['/', '=', '\0', '\n', '\r', '\t'];
if let Some(c) = label.chars().find(|c| FORBIDDEN.contains(c)) {
let display = match c {
'\0' => "\\0 (null)".to_string(),
'\n' => "\\n (newline)".to_string(),
'\r' => "\\r (carriage return)".to_string(),
'\t' => "\\t (tab)".to_string(),
c => format!("'{c}'"),
};
return Err(format!("genome label contains forbidden character {display}"));
}
Ok(())
}
src/obikindex/src/index/mod.rs
+2 -28
View File
@@ -1,39 +1,13 @@
pub mod error;
pub mod meta;
pub mod predicate;
pub mod state;
mod builder;
mod common;
mod distance;
mod dump;
mod dump_layer;
pub mod filter;
mod graph_pipeline;
mod kmer_index;
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 builder::IndexBuilder;
pub use common::olm_to_sk;
pub use graph_pipeline::{materialize_layer, write_graph_as_unitigs};
pub use distance::{DistanceMetric, DistanceOutput};
pub use filter::{GroupQuorumFilter, KmerFilter, passes_all};
pub use common::ColBuilder;
pub use kmer_index::KmerIndex;
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 meta::{GenomeInfo, IndexConfig, IndexMeta, META_FILENAME};
pub use state::IndexState;
pub use stats::IndexBitsPerKmer;
pub use numa::PartitionRunner;
src/obikindex/src/index/numa/mod.rs
-17
View File
@@ -1,17 +0,0 @@
//! NUMA-aware partition runner via hwlocality.
//!
//! Detects NUMA topology using hwloc (cross-platform: Linux, macOS, etc.) and
//! builds one Rayon ThreadPool per NUMA node with threads pinned to that node's
//! CPUs. Linux first-touch policy then places graph allocations in local DRAM
//! automatically — no explicit memory binding needed.
//!
//! UMA systems (single socket, Apple Silicon, etc.) are the degenerate case:
//! one synthetic node containing all cores, no pool, no pinning.
//!
//! Submodules: [`topology`] (NUMA detection, per-node pools, thread pinning),
//! [`runner`] ([`PartitionRunner`], the adaptive worker-activation scheduler).
mod runner;
mod topology;
pub use runner::PartitionRunner;
src/obikindex/src/index/numa/runner.rs
-405
View File
@@ -1,405 +0,0 @@
use std::sync::Arc;
use std::time::{Duration, Instant};
use crossbeam_channel::unbounded;
use obisys::{CpuSample, IoSample};
use tracing::debug;
use super::topology::{build, pin_current_thread};
// ── PartitionRunner ─────────────────────────────────────────────────────────
/// Growth step (fraction of a node's worker capacity added per activation
/// event, see [`NodeActivation::grow`]).
const GROWTH_DIVISOR: usize = 8;
/// Minimum CPU efficiency growth to activate more workers, as a fraction of
/// the size of the *last growth step* (e.g. `0.2` after adding 8 workers
/// requires the next check to show at least +1.6 cores of growth — 20 % of
/// the ~8 cores those 8 workers should contribute if the workload is truly
/// CPU-bound). Scaling by the last step's size — not the cumulative total —
/// keeps the bar meaningful regardless of how many workers are already
/// active, instead of demanding an ever-larger absolute jump as the pool
/// grows.
const CPU_SPAWN_THRESHOLD: f64 = 0.2;
/// Minimum I/O throughput growth (relative) to activate more workers.
const IO_SPAWN_THRESHOLD: f64 = 0.2;
struct NodeConfig {
pool: Option<Arc<rayon::ThreadPool>>,
cpu_ids: Vec<usize>,
max_workers: usize,
}
/// Generic NUMA-aware runner for partition-level parallel work.
///
/// Workers are distributed evenly across NUMA nodes and pinned to their
/// node's CPUs. UMA is the degenerate case: one node, no pinning.
///
/// Workers are pre-spawned dormant, one activation channel per node so
/// growth always targets a specific node rather than whichever dormant
/// worker happens to wake up first on a shared channel. Growth (both the
/// initial count and each subsequent step) is expressed as a fraction of
/// `workers_per_node`, applied identically to every node, so the pace of
/// ramp-up depends on node size rather than node count — a single-NUMA-node
/// (UMA) machine ramps just as fast as an 8-node one.
///
/// # Termination
///
/// ```text
/// drop(part_tx) → part_rx drains → workers exit → drop their result_tx
/// drop(result_tx) → result_rx closes → controller loop exits
/// drop(activate_txs) → dormant workers exit cleanly
/// ```
pub struct PartitionRunner {
nodes: Vec<NodeConfig>,
}
impl PartitionRunner {
/// Total worker slots across all nodes.
pub fn max_workers(&self) -> usize {
self.nodes.iter().map(|n| n.max_workers).sum()
}
/// Detect topology and build. Always succeeds.
pub fn new() -> Self {
let ns = build();
let wpn = ns.workers_per_node();
debug!(
"PartitionRunner: {} node(s) × {} worker(s)/node max",
ns.pools.len(),
wpn,
);
let nodes = ns
.pools
.into_iter()
.zip(ns.cpus_per_node)
.map(|(pool, cpu_ids)| NodeConfig {
pool,
cpu_ids,
max_workers: wpn,
})
.collect();
Self { nodes }
}
/// Like [`new`](Self::new), but caps total worker slots (summed across
/// nodes) at `max_total_workers` — split evenly across nodes, each
/// further capped by that node's actual core count. For callers whose
/// own per-worker closure does further internal parallel work (so the
/// natural per-node core count would oversubscribe if used as the
/// *outer* degree of parallelism too).
pub fn new_capped(max_total_workers: usize) -> Self {
let ns = build();
let n_nodes = ns.pools.len().max(1);
let per_node_cap = (max_total_workers / n_nodes).max(1);
debug!(
"PartitionRunner (capped): {} node(s) × up to {} worker(s)/node ({} total requested)",
n_nodes, per_node_cap, max_total_workers,
);
let nodes = ns
.pools
.into_iter()
.zip(ns.cpus_per_node)
.map(|(pool, cpu_ids)| {
let node_cores = cpu_ids.len().max(1);
NodeConfig { pool, cpu_ids, max_workers: per_node_cap.min(node_cores) }
})
.collect();
Self { nodes }
}
/// Run `f(i)` for every index in `order`.
///
/// Workers are pre-spawned dormant and activated adaptively, per node:
/// `(workers_per_node / INITIAL_DIVISOR).max(1)` are woken immediately on
/// every node, then `(workers_per_node / GROWTH_DIVISOR).max(1)` more per
/// node each time the check below fires. A timer thread fires that check
/// every `TIMER_SECS` seconds; each completed partition resets that timer
/// (forcing an immediate check) and also triggers its own inline check. A
/// growth step happens whenever CPU efficiency grows by at least
/// `CPU_SPAWN_THRESHOLD` of what the last growth step should have
/// contributed, or I/O throughput grows by at least `IO_SPAWN_THRESHOLD`
/// (relative) since the last check — whichever resource is the actual
/// bottleneck still shows headroom.
///
/// `on_done(i, result, elapsed)` is called from the controller thread as
/// each partition completes — suitable for progress bars and result
/// aggregation.
///
/// Returns the first error produced by `f`, if any.
pub fn run<F, R, E, C>(&self, order: &[usize], f: F, mut on_done: C) -> Result<(), E>
where
F: Fn(usize) -> Result<R, E> + Send + Sync,
R: Send,
E: Send,
C: FnMut(usize, R, Duration) + Send,
{
let n_total = order.len();
if n_total == 0 {
return Ok(());
}
const TIMER_SECS: u64 = 30;
const INITIAL_DIVISOR: usize = 4;
// ── Channels ──────────────────────────────────────────────────────────
let (part_tx, part_rx) = unbounded::<usize>();
// reset_tx: controller → timer ("reset the 30 s window")
let (reset_tx, reset_rx) = unbounded::<()>();
// event_tx: workers + timer → controller (unified event stream)
let (event_tx, event_rx) = unbounded::<WorkerEvent<R, E>>();
// One activation channel per node: growth always targets a specific
// node, rather than whichever dormant worker happens to win the race
// on a channel shared across all nodes.
let (activate_txs, activate_rxs): (Vec<_>, Vec<_>) =
(0..self.nodes.len()).map(|_| unbounded::<()>()).unzip();
for &i in order {
part_tx.send(i).ok();
}
drop(part_tx);
let max_workers = self.max_workers();
let node_caps: Vec<usize> = self.nodes.iter().map(|n| n.max_workers).collect();
let f = &f;
let mut first_err: Option<E> = None;
std::thread::scope(|s| {
// ── Timer thread ──────────────────────────────────────────────────
// Sends TimerTick every TIMER_SECS seconds. Resets its window each
// time reset_rx receives a message (i.e. on partition completion).
let timer_tx = event_tx.clone();
s.spawn(move || {
let period = Duration::from_secs(TIMER_SECS);
loop {
crossbeam_channel::select! {
recv(reset_rx) -> r => {
if r.is_err() { break; } // reset_tx dropped → exit
}
default(period) => {
if timer_tx.send(WorkerEvent::TimerTick).is_err() { break; }
}
}
}
});
// ── Pre-spawn workers dormant, grouped by node ────────────────────
// Each worker listens on its own node's activation channel only.
for (node, arx) in self.nodes.iter().zip(activate_rxs.iter()) {
let cpu_ids = &node.cpu_ids;
for _ in 0..node.max_workers {
let prx = part_rx.clone();
let etx = event_tx.clone();
let arx = arx.clone();
let pool = node.pool.clone();
s.spawn(move || {
let tid = std::thread::current().id();
debug!(?tid, "PartitionRunner worker: waiting on activation");
if arx.recv().is_err() {
debug!(?tid, "PartitionRunner worker: activation channel closed, exiting");
return;
}
debug!(?tid, "PartitionRunner worker: activated");
if !cpu_ids.is_empty() {
pin_current_thread(cpu_ids);
}
for i in &prx {
debug!(?tid, partition = i, "PartitionRunner worker: picked partition");
let t = Instant::now();
let r = match &pool {
Some(p) => p.install(|| f(i)),
None => f(i),
};
debug!(?tid, partition = i, "PartitionRunner worker: partition done");
etx.send(WorkerEvent::Completed(i, r, t.elapsed())).ok();
}
debug!(?tid, "PartitionRunner worker: no more partitions, exiting");
});
}
}
// Drop controller's event_tx: event_rx closes when all workers +
// timer have exited.
drop(event_tx);
// ── Controller ────────────────────────────────────────────────────
let mut activation = NodeActivation::new(&activate_txs, &node_caps, max_workers);
activation.activate_initial(INITIAL_DIVISOR, n_total);
debug!(n_total, activated = activation.total(), "PartitionRunner controller: initial activation");
let mut cpu_sample = CpuSample::now();
let mut io_sample = IoSample::now();
let mut completed = 0usize;
while completed < n_total {
debug!(completed, n_total, "PartitionRunner controller: waiting for an event");
let Ok(event) = event_rx.recv() else {
debug!("PartitionRunner controller: event channel closed, stopping");
break;
};
match event {
WorkerEvent::Completed(i, r, dur) => {
match r {
Ok(v) => on_done(i, v, dur),
Err(e) => {
if first_err.is_none() {
first_err = Some(e);
}
}
}
completed += 1;
// Reset the 30 s timer.
reset_tx.send(()).ok();
// Inline check: same logic as a timer tick.
maybe_activate(
&mut activation,
&mut cpu_sample,
&mut io_sample,
completed,
n_total,
);
}
WorkerEvent::TimerTick => {
maybe_activate(
&mut activation,
&mut cpu_sample,
&mut io_sample,
completed,
n_total,
);
}
}
}
// Dormant workers exit once every sender for their node's channel
// is dropped — `activate_txs` holds the only ones.
drop(activate_txs);
// Timer thread exits when reset_tx closes.
drop(reset_tx);
});
match first_err {
Some(e) => Err(e),
None => Ok(()),
}
}
}
// ── Internal event type ───────────────────────────────────────────────────────
enum WorkerEvent<R, E> {
Completed(usize, Result<R, E>, Duration),
TimerTick,
}
/// Tracks how many of each node's dormant workers have been woken, and
/// grows every node by the same amount at each step (capped by that node's
/// remaining dormant workers and by the run's total budget) so load stays
/// balanced across nodes at every point in time — never just "one more
/// worker somewhere". Also remembers the size of the last real growth step
/// (`last_step`), used to scale the CPU activation threshold to what that
/// step could plausibly have contributed (see `maybe_activate`).
struct NodeActivation<'a> {
txs: &'a [crossbeam_channel::Sender<()>],
caps: &'a [usize],
active: Vec<usize>,
total: usize,
max: usize,
last_step: usize,
}
impl<'a> NodeActivation<'a> {
fn new(txs: &'a [crossbeam_channel::Sender<()>], caps: &'a [usize], max: usize) -> Self {
Self {
txs,
caps,
active: vec![0; txs.len()],
total: 0,
max,
last_step: 0,
}
}
fn total(&self) -> usize {
self.total
}
fn last_step(&self) -> usize {
self.last_step
}
fn max(&self) -> usize {
self.max
}
fn is_full(&self) -> bool {
self.total >= self.max
}
/// Wake up to `(node_cap / divisor).max(1)` dormant workers on every
/// node, capped by `n_total`. Called once at startup, unconditionally.
fn activate_initial(&mut self, divisor: usize, n_total: usize) {
self.grow(divisor, n_total);
}
/// Same per-node sizing as [`activate_initial`](Self::activate_initial),
/// applied as a growth step. Returns the number of workers actually
/// activated (may be less than requested once a node or the total
/// budget is exhausted). Updates `last_step` when it actually grew.
fn grow(&mut self, divisor: usize, n_total: usize) -> usize {
let before = self.total;
for idx in 0..self.txs.len() {
let wanted = (self.caps[idx] / divisor).max(1);
let room = self.caps[idx].saturating_sub(self.active[idx]);
let grow = wanted.min(room).min(n_total.saturating_sub(self.total));
for _ in 0..grow {
self.txs[idx].send(()).ok();
}
self.active[idx] += grow;
self.total += grow;
}
let grew = self.total - before;
if grew > 0 {
self.last_step = grew;
}
grew
}
}
fn maybe_activate(
activation: &mut NodeActivation,
cpu_sample: &mut CpuSample,
io_sample: &mut IoSample,
completed: usize,
n_total: usize,
) {
if activation.is_full() || completed >= n_total {
return;
}
// Expect roughly 1 core of extra efficiency per worker activated in the
// last growth step (CPU-bound case); require at least CPU_SPAWN_THRESHOLD
// (20 %) of that expected gain before growing again. Scaling by the last
// step's size — not the cumulative total — keeps the bar meaningful
// regardless of how many workers are already active: growing by 8 should
// always take ~+1.6 cores to confirm, whether that's the 2nd growth step
// or the 20th.
let cpu_threshold = CPU_SPAWN_THRESHOLD * activation.last_step() as f64;
// Call both unconditionally (no `||` short-circuit): each sampler must
// advance its own window every tick, regardless of what the other one
// reports, or it would starve behind whichever signal fires first.
let cpu_wants_more = cpu_sample.do_i_activate(cpu_threshold);
let io_wants_more = io_sample.do_i_activate(IO_SPAWN_THRESHOLD * activation.last_step() as f64);
if !(cpu_wants_more || io_wants_more) {
return;
}
let grew = activation.grow(GROWTH_DIVISOR, n_total);
if grew > 0 {
debug!(
"activated {} worker(s) — {}/{} active",
grew,
activation.total(),
activation.max()
);
}
}
src/obikindex/src/index/numa/topology.rs
-121
View File
@@ -1,121 +0,0 @@
use std::sync::Arc;
#[cfg(feature = "numa")]
use hwlocality::Topology;
#[cfg(feature = "numa")]
use hwlocality::cpu::binding::CpuBindingFlags;
#[cfg(feature = "numa")]
use hwlocality::cpu::cpuset::CpuSet;
#[cfg(feature = "numa")]
use hwlocality::object::types::ObjectType;
use tracing::debug;
// ── Public interface ──────────────────────────────────────────────────────────
pub struct NumaSetup {
/// One entry per NUMA node. `None` on UMA systems (no pool, no pinning).
pub pools: Vec<Option<Arc<rayon::ThreadPool>>>,
/// CPU indices for each NUMA node, in node order.
pub cpus_per_node: Vec<Vec<usize>>,
}
impl NumaSetup {
/// Maximum worker slots per node (one per physical core in the node).
pub fn workers_per_node(&self) -> usize {
self.cpus_per_node
.first()
.map(|c| c.len().max(1))
.unwrap_or(1)
}
}
/// Detect NUMA topology and build per-node Rayon pools.
/// Always succeeds: falls back to a single synthetic UMA node on failure.
#[cfg(feature = "numa")]
pub fn build() -> NumaSetup {
if let Ok(topology) = Topology::new() {
let nodes: Vec<Vec<usize>> = topology
.objects_with_type(ObjectType::NUMANode)
.filter_map(|obj| obj.cpuset())
.map(|cpuset| {
cpuset
.iter_set()
.map(|idx| usize::from(idx))
.collect::<Vec<_>>()
})
.filter(|v| !v.is_empty())
.collect();
if nodes.len() > 1 {
if let Some(pools) = nodes
.iter()
.map(|cpus| build_pool(cpus).map(|p| Some(Arc::new(p))))
.collect::<Option<Vec<_>>>()
{
debug!(
"NUMA topology: {} node(s), {} core(s)/node",
nodes.len(),
nodes.first().map_or(0, |v| v.len()),
);
return NumaSetup {
pools,
cpus_per_node: nodes,
};
}
}
}
// UMA fallback: single synthetic node, all cores, no pool, no pinning.
let n_cores = obisys::effective_parallelism();
debug!("UMA: single synthetic node, {} core(s)", n_cores);
NumaSetup {
pools: vec![None],
cpus_per_node: vec![(0..n_cores).collect()],
}
}
#[cfg(not(feature = "numa"))]
pub fn build() -> NumaSetup {
let n_cores = obisys::effective_parallelism();
debug!("UMA: single synthetic node, {} core(s)", n_cores);
NumaSetup {
pools: vec![None],
cpus_per_node: vec![(0..n_cores).collect()],
}
}
/// Bind the calling thread to `cpu_indices` using hwloc.
/// Silently returns on any error so the thread still runs, just unbound.
#[cfg(feature = "numa")]
pub fn pin_current_thread(cpu_indices: &[usize]) {
let Ok(topology) = Topology::new() else {
return;
};
let mut cpuset = CpuSet::new();
for &idx in cpu_indices {
cpuset.set(idx);
}
let _ = topology.bind_cpu(&cpuset, CpuBindingFlags::THREAD);
}
#[cfg(not(feature = "numa"))]
pub fn pin_current_thread(_cpu_indices: &[usize]) {}
// ── Internal helpers ──────────────────────────────────────────────────────────
#[cfg(feature = "numa")]
fn build_pool(cpus: &[usize]) -> Option<rayon::ThreadPool> {
let cpus = cpus.to_vec();
rayon::ThreadPoolBuilder::new()
.num_threads(cpus.len())
.spawn_handler(move |thread| {
let cpus = cpus.clone();
std::thread::Builder::new().spawn(move || {
pin_current_thread(&cpus);
thread.run();
})?;
Ok(())
})
.build()
.ok()
}
src/obikindex/src/index/predicate.rs
-254
View File
@@ -1,254 +0,0 @@
use std::collections::HashMap;
use crate::index::GroupQuorumFilter;
use obitaxonomy::{TaxPath, TaxPattern};
use crate::index::meta::{GenomeInfo, IndexMeta};
// ── Operator ──────────────────────────────────────────────────────────────────
enum PredOp { Wildcard, Eq, Ne, Matches, NotMatches }
// ── MetaPred ──────────────────────────────────────────────────────────────────
/// A single predicate on genome metadata: `key OP val1|val2|…`
///
/// Operators: `=` (exact), `!=` (not equal), `~` (path ancestor), `!~` (not ancestor).
/// Multiple values separated by `|` are OR'd.
pub struct MetaPred {
key: String,
op: PredOp,
values: Vec<String>,
}
impl MetaPred {
/// Parse a predicate string of the form `key=v1|v2`, `key!=v`, `key~path`, `key!~path`.
/// The special values `*` and `all` (case-insensitive) match every genome.
pub fn parse(s: &str) -> Result<Self, String> {
let t = s.trim();
if t == "*" || t.eq_ignore_ascii_case("all") {
return Ok(Self { key: String::new(), op: PredOp::Wildcard, values: vec![] });
}
let (op, key, rhs) =
if let Some(pos) = s.find("!=") {
(PredOp::Ne, &s[..pos], &s[pos+2..])
} else if let Some(pos) = s.find("!~") {
(PredOp::NotMatches, &s[..pos], &s[pos+2..])
} else if let Some(pos) = s.find('=') {
(PredOp::Eq, &s[..pos], &s[pos+1..])
} else if let Some(pos) = s.find('~') {
(PredOp::Matches, &s[..pos], &s[pos+1..])
} else {
return Err(format!("no operator found in predicate: {s}"));
};
let key = key.trim().to_string();
if key.is_empty() { return Err(format!("empty key in predicate: {s}")); }
let values: Vec<String> = rhs.split('|').map(|v| v.trim().to_string()).collect();
if values.iter().any(|v| v.is_empty()) {
return Err(format!("empty value in predicate: {s}"));
}
Ok(Self { key, op, values })
}
/// Evaluate against one genome's metadata.
/// Returns `None` when the key is absent (NA propagation).
pub(crate) fn eval(&self, meta: &HashMap<String, String>) -> Option<bool> {
if matches!(self.op, PredOp::Wildcard) { return Some(true); }
let value = meta.get(&self.key)?;
Some(match self.op {
PredOp::Wildcard => unreachable!(),
PredOp::Eq => self.values.iter().any(|v| v == value),
PredOp::Ne => self.values.iter().all(|v| v != value),
PredOp::Matches => self.values.iter().any(|v| path_matches(value, v)),
PredOp::NotMatches => self.values.iter().all(|v| !path_matches(value, v)),
})
}
}
impl GenomeInfo {
/// Evaluate a single metadata predicate against this genome.
/// Returns `None` when the predicate's key is absent (NA propagation).
pub fn matches(&self, pred: &MetaPred) -> Option<bool> {
pred.eval(&self.meta)
}
}
// ── Path matching ─────────────────────────────────────────────────────────────
/// True if the stored taxonomy `value` matches `pattern`.
///
/// `value` must be a valid `TaxPath` (starts with `taxonomy:/`).
/// `pattern` is a `TaxPattern` query (see `obitaxonomy::TaxPattern` for syntax).
/// Returns `false` if either fails to parse.
fn path_matches(value: &str, pattern: &str) -> bool {
let Ok(path) = TaxPath::parse(value) else { return false };
let Ok(pat) = TaxPattern::parse(pattern) else { return false };
pat.matches(&path)
}
// ── Three-value group evaluation ──────────────────────────────────────────────
/// AND of all predicates (ingroup semantics).
/// Short-circuits on `Some(false)`; propagates `None` if no predicate returns `false`.
fn eval_and(preds: &[MetaPred], meta: &HashMap<String, String>) -> Option<bool> {
let mut has_na = false;
for pred in preds {
match pred.eval(meta) {
Some(false) => return Some(false),
Some(true) => {}
None => has_na = true,
}
}
if has_na { None } else { Some(true) }
}
/// OR of all predicates (outgroup semantics).
/// Short-circuits on `Some(true)`; propagates `None` if no predicate returns `true`.
fn eval_or(preds: &[MetaPred], meta: &HashMap<String, String>) -> Option<bool> {
let mut has_na = false;
for pred in preds {
match pred.eval(meta) {
Some(true) => return Some(true),
Some(false) => {}
None => has_na = true,
}
}
if has_na { None } else { Some(false) }
}
// ── Genome classification ─────────────────────────────────────────────────────
enum Membership { Ingroup, Outgroup, Uncategorized }
fn classify(
genomes: &[GenomeInfo],
ingroup: &[MetaPred],
outgroup: &[MetaPred],
) -> Vec<Membership> {
genomes.iter().map(|g| {
let in_r = if ingroup.is_empty() { None } else { eval_and(ingroup, &g.meta) };
let out_r = if outgroup.is_empty() { None } else { eval_or(outgroup, &g.meta) };
// Ingroup wins over outgroup.
if in_r == Some(true) { return Membership::Ingroup; }
if out_r == Some(true) { return Membership::Outgroup; }
Membership::Uncategorized
}).collect()
}
// ── Group quorum filter construction ──────────────────────────────────────────
pub struct GroupFilterParams {
pub threshold: u32,
pub min_count: Option<isize>,
pub max_count: Option<isize>,
pub min_frac: Option<f64>,
pub max_frac: Option<f64>,
pub min_outgroup_count: Option<isize>,
pub max_outgroup_count: Option<isize>,
pub min_outgroup_frac: Option<f64>,
pub max_outgroup_frac: Option<f64>,
}
impl IndexMeta {
/// Returns indices of genomes matching `pred_str` (single predicate).
pub fn matching_genome_indices(&self, pred_str: &str) -> Result<Vec<usize>, String> {
let pred = MetaPred::parse(pred_str)?;
let genomes = self.genomes().map_err(|e| e.to_string())?;
Ok(genomes.iter().enumerate()
.filter_map(|(i, g)| {
if g.matches(&pred) == Some(true) { Some(i) } else { std::option::Option::None }
})
.collect())
}
/// Build a `GroupQuorumFilter` from parsed predicates, evaluated against `self.genomes`.
///
/// - No groups defined: `ingroup_idx` = all genomes (implicit ingroup).
/// - `ingroup` predicates only: outgroup indices are empty.
/// - `outgroup` predicates only: ingroup indices are empty.
/// - Both defined: ingroup wins on overlap; uncategorized genomes are ignored.
pub fn build_group_filter(
&self,
ingroup_preds: &[MetaPred],
outgroup_preds: &[MetaPred],
p: GroupFilterParams,
) -> Result<GroupQuorumFilter, String> {
let genomes = self.genomes().map_err(|e| e.to_string())?;
let (ingroup_idx, outgroup_idx) = if ingroup_preds.is_empty() && outgroup_preds.is_empty() {
((0..genomes.len()).collect(), vec![])
} else {
let members = classify(&genomes, ingroup_preds, outgroup_preds);
let in_idx: Vec<usize> = members.iter().enumerate()
.filter(|(_, m)| matches!(m, Membership::Ingroup))
.map(|(i, _)| i).collect();
let out_idx: Vec<usize> = members.iter().enumerate()
.filter(|(_, m)| matches!(m, Membership::Outgroup))
.map(|(i, _)| i).collect();
(in_idx, out_idx)
};
let in_size = ingroup_idx.len();
let out_size = outgroup_idx.len();
let ingroup_quorum_explicit = p.min_count.is_some() || p.max_count.is_some()
|| p.min_frac.is_some() || p.max_frac.is_some();
let outgroup_quorum_explicit = p.min_outgroup_count.is_some() || p.max_outgroup_count.is_some()
|| p.min_outgroup_frac.is_some() || p.max_outgroup_frac.is_some();
let default_min_frac = if !ingroup_preds.is_empty() && !ingroup_quorum_explicit { 1.0 } else { 0.0 };
let default_max_outgroup_count = if !outgroup_preds.is_empty() && !outgroup_quorum_explicit { 0 } else { out_size };
// Resolve a signed count: negative means an offset from the group size
// (e.g. -1 = all but one), floored at 1 so the negative form always keeps
// constraining the group — even a singleton group, where n-1 would be 0
// and would otherwise drop the constraint entirely.
let resolve = |v: isize, size: usize| -> usize {
if v < 0 { (size as isize + v).max(1) as usize } else { v as usize }
};
let min_count = p.min_count.map(|v| resolve(v, in_size)).unwrap_or(0);
let max_count = p.max_count.map(|v| resolve(v, in_size)).unwrap_or(in_size);
let min_frac = p.min_frac.unwrap_or(default_min_frac);
let max_frac = p.max_frac.unwrap_or(1.0);
let min_outgroup_count = p.min_outgroup_count.map(|v| resolve(v, out_size)).unwrap_or(0);
let max_outgroup_count = p.max_outgroup_count.map(|v| resolve(v, out_size)).unwrap_or(default_max_outgroup_count);
let min_outgroup_frac = p.min_outgroup_frac.unwrap_or(0.0);
let max_outgroup_frac = p.max_outgroup_frac.unwrap_or(1.0);
for (v, lo, hi) in [
("--min-frac/--max-frac", min_frac, max_frac),
("--min-outgroup-frac/--max-outgroup-frac", min_outgroup_frac, max_outgroup_frac),
] {
if !(0.0..=1.0).contains(&lo) || !(0.0..=1.0).contains(&hi) {
return Err(format!("{v}: fraction values must be in [0.0, 1.0]"));
}
if lo > hi {
return Err(format!("{v}: min ({lo}) is greater than max ({hi})"));
}
}
if min_count > max_count {
return Err(format!("--min-count/--max-count: min ({min_count}) is greater than max ({max_count})"));
}
if min_outgroup_count > max_outgroup_count {
return Err(format!("--min-outgroup-count/--max-outgroup-count: min ({min_outgroup_count}) is greater than max ({max_outgroup_count})"));
}
Ok(GroupQuorumFilter {
ingroup_idx,
outgroup_idx,
threshold: p.threshold,
min_count,
max_count,
min_frac,
max_frac,
min_outgroup_count,
max_outgroup_count,
min_outgroup_frac,
max_outgroup_frac,
})
}
}
src/obikindex/src/index/progress.rs
-21
View File
@@ -1,21 +0,0 @@
use std::sync::OnceLock;
use indicatif::MultiProgress;
static MULTI: OnceLock<MultiProgress> = OnceLock::new();
/// Initialise the shared progress display. Call once from the binary before
/// any index operation. Subsequent calls are silently ignored.
pub fn init(multi: MultiProgress) {
let _ = MULTI.set(multi);
}
/// Return the shared `MultiProgress`, creating a plain default one if the
/// binary never called [`init`].
pub fn get() -> &'static MultiProgress {
MULTI.get_or_init(MultiProgress::new)
}
pub(crate) fn multi() -> &'static MultiProgress {
get()
}
src/obikindex/src/index/query_layer.rs
-240
View File
@@ -1,240 +0,0 @@
use std::collections::HashMap;
use std::path::Path;
use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix};
use obikseq::CanonicalKmer;
use crate::layer::{MphfLayer, OLMError};
use obiskio::{SKError, SKResult};
use crate::index::kmer_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) -> SKResult<Self> {
let mphf = MphfLayer::open(layer_dir).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))
.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;
src/obikindex/src/index/rebuild.rs
-77
View File
@@ -1,77 +0,0 @@
use std::path::Path;
use crate::index::builder::IndexBuilder;
use crate::index::{KmerFilter, MergeMode};
use obisys::{Reporter, Stage, progress_bar};
use tracing::info;
use crate::index::error::{OKIError, OKIResult};
use crate::index::kmer_index::KmerIndex;
use crate::index::state::IndexState;
impl KmerIndex {
/// Rebuild `src` into a new compact single-layer index at `output`.
///
/// Only k-mers whose per-genome row passes every filter in `filters` are
/// written. If `filters` is empty every k-mer is kept (pure compaction).
///
/// `mode` controls whether the output stores counts or presence/absence.
/// A count source may be rebuilt in presence mode; a presence source
/// cannot be rebuilt in count mode.
pub fn rebuild<P: AsRef<Path>>(
output: P,
src: &KmerIndex,
filters: &[Box<dyn KmerFilter>],
mode: MergeMode,
force: bool,
rep: &mut Reporter,
) -> OKIResult<Self> {
let output = output.as_ref();
if src.state()? != IndexState::Indexed {
return Err(OKIError::NotIndexed(src.root_path.clone()));
}
if mode == MergeMode::Count && !src.meta.config.with_counts {
return Err(OKIError::InvalidInput(
"cannot rebuild in count mode from a presence-only source index".into(),
));
}
KmerIndex::clear_output_for_create(output, force)?;
// ── Create output directory + metadata ────────────────────────────────
let mut config = src.meta.config.clone();
config.with_counts = mode == MergeMode::Count;
let genomes = src.genomes()?;
let n_genomes = genomes.len();
let n_partitions = src.n_partitions();
let block_bits = config.block_bits;
// ── Create an empty destination KmerPartition ─────────────────────────
let dst_partition = KmerIndex::create_skeleton(output, config, genomes)?;
info!(
"rebuild: {} partition(s), {} genome(s), mode={:?}",
n_partitions, n_genomes, mode,
);
let t = Stage::start("rebuild");
let pb = progress_bar("rebuild", n_partitions as u64, "partitions");
let order: Vec<usize> = (0..n_partitions).collect();
let runner = crate::index::numa::PartitionRunner::new();
runner.run(
&order,
|i| dst_partition.rebuild_partition(src, i, filters, mode, n_genomes, block_bits),
|_, _, _| { pb.inc(1); },
).map_err(OKIError::Partition)?;
pb.finish_and_clear();
rep.push(t.stop());
KmerIndex::finalize_indexed(output, rep)
}
}
src/obikindex/src/index/rebuild_layer.rs
-270
View File
@@ -1,270 +0,0 @@
use std::path::Path;
use obicompactvec::{
FilterMask, PersistentBitMatrixBuilder, PersistentBitVecBuilder,
PersistentCompactIntMatrixBuilder, PersistentCompactIntVecBuilder, eval_filter_mask,
};
use obidebruinj::GraphDeBruijn;
use obikseq::CanonicalKmer;
use crate::layer::meta::PartitionMeta;
use crate::layer::{IndexMode, MphfLayer, layer_dir};
use obiskio::{SKError, SKResult, UnitigFileReader};
use crate::index::common::{load_meta, olm_to_sk};
use crate::index::filter::KmerFilter;
use crate::index::graph_pipeline::materialize_layer;
use crate::index::merge_layer::{MergeMode, SrcLayerData};
use crate::index::kmer_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)
.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(())
}
}
src/obikindex/src/index/reindex.rs
-130
View File
@@ -1,130 +0,0 @@
use crate::index::builder::IndexBuilder;
use crate::index::meta::IndexMeta;
use crate::layer::{IndexMode, TypedLayer};
use obisys::{Reporter, Stage, progress_bar};
use std::fs;
use std::path::Path;
use std::sync::Arc;
use tracing::info;
use crate::index::error::{OKIError, OKIResult};
use crate::index::kmer_index::KmerIndex;
use crate::index::state::IndexState;
const EVIDENCE_FILE: &str = "evidence.bin";
const FINGERPRINT_FILE: &str = "fingerprint.bin";
const UNITIG_IDX_FILE: &str = "unitigs.bin.idx";
fn olm_to_oki(e: crate::layer::OLMError) -> OKIError {
OKIError::InvalidInput(e.to_string())
}
impl KmerIndex {
/// Convert every layer's evidence bundle to `target` in-place.
///
/// - `Exact` → builds `evidence.bin` + `unitigs.bin.idx`, removes `fingerprint.bin`
/// - `Approx` → builds `fingerprint.bin`, removes `evidence.bin` + `unitigs.bin.idx`
///
/// The MPHF (`mphf.bin`) and unitigs (`unitigs.bin`) are never touched.
/// `index.meta` is updated with the new evidence kind on success.
pub fn reindex(
&mut self,
target: IndexMode,
block_bits: u8,
rep: &mut Reporter,
) -> OKIResult<()> {
if self.state()? != IndexState::Indexed {
return Err(OKIError::NotIndexed(self.root_path.clone()));
}
let n = self.n_partitions();
info!(
"reindex {} partition(s): {:?} → {:?}",
n, self.meta.config.evidence, target,
);
let t = Stage::start("reindex");
let pb = progress_bar("reindex", n as u64, "partitions");
let order: Vec<usize> = (0..n).collect();
let runner = crate::index::numa::PartitionRunner::new();
runner.run(
&order,
|i| {
reindex_partition(self, i, &target, block_bits)
.map_err(|e| OKIError::InvalidInput(format!("partition {i}: {e}")))
},
|_, _, _| {
pb.inc(1);
},
)?;
pb.finish_and_clear();
let mut config = self.meta.config.clone();
config.evidence = target;
if matches!(config.evidence, IndexMode::Exact) {
config.block_bits = block_bits;
}
let genomes = self.meta.genomes().map_err(OKIError::Io)?;
self.meta.rewrite_config(config, genomes).map_err(OKIError::Io)?;
self.meta = Arc::new(IndexMeta::open(self).map_err(OKIError::Io)?);
rep.push(t.stop());
Ok(())
}
}
/// Process all layers of one partition's index directory.
fn reindex_partition(
index: &KmerIndex,
i: usize,
target: &IndexMode,
block_bits: u8,
) -> OKIResult<()> {
if !index.index_dir(i).exists() {
return Ok(());
}
let n_layers = index
.n_layers(i)
.map_err(|e| OKIError::InvalidInput(e.to_string()))?;
for layer_idx in 0..n_layers {
reindex_layer(&index.layer_dir(i, layer_idx), target, block_bits)?;
}
Ok(())
}
fn reindex_layer(layer_dir: &Path, target: &IndexMode, block_bits: u8) -> OKIResult<()> {
match target {
IndexMode::Exact => {
TypedLayer::<()>::build_exact_evidence(layer_dir, block_bits).map_err(olm_to_oki)?;
}
IndexMode::Approx { b, z } | IndexMode::Hybrid { b, z } => {
TypedLayer::<()>::build_approx_evidence(layer_dir, *b, *z).map_err(olm_to_oki)?;
}
}
remove_stale_evidence(layer_dir, target)
}
fn remove_stale_evidence(layer_dir: &Path, target: &IndexMode) -> OKIResult<()> {
match target {
IndexMode::Exact => {
remove_if_exists(&layer_dir.join(FINGERPRINT_FILE));
}
IndexMode::Approx { .. } => {
remove_if_exists(&layer_dir.join(EVIDENCE_FILE));
remove_if_exists(&layer_dir.join(UNITIG_IDX_FILE));
}
IndexMode::Hybrid { .. } => {
// both bundles kept — nothing to remove
}
}
Ok(())
}
fn remove_if_exists(path: &Path) {
if let Err(e) = fs::remove_file(path) {
if e.kind() != std::io::ErrorKind::NotFound {
eprintln!("warning: could not remove {}: {e}", path.display());
}
}
}
src/obikindex/src/index/select.rs
-154
View File
@@ -1,154 +0,0 @@
use std::path::Path;
use std::sync::Arc;
use crate::index::builder::IndexBuilder;
use crate::index::OutputCol;
use obisys::{Reporter, Stage, progress_bar};
use tracing::info;
use crate::index::error::{OKIError, OKIResult};
use crate::index::kmer_index::KmerIndex;
use crate::index::meta::{GenomeInfo, IndexMeta};
use crate::index::state::IndexState;
impl KmerIndex {
/// Create a new index at `output` by projecting/aggregating the genome columns
/// of `src` according to `specs`.
///
/// `output_presence` — if true, output uses bit matrices (0/1), regardless of
/// whether the source stores counts. The caller is responsible for ensuring all
/// specs use logical operators when `output_presence=true` on a count source.
pub fn select<P: AsRef<Path>>(
output: P,
src: &KmerIndex,
specs: &[OutputCol],
threshold: u32,
output_presence: bool,
force: bool,
rep: &mut Reporter,
) -> OKIResult<Self> {
let output = output.as_ref();
if src.state()? != IndexState::Indexed {
return Err(OKIError::NotIndexed(src.root_path.clone()));
}
KmerIndex::clear_output_for_create(output, force)?;
let mut config = src.meta.config.clone();
config.with_counts = !output_presence;
let genomes: Vec<GenomeInfo> = specs
.iter()
.map(|s| GenomeInfo::new(s.label.clone()))
.collect();
let n_src_genomes = src.meta.genomes().map_err(OKIError::Io)?.len();
let n_partitions = src.n_partitions();
let dst_partition = KmerIndex::create_skeleton(output, config, genomes)?;
info!(
"select: {} partition(s), {} source genome(s) → {} output column(s)",
n_partitions,
n_src_genomes,
specs.len(),
);
let t = Stage::start("select");
let pb = progress_bar("select", n_partitions as u64, "partitions");
let order: Vec<usize> = (0..n_partitions).collect();
let runner = crate::index::numa::PartitionRunner::new();
runner
.run(
&order,
|i| {
dst_partition.select_partition(
src,
i,
specs,
n_src_genomes,
threshold,
output_presence,
false,
)
},
|_, _, _| {
pb.inc(1);
},
)
.map_err(OKIError::Partition)?;
pb.finish_and_clear();
rep.push(t.stop());
KmerIndex::finalize_indexed(output, rep)
}
/// Rewrite the genome columns of this index in-place according to `specs`.
///
/// The MPHF and unitig files are unchanged; only data matrices are rewritten.
pub fn select_in_place(
&mut self,
specs: &[OutputCol],
threshold: u32,
output_presence: bool,
rep: &mut Reporter,
) -> OKIResult<()> {
if self.state()? != IndexState::Indexed {
return Err(OKIError::NotIndexed(self.root_path.clone()));
}
let n_src_genomes = self.meta.genomes().map_err(OKIError::Io)?.len();
let n_partitions = self.n_partitions();
info!(
"select (in-place): {} partition(s), {} source genome(s) → {} output column(s)",
n_partitions,
n_src_genomes,
specs.len(),
);
let t = Stage::start("select");
let pb = progress_bar("select", n_partitions as u64, "partitions");
let order: Vec<usize> = (0..n_partitions).collect();
let runner = crate::index::numa::PartitionRunner::new();
runner
.run(
&order,
|i| {
self.select_partition(
self,
i,
specs,
n_src_genomes,
threshold,
output_presence,
true,
)
},
|_, _, _| {
pb.inc(1);
},
)
.map_err(OKIError::Partition)?;
pb.finish_and_clear();
rep.push(t.stop());
let mut config = self.meta.config.clone();
config.with_counts = !output_presence;
let genomes: Vec<GenomeInfo> = specs
.iter()
.map(|s| GenomeInfo::new(s.label.clone()))
.collect();
self.meta.rewrite_config(config, genomes).map_err(OKIError::Io)?;
self.meta = Arc::new(IndexMeta::open(self).map_err(OKIError::Io)?);
let t_pack = Stage::start("pack");
self.pack_matrices(false)?;
rep.push(t_pack.stop());
Ok(())
}
}
src/obikindex/src/index/select_layer.rs
-309
View File
@@ -1,309 +0,0 @@
use std::fs;
use std::io;
use std::path::{Path, PathBuf};
use obicompactvec::{
ColGroup, MatrixGroupOps, PersistentBitMatrix, PersistentBitMatrixBuilder,
PersistentCompactIntMatrix, PersistentCompactIntMatrixBuilder,
};
use crate::layer::OLMError;
use obiskio::{SKError, SKResult};
use crate::index::kmer_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(())
}
}
src/obikindex/src/index/stats.rs
-187
View File
@@ -1,187 +0,0 @@
use std::fs;
use std::path::Path;
use obicompactvec::{LayerMeta, PersistentBitMatrix, PersistentCompactIntMatrix};
use obicompactvec::traits::ColumnWeights;
use rayon::prelude::*;
use crate::index::error::{OKIError, OKIResult};
use crate::index::kmer_index::KmerIndex;
/// Bits per kmer broken down by index component.
pub struct IndexBitsPerKmer {
/// Total distinct k-mers across all partitions and layers.
pub n_kmers: usize,
/// Number of genomes in the index.
pub n_genomes: usize,
/// Bits used by the minimal perfect hash function (`mphf.bin`).
pub mphf: f64,
/// Bits used by the evidence files (`evidence.bin`, `unitigs.bin*`,
/// `fingerprint.bin`).
pub evidence: f64,
/// Bits used by the count/presence matrices (`counts/` and `presence/`),
/// normalised by k-mers only.
pub matrix: f64,
/// `matrix` divided by the number of genomes — intrinsic encoding
/// efficiency, independent of index size.
pub matrix_per_genome: f64,
/// Sum of mphf + evidence + matrix.
pub total: f64,
}
// ── File-size helpers ─────────────────────────────────────────────────────────
fn file_bytes(path: &Path) -> u64 {
fs::metadata(path).map(|m| m.len()).unwrap_or(0)
}
fn dir_bytes(dir: &Path) -> u64 {
if !dir.exists() { return 0; }
fs::read_dir(dir)
.map(|entries| {
entries
.filter_map(|e| e.ok())
.filter_map(|e| e.metadata().ok())
.filter(|m| m.is_file())
.map(|m| m.len())
.sum()
})
.unwrap_or(0)
}
// ── Per-layer accounting ──────────────────────────────────────────────────────
struct LayerBytes {
n_kmers: usize,
mphf: u64,
evidence: u64,
matrix: u64,
}
fn layer_bytes(layer_dir: &Path) -> LayerBytes {
let n_kmers = LayerMeta::load(layer_dir).map(|m| m.n).unwrap_or(0);
let mphf = file_bytes(&layer_dir.join("mphf.bin"));
let evidence = file_bytes(&layer_dir.join("unitigs.bin"))
+ file_bytes(&layer_dir.join("unitigs.bin.idx"))
+ file_bytes(&layer_dir.join("evidence.bin"))
+ file_bytes(&layer_dir.join("fingerprint.bin"));
let matrix = dir_bytes(&layer_dir.join("counts"))
+ dir_bytes(&layer_dir.join("presence"));
LayerBytes { n_kmers, mphf, evidence, matrix }
}
// ── KmerIndex::bits_per_kmer ──────────────────────────────────────────────────
impl KmerIndex {
/// Compute bits-per-kmer statistics for the built index.
///
/// File sizes are read directly from disk; kmer counts come from
/// `layer_meta.json` (no need to scan the MPHF or unitig files).
/// Computation is parallelised across partitions.
pub fn bits_per_kmer(&self) -> OKIResult<IndexBitsPerKmer> {
let n = self.n_partitions();
let n_genomes = self.meta().genomes().map_err(OKIError::Io)?.len().max(1);
let (n_kmers, mphf_b, evidence_b, matrix_b) = (0..n)
.into_par_iter()
.map(|i| {
let index_dir = self.index_dir(i);
if !index_dir.exists() { return (0usize, 0u64, 0u64, 0u64); }
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.layer_dir(i, l));
(acc.0 + lb.n_kmers, acc.1 + lb.mphf, acc.2 + lb.evidence, acc.3 + lb.matrix)
})
})
.reduce(|| (0, 0, 0, 0), |a, b| (a.0 + b.0, a.1 + b.1, a.2 + b.2, a.3 + b.3));
if n_kmers == 0 {
return Ok(IndexBitsPerKmer {
n_kmers: 0, n_genomes,
mphf: 0.0, evidence: 0.0,
matrix: 0.0, matrix_per_genome: 0.0, total: 0.0,
});
}
let bpk = |bytes: u64| bytes as f64 * 8.0 / n_kmers as f64;
let matrix = bpk(matrix_b);
Ok(IndexBitsPerKmer {
n_kmers,
n_genomes,
mphf: bpk(mphf_b),
evidence: bpk(evidence_b),
matrix,
matrix_per_genome: matrix / n_genomes as f64,
total: bpk(mphf_b + evidence_b + matrix_b),
})
}
/// Return `(total_distinct_kmers, per_genome_kmer_counts)`.
///
/// For each genome, the count is the number of distinct k-mers for which
/// that genome has a non-zero value (presence = 1, count > 0).
/// Partitions are scanned in parallel; results are summed across partitions.
pub fn genome_kmer_counts(&self) -> OKIResult<(usize, Vec<u64>)> {
let n = self.n_partitions();
let n_genomes = self.meta.genomes().map_err(OKIError::Io)?.len();
let partials: Vec<(usize, Vec<u64>)> = (0..n)
.into_par_iter()
.map(|i| {
let mut counts = vec![0u64; n_genomes];
let mut n_kmers = 0usize;
let index_dir = self.index_dir(i);
if !index_dir.exists() { return (0, counts); }
let n_layers = self.n_layers(i).unwrap_or(0);
for l in 0..n_layers {
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);
let mat: Box<dyn ColumnWeights> =
if this_layer_dir.join("counts").exists()
&& !this_layer_dir.join("presence").exists()
{
match crate::layer::open_data::<PersistentCompactIntMatrix>(&index_dir, l) {
Ok(m) => Box::new(m),
Err(_) => continue,
}
} else {
match crate::layer::open_data::<PersistentBitMatrix>(&index_dir, l) {
Ok(m) => Box::new(m),
Err(_) => continue,
}
};
let col_counts = mat.partial_kmer_counts();
for (c, &v) in col_counts.iter().enumerate() {
if c < n_genomes { counts[c] += v; }
}
}
(n_kmers, counts)
})
.collect();
let total_kmers: usize = partials.iter().map(|(n, _)| n).sum();
let mut total_counts = vec![0u64; n_genomes];
for (_, counts) in partials {
for (i, v) in counts.into_iter().enumerate() {
total_counts[i] += v;
}
}
Ok((total_kmers, total_counts))
}
}
src/obikindex/src/index/tests/query_layer.rs
-94
View File
@@ -1,94 +0,0 @@
use super::*;
use crate::index::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: crate::layer::IndexMode::Exact,
block_bits: 0,
};
let index = KmerIndex::create(tmp.path().join("idx"), config, None).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: crate::layer::IndexMode::Exact,
block_bits: 0,
};
let index = KmerIndex::create(tmp.path().join("idx"), config, None).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());
}
src/obikindex/src/layer/content_layer.rs
+108 -61
View File
@@ -24,77 +24,103 @@
use std::path::{Path, PathBuf};
use crate::layer::utils::LAYERNAME_SUFFIX;
use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix};
use obikseq::CanonicalKmer;
use crate::layer::error::OLMResult;
use crate::layer::meta::IndexMode;
use crate::index::error::OKIResult;
use crate::layer::mphf_layer::{
EVIDENCE_FILE, EvidenceKind, FINGERPRINT_FILE, MPHF_FILE, UNITIGS_FILE,
};
use crate::layer::typed_layer::{COUNTS_DIR, LayerContent, PRESENCE_DIR, TypedLayer};
use crate::partition::KmerPartition;
/// Superkmer-file extension used by every pre-layer-construction artifact
/// below (`raw`/`dereplicated`) — an implementation detail of the SK file
/// format, never meant to leak as a literal string past this module.
const SK_EXT: &str = "skmer.zst";
/// One layer, at any point in its life — see the module docs. Only
/// [`Empty`](Layer::Empty) and the two ready-to-read states
/// (`Count`/`Presence`) exist so far; states in between (unitigs written,
/// MPHF built, evidence built, matrix not yet built) are not represented
/// yet.
///
/// Every variant carries `dir` + `id` — a layer's identity never depends on
/// its state. `id` is its own layer number within the partition (not yet
/// consumed by anything — `KmerPartition` doesn't hand it in yet, that
/// lands with `KmerLayer::at`). `dir`/`id` live here, on `KmerLayer` itself,
/// not inside `TypedLayer`/`MphfLayer`: those only ever need `dir`
/// transiently, while opening — `KmerLayer` is what has to answer `.dir()`
/// after the fact, so it's the one that keeps it.
pub enum KmerLayer {
/// Just a directory — nothing built yet. Every method below other than
/// the path accessors panics on this variant: calling them means the
/// caller assumed a layer was ready when it wasn't, an implementation
/// error to surface loudly, not paper over with a default value.
Empty {
Empty { dir: PathBuf, id: usize },
Count {
dir: PathBuf,
id: usize,
layer: TypedLayer<PersistentCompactIntMatrix>,
},
Presence {
dir: PathBuf,
id: usize,
layer: TypedLayer<PersistentBitMatrix>,
},
Count(TypedLayer<PersistentCompactIntMatrix>),
Presence(TypedLayer<PersistentBitMatrix>),
}
impl KmerLayer {
pub fn new(partition: &KmerPartition, i: usize) -> Self {
let dir = partition.dir().join(format!("{LAYERNAME_SUFFIX}{i}"));
KmerLayer::Empty { dir, id: i }
}
/// An empty shell at `dir` — creates the directory (if it doesn't
/// already exist) but nothing inside it. The starting state for
/// building a new layer.
pub fn create(dir: &Path) -> std::io::Result<Self> {
std::fs::create_dir_all(dir)?;
Ok(KmerLayer::Empty {
dir: dir.to_owned(),
})
pub fn create(self) -> std::io::Result<Self> {
std::fs::create_dir_all(self.dir())?;
Ok(self)
}
/// Open one layer, auto-detecting count vs. presence from what is
/// actually on disk (`counts/` present and wanted, else presence) — the
/// single source of truth every caller that opens a layer's own matrix
/// should share, so two callers can never disagree about which format a
/// given layer is. Dense vs. sparse presence storage is
/// `PersistentBitMatrix::open`'s own concern (it's a 4-way
/// `Columnar`/`Packed`/`Sparse`/`Implicit` enum internally), not decided
/// here.
pub fn open(layer_dir: &Path, mode: &IndexMode, with_counts: bool) -> OLMResult<Self> {
if with_counts && layer_dir.join(COUNTS_DIR).exists() {
return TypedLayer::<PersistentCompactIntMatrix>::open(layer_dir, mode)
.map(KmerLayer::Count);
/// Open this layer, auto-detecting count vs. presence from what is
/// actually on disk (`counts/` present, else presence) — the single
/// source of truth every caller that opens a layer's own matrix should
/// share, so two callers can never disagree about which format a given
/// layer is. Dense vs. sparse presence storage is `PersistentBitMatrix::
/// open`'s own concern (a 4-way `Columnar`/`Packed`/`Sparse`/`Implicit`
/// enum internally), not decided here. A no-op (returns `self`
/// unchanged) if already `Count`/`Presence`.
pub fn open(self) -> OKIResult<Self> {
let (dir, id) = match self {
KmerLayer::Empty { dir, id } => (dir, id),
already_open => return Ok(already_open),
};
if dir.join(COUNTS_DIR).exists() {
let layer = TypedLayer::<PersistentCompactIntMatrix>::open(&dir)?;
Ok(KmerLayer::Count { dir, id, layer })
} else {
let layer = TypedLayer::<PersistentBitMatrix>::open(&dir)?;
Ok(KmerLayer::Presence { dir, id, layer })
}
TypedLayer::<PersistentBitMatrix>::open(layer_dir, mode).map(KmerLayer::Presence)
}
// ── Paths — only meaningful before anything is built ───────────────
// ── Paths — valid in every state, `dir` never varies with content ────
//
// Once a layer is `Count`/`Presence`, its caller already knows the
// directory (it had to pass it to `open`) — these exist for builder
// code holding an `Empty` layer, so the `mphf.bin`/`unitigs.bin`/…
// naming stays defined once, here, rather than re-declared as
// string literals at every call site that writes into a layer
// directory (the same duplication `layer_dir`/`index_dir` fixed one
// level up — see `DevDocMD/implementation/partition_layer_cache.md`).
// The `mphf.bin`/`unitigs.bin`/… naming stays defined once, here,
// rather than re-declared as string literals at every call site that
// writes into a layer directory (the same duplication `layer_dir`/
// `index_dir` fixed one level up — see `DevDocMD/implementation/
// partition_layer_cache.md`).
/// This layer's own directory. Panics on `Count`/`Presence` — by that
/// point the caller already has the directory it opened with; asking
/// again here would mean it lost track of its own state.
/// This layer's own directory.
pub fn dir(&self) -> &Path {
match self {
KmerLayer::Empty { dir } => dir,
_ => panic!("Layer::dir() only available on Empty — caller already has this path"),
KmerLayer::Empty { dir, .. } => dir,
KmerLayer::Count { dir, .. } => dir,
KmerLayer::Presence { dir, .. } => dir,
}
}
@@ -117,36 +143,57 @@ impl KmerLayer {
self.dir().join(PRESENCE_DIR)
}
/// Path of this layer's raw, not-yet-dereplicated superkmer file —
/// written by whichever algorithm routes superkmers into this layer
/// (today: `obikindexer::algorithms::partitionner::PartitionRouter`),
/// read by whichever algorithm dereplicates it (today:
/// `obikindexer::algorithms::dereplicator::Dereplicator`). Naming this
/// once here, rather than in either algorithm submodule, is what lets
/// the two agree on the filename without depending on each other
/// directly — see `DevDocMD/implementation/partition_layer_cache.md`.
pub fn raw_superkmers_path(&self) -> PathBuf {
self.dir().join(format!("raw.{SK_EXT}"))
}
/// Path of this layer's dereplicated superkmer file — written by
/// `obikindexer::algorithms::dereplicator::Dereplicator`, read by
/// whichever algorithm counts kmer abundances from it (today:
/// `obikindexer::algorithms::partitionner::PartitionRouter::count_kmer`)
/// and, later, by `obikindex::build_index_layer` to build the real layer.
pub fn dereplicated_superkmers_path(&self) -> PathBuf {
self.dir().join(format!("dereplicated.{SK_EXT}"))
}
// ── Ready-only surface ───────────────────────────────────────────────
pub fn content(&self) -> LayerContent {
match self {
KmerLayer::Count(_) => LayerContent::Count,
KmerLayer::Presence(_) => LayerContent::Presence,
KmerLayer::Count { .. } => LayerContent::Count,
KmerLayer::Presence { .. } => LayerContent::Presence,
KmerLayer::Empty { .. } => panic!("Layer::content() called on an Empty layer"),
}
}
pub fn evidence_kind(&self) -> EvidenceKind {
match self {
KmerLayer::Count(l) => l.evidence_kind(),
KmerLayer::Presence(l) => l.evidence_kind(),
KmerLayer::Count { layer, .. } => layer.evidence_kind(),
KmerLayer::Presence { layer, .. } => layer.evidence_kind(),
KmerLayer::Empty { .. } => panic!("Layer::evidence_kind() called on an Empty layer"),
}
}
pub fn n(&self) -> usize {
match self {
KmerLayer::Count(l) => l.n(),
KmerLayer::Presence(l) => l.n(),
KmerLayer::Count { layer, .. } => layer.n(),
KmerLayer::Presence { layer, .. } => layer.n(),
KmerLayer::Empty { .. } => panic!("Layer::n() called on an Empty layer"),
}
}
pub fn find_slot(&self, kmer: CanonicalKmer) -> Option<usize> {
match self {
KmerLayer::Count(l) => l.find_slot(kmer),
KmerLayer::Presence(l) => l.find_slot(kmer),
KmerLayer::Count { layer, .. } => layer.find_slot(kmer),
KmerLayer::Presence { layer, .. } => layer.find_slot(kmer),
KmerLayer::Empty { .. } => panic!("Layer::find_slot() called on an Empty layer"),
}
}
@@ -156,16 +203,16 @@ impl KmerLayer {
/// the evidence check `find_slot`/`find` would perform is redundant.
pub fn hash_batch(&self, kmers: &[CanonicalKmer]) -> Vec<usize> {
match self {
KmerLayer::Count(l) => l.hash_batch(kmers),
KmerLayer::Presence(l) => l.hash_batch(kmers),
KmerLayer::Count { layer, .. } => layer.hash_batch(kmers),
KmerLayer::Presence { layer, .. } => layer.hash_batch(kmers),
KmerLayer::Empty { .. } => panic!("Layer::hash_batch() called on an Empty layer"),
}
}
pub fn n_cols(&self) -> usize {
match self {
KmerLayer::Count(l) => l.n_cols(),
KmerLayer::Presence(l) => l.n_cols(),
KmerLayer::Count { layer, .. } => layer.n_cols(),
KmerLayer::Presence { layer, .. } => layer.n_cols(),
KmerLayer::Empty { .. } => panic!("Layer::n_cols() called on an Empty layer"),
}
}
@@ -178,10 +225,10 @@ impl KmerLayer {
/// presence (`!= 0`) in place.
pub fn fill_sub_matrix_carries(&self, slots: &[usize], out: &mut [Vec<bool>]) {
match self {
KmerLayer::Presence(l) => l.fill_sub_matrix(slots, out),
KmerLayer::Count(l) => {
KmerLayer::Presence { layer, .. } => layer.fill_sub_matrix(slots, out),
KmerLayer::Count { layer, .. } => {
let mut counts: Vec<Vec<u32>> = out.iter().map(|_| Vec::new()).collect();
l.fill_sub_matrix(slots, &mut counts);
layer.fill_sub_matrix(slots, &mut counts);
for (o, c) in out.iter_mut().zip(counts.iter()) {
o.clear();
o.extend(c.iter().map(|&v| v != 0));
@@ -196,8 +243,8 @@ impl KmerLayer {
/// Raw MPHF lookup: kmer → slot, no membership check.
pub fn hash(&self, kmer: CanonicalKmer) -> usize {
match self {
KmerLayer::Count(l) => l.hash(kmer),
KmerLayer::Presence(l) => l.hash(kmer),
KmerLayer::Count { layer, .. } => layer.hash(kmer),
KmerLayer::Presence { layer, .. } => layer.hash(kmer),
KmerLayer::Empty { .. } => panic!("Layer::hash() called on an Empty layer"),
}
}
@@ -205,8 +252,8 @@ impl KmerLayer {
/// Iterate over all canonical kmers in the layer, in deterministic order.
pub fn iter_kmers(&self) -> crate::layer::mphf_layer::KmerIter {
match self {
KmerLayer::Count(l) => l.iter_kmers(),
KmerLayer::Presence(l) => l.iter_kmers(),
KmerLayer::Count { layer, .. } => layer.iter_kmers(),
KmerLayer::Presence { layer, .. } => layer.iter_kmers(),
KmerLayer::Empty { .. } => panic!("Layer::iter_kmers() called on an Empty layer"),
}
}
@@ -215,8 +262,8 @@ impl KmerLayer {
/// sequence index in `unitigs.bin`.
pub fn enumerate_kmers(&self) -> std::iter::Enumerate<crate::layer::mphf_layer::KmerIter> {
match self {
KmerLayer::Count(l) => l.enumerate_kmers(),
KmerLayer::Presence(l) => l.enumerate_kmers(),
KmerLayer::Count { layer, .. } => layer.enumerate_kmers(),
KmerLayer::Presence { layer, .. } => layer.enumerate_kmers(),
KmerLayer::Empty { .. } => panic!("Layer::enumerate_kmers() called on an Empty layer"),
}
}
@@ -224,8 +271,8 @@ impl KmerLayer {
/// Iterate over the layer's canonical kmers in batches of `n`.
pub fn iter_kmers_batch(&self, n: usize) -> crate::layer::mphf_layer::KmerBatchIter {
match self {
KmerLayer::Count(l) => l.iter_kmers_batch(n),
KmerLayer::Presence(l) => l.iter_kmers_batch(n),
KmerLayer::Count { layer, .. } => layer.iter_kmers_batch(n),
KmerLayer::Presence { layer, .. } => layer.iter_kmers_batch(n),
KmerLayer::Empty { .. } => panic!("Layer::iter_kmers_batch() called on an Empty layer"),
}
}
@@ -237,8 +284,8 @@ impl KmerLayer {
n: usize,
) -> Box<dyn Iterator<Item = (usize, Vec<CanonicalKmer>)> + Send + 'static> {
match self {
KmerLayer::Count(l) => Box::new(l.enumerate_kmers_batch(n)),
KmerLayer::Presence(l) => Box::new(l.enumerate_kmers_batch(n)),
KmerLayer::Count { layer, .. } => Box::new(layer.enumerate_kmers_batch(n)),
KmerLayer::Presence { layer, .. } => Box::new(layer.enumerate_kmers_batch(n)),
KmerLayer::Empty { .. } => {
panic!("Layer::enumerate_kmers_batch() called on an Empty layer")
}
src/obikindex/src/layer/error.rs
-50
View File
@@ -1,50 +0,0 @@
use std::fmt;
use std::io;
#[derive(Debug)]
pub enum OLMError {
Io(io::Error),
Json(serde_json::Error),
Mphf(String),
InvalidLayer(String),
}
pub type OLMResult<T> = Result<T, OLMError>;
impl fmt::Display for OLMError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
OLMError::Io(e) => write!(f, "I/O error: {e}"),
OLMError::Json(e) => write!(f, "JSON error: {e}"),
OLMError::Mphf(s) => write!(f, "MPHF error: {s}"),
OLMError::InvalidLayer(s) => write!(f, "invalid layer: {s}"),
}
}
}
impl std::error::Error for OLMError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
OLMError::Io(e) => Some(e),
OLMError::Json(e) => Some(e),
_ => None,
}
}
}
impl From<io::Error> for OLMError {
fn from(e: io::Error) -> Self { OLMError::Io(e) }
}
impl From<serde_json::Error> for OLMError {
fn from(e: serde_json::Error) -> Self { OLMError::Json(e) }
}
impl From<obiskio::SKError> for OLMError {
fn from(e: obiskio::SKError) -> Self {
match e {
obiskio::SKError::Io(io_err) => OLMError::Io(io_err),
other => OLMError::InvalidLayer(other.to_string()),
}
}
}
src/obikindex/src/layer/evidence.rs
+10 -10
View File
@@ -6,14 +6,14 @@ use std::path::Path;
use memmap2::Mmap;
use crate::layer::error::{OLMError, OLMResult};
use crate::index::error::{OKIError, OKIResult};
pub struct Evidence {
mmap: Mmap,
}
impl Evidence {
pub fn open(path: &Path) -> OLMResult<Self> {
pub fn open(path: &Path) -> OKIResult<Self> {
let f = File::open(path)?;
let mmap = unsafe { Mmap::map(&f)? };
Ok(Self { mmap })
@@ -26,16 +26,16 @@ impl Evidence {
(raw >> 7, (raw & 0x7F) as u8)
}
#[inline]
pub fn encode(chunk_id: u32, rank: u8) -> u32 {
(chunk_id << 7) | (rank as u32 & 0x7F)
}
pub fn len(&self) -> usize {
self.mmap.len() / 4
}
}
#[inline]
pub fn encode(chunk_id: u32, rank: u8) -> u32 {
(chunk_id << 7) | (rank as u32 & 0x7F)
}
pub struct EvidenceWriter {
buf: Vec<u32>,
}
@@ -47,13 +47,13 @@ impl EvidenceWriter {
#[inline]
pub fn set(&mut self, slot: usize, chunk_id: u32, rank: u8) {
self.buf[slot] = encode(chunk_id, rank);
self.buf[slot] = Evidence::encode(chunk_id, rank);
}
pub fn write(self, path: &Path) -> OLMResult<()> {
pub fn write(self, path: &Path) -> OKIResult<()> {
let mut f = BufWriter::new(File::create(path)?);
for v in self.buf {
f.write_all(&v.to_le_bytes()).map_err(OLMError::Io)?;
f.write_all(&v.to_le_bytes()).map_err(OKIError::Io)?;
}
Ok(())
}
src/obikindex/src/layer/fingerprint.rs
+10 -10
View File
@@ -14,7 +14,7 @@ use std::path::Path;
use bitvec::prelude::*;
use memmap2::Mmap;
use crate::layer::error::{OLMError, OLMResult};
use crate::index::error::{OKIError, OKIResult};
const MAGIC: &[u8; 4] = b"FPVF";
const HEADER: usize = 16;
@@ -30,15 +30,15 @@ pub struct FingerprintVec {
}
impl FingerprintVec {
pub fn open(path: &Path) -> OLMResult<Self> {
pub fn open(path: &Path) -> OKIResult<Self> {
let f = File::open(path)?;
let mmap = unsafe { Mmap::map(&f)? };
if mmap.len() < HEADER || &mmap[..4] != MAGIC {
return Err(OLMError::InvalidLayer("bad fingerprint magic".into()));
return Err(OKIError::InvalidData { context: "layer", detail: "bad fingerprint magic".into() });
}
let b = mmap[4];
if b == 0 || b > 64 {
return Err(OLMError::InvalidLayer("invalid fingerprint width".into()));
return Err(OKIError::InvalidData { context: "layer", detail: "invalid fingerprint width".into() });
}
let n = u64::from_le_bytes(mmap[8..16].try_into().unwrap()) as usize;
let mask: u64 = if b == 64 { u64::MAX } else { (1u64 << b) - 1 };
@@ -89,12 +89,12 @@ impl FingerprintVecWriter {
bits[lo .. lo + self.b as usize].store_le(fingerprint);
}
pub fn write(self, path: &Path) -> OLMResult<()> {
let mut f = BufWriter::new(File::create(path).map_err(OLMError::Io)?);
f.write_all(MAGIC).map_err(OLMError::Io)?;
f.write_all(&[self.b, 0, 0, 0]).map_err(OLMError::Io)?;
f.write_all(&(self.n as u64).to_le_bytes()).map_err(OLMError::Io)?;
f.write_all(&self.buf).map_err(OLMError::Io)?;
pub fn write(self, path: &Path) -> OKIResult<()> {
let mut f = BufWriter::new(File::create(path).map_err(OKIError::Io)?);
f.write_all(MAGIC).map_err(OKIError::Io)?;
f.write_all(&[self.b, 0, 0, 0]).map_err(OKIError::Io)?;
f.write_all(&(self.n as u64).to_le_bytes()).map_err(OKIError::Io)?;
f.write_all(&self.buf).map_err(OKIError::Io)?;
Ok(())
}
}
src/obikindex/src/layer/layered_store.rs
-100
View File
@@ -1,100 +0,0 @@
use ndarray::{Array1, Array2};
use rayon::prelude::*;
use obicompactvec::traits::{BitPartials, ColumnWeights, CountPartials};
/// A store that aggregates a `Vec<S>` — one entry per layer (within a partition)
/// or one entry per partition.
///
/// Blanket impls of `ColumnWeights`, `CountPartials`, and `BitPartials` propagate
/// automatically: `LayeredStore<LayeredStore<S>>` implements the same traits as
/// `LayeredStore<S>`, giving the partitioned level for free.
pub struct LayeredStore<S>(pub Vec<S>);
impl<S> LayeredStore<S> {
pub fn new(layers: Vec<S>) -> Self { Self(layers) }
pub fn layers(&self) -> &[S] { &self.0 }
pub fn n_layers(&self) -> usize { self.0.len() }
pub fn is_empty(&self) -> bool { self.0.is_empty() }
}
// ── ColumnWeights ─────────────────────────────────────────────────────────────
impl<S: ColumnWeights> ColumnWeights for LayeredStore<S> {
fn col_weights(&self) -> Array1<u64> {
self.0.par_iter()
.map(|s| s.col_weights())
.reduce_with(|a, b| a + b)
.unwrap_or_else(|| Array1::zeros(0))
}
}
// ── CountPartials ─────────────────────────────────────────────────────────────
impl<S: CountPartials> CountPartials for LayeredStore<S> {
fn partial_bray(&self) -> Array2<u64> {
self.0.par_iter()
.map(|s| s.partial_bray())
.reduce_with(|a, b| a + b)
.unwrap()
}
fn partial_euclidean(&self) -> Array2<f64> {
self.0.par_iter()
.map(|s| s.partial_euclidean())
.reduce_with(|a, b| a + b)
.unwrap()
}
fn partial_threshold_jaccard(&self, threshold: u32) -> (Array2<u64>, Array2<u64>) {
self.0.par_iter()
.map(|s| s.partial_threshold_jaccard(threshold))
.reduce_with(|(ai, au), (bi, bu)| (ai + bi, au + bu))
.unwrap()
}
fn partial_relfreq_bray(&self, global: &Array1<u64>) -> Array2<f64> {
self.0.par_iter()
.map(|s| s.partial_relfreq_bray(global))
.reduce_with(|a, b| a + b)
.unwrap()
}
fn partial_relfreq_euclidean(&self, global: &Array1<u64>) -> Array2<f64> {
self.0.par_iter()
.map(|s| s.partial_relfreq_euclidean(global))
.reduce_with(|a, b| a + b)
.unwrap()
}
fn partial_hellinger(&self, global: &Array1<u64>) -> Array2<f64> {
self.0.par_iter()
.map(|s| s.partial_hellinger(global))
.reduce_with(|a, b| a + b)
.unwrap()
}
}
// ── BitPartials ───────────────────────────────────────────────────────────────
impl<S: BitPartials> BitPartials for LayeredStore<S> {
fn partial_jaccard(&self) -> (Array2<u64>, Array2<u64>) {
self.0.par_iter()
.map(|s| s.partial_jaccard())
.reduce_with(|(ai, au), (bi, bu)| (ai + bi, au + bu))
.unwrap()
}
fn partial_hamming(&self) -> Array2<u64> {
self.0.par_iter()
.map(|s| s.partial_hamming())
.reduce_with(|a, b| a + b)
.unwrap()
}
}
// ── Tests ─────────────────────────────────────────────────────────────────────
#[cfg(test)]
#[path = "tests/layered_store.rs"]
mod tests;
src/obikindex/src/layer/map.rs
-128
View File
@@ -1,128 +0,0 @@
use std::collections::HashMap;
use std::fs;
use std::path::{Path, PathBuf};
use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix, StorageKind};
use obikseq::CanonicalKmer;
use obiskio::{UnitigFileWriter, DEFAULT_BLOCK_BITS};
use crate::layer::error::{OLMError, OLMResult};
use crate::layer::typed_layer::{layer_dir, Hit, TypedLayer, LayerContent, LayerData};
use crate::layer::meta::{IndexMode, PartitionMeta};
use crate::layer::mphf_layer::EvidenceKind;
/// Layered kmer index for a single partition.
///
/// Each layer covers a disjoint kmer set. Queries probe layers in order;
/// the first match wins. Adding a dataset appends a new layer without
/// rebuilding existing ones.
pub struct LayeredMap<D: LayerData = ()> {
root: PathBuf,
meta: PartitionMeta,
layers: Vec<TypedLayer<D>>,
}
// ── Common methods ────────────────────────────────────────────────────────────
impl<D: LayerData> LayeredMap<D> {
/// Open an existing layered index at `root`.
/// The mode is read once from `PartitionMeta` and applied to all layers.
pub fn open(root: &Path) -> OLMResult<Self> {
let meta = PartitionMeta::load(root)?;
let layers = (0..meta.n_layers)
.map(|i| TypedLayer::<D>::open(&layer_dir(root, i), &meta.mode))
.collect::<OLMResult<Vec<_>>>()?;
Ok(Self { root: root.to_owned(), meta, layers })
}
/// Create a new, empty layered index at `root` with the given mode.
pub fn create(root: &Path, mode: IndexMode) -> OLMResult<Self> {
fs::create_dir_all(root)?;
let meta = PartitionMeta::new(mode);
meta.save(root)?;
Ok(Self { root: root.to_owned(), meta, layers: Vec::new() })
}
pub fn n_layers(&self) -> usize { self.layers.len() }
pub fn layer(&self, i: usize) -> &TypedLayer<D> { &self.layers[i] }
pub fn mode(&self) -> &IndexMode { &self.meta.mode }
/// TypedLayer `i`'s content (`Count`/`Presence`) — a lightweight disk probe,
/// no MPHF/matrix opened. For picking a `D` before committing to
/// [`TypedLayer::open`], or for reporting/diagnostics over an already-open
/// `LayeredMap` without re-opening every layer under a different `D`.
pub fn detect_layer_content(&self, i: usize) -> LayerContent {
LayerContent::detect(&layer_dir(&self.root, i))
}
/// TypedLayer `i`'s storage format — a lightweight disk probe, no
/// MPHF/matrix opened. See [`detect_layer_content`](Self::detect_layer_content).
pub fn detect_layer_storage(&self, i: usize) -> OLMResult<StorageKind> {
let dir = layer_dir(&self.root, i);
match LayerContent::detect(&dir) {
LayerContent::Count => PersistentCompactIntMatrix::detect_storage(&dir).map_err(OLMError::Io),
LayerContent::Presence => PersistentBitMatrix::detect_storage(&dir).map_err(OLMError::Io),
}
}
/// TypedLayer `i`'s evidence mode — a lightweight disk probe, no
/// MPHF/matrix opened. See [`detect_layer_content`](Self::detect_layer_content).
pub fn detect_layer_evidence(&self, i: usize) -> OLMResult<EvidenceKind> {
EvidenceKind::detect(&layer_dir(&self.root, i))
}
/// Query `kmer` across all layers. Returns `(layer_index, Hit)` on match.
pub fn query(&self, kmer: CanonicalKmer) -> Option<(usize, Hit<D::Item>)> {
self.layers
.iter()
.enumerate()
.find_map(|(i, layer)| layer.query(kmer).map(|hit| (i, hit)))
}
pub fn next_layer_writer(&self) -> OLMResult<UnitigFileWriter> {
let dir = layer_dir(&self.root, self.layers.len());
TypedLayer::<D>::unitig_writer(&dir)
}
fn append_layer(&mut self) -> OLMResult<()> {
let i = self.layers.len();
let dir = layer_dir(&self.root, i);
self.layers.push(TypedLayer::<D>::open(&dir, &self.meta.mode)?);
self.meta.n_layers = self.layers.len();
self.meta.save(&self.root)?;
Ok(())
}
}
// ── Mode 1 — set membership ───────────────────────────────────────────────────
impl LayeredMap<()> {
pub fn push_layer(&mut self) -> OLMResult<usize> {
let i = self.layers.len();
let dir = layer_dir(&self.root, i);
TypedLayer::<()>::build(&dir, DEFAULT_BLOCK_BITS, &self.meta.mode)?;
self.append_layer()?;
Ok(i)
}
}
// ── Mode 2 — count matrix ─────────────────────────────────────────────────────
impl LayeredMap<PersistentCompactIntMatrix> {
pub fn push_layer(&mut self, count_of: impl Fn(CanonicalKmer) -> u32) -> OLMResult<usize> {
let i = self.layers.len();
let dir = layer_dir(&self.root, i);
TypedLayer::<PersistentCompactIntMatrix>::build(&dir, DEFAULT_BLOCK_BITS, &self.meta.mode, count_of)?;
self.append_layer()?;
Ok(i)
}
pub fn push_layer_from_map(&mut self, counts: &HashMap<CanonicalKmer, u32>) -> OLMResult<usize> {
self.push_layer(|kmer| counts.get(&kmer).copied().unwrap_or(0))
}
}
#[cfg(test)]
#[path = "tests/map.rs"]
mod tests;
src/obikindex/src/layer/meta.rs
+30 -35
View File
@@ -1,18 +1,20 @@
use std::fs::File;
use std::path::Path;
use serde::{Deserialize, Serialize};
use crate::layer::error::OLMResult;
const META_FILE: &str = "meta.json";
use crate::index::error::OKIResult;
use crate::layer::fingerprint::FingerprintVec;
use crate::layer::mphf_layer::{EvidenceKind, FINGERPRINT_FILE};
use crate::layer::utils::layer_dir;
// ── IndexMode ─────────────────────────────────────────────────────────────────
/// Evidence mode for an entire partitioned index — homogeneous across all layers.
///
/// Determined once at build time; stored in `PartitionMeta` (`meta.json`).
/// All layers within an index share the same mode.
/// Determined once at build time. All layers within an index share the same
/// mode, so it is never persisted separately — [`IndexMode::detect`] reads it
/// straight back from whichever evidence file(s) layer 0 actually has on
/// disk, the same probe [`EvidenceKind::detect`] uses.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum IndexMode {
@@ -20,7 +22,9 @@ pub enum IndexMode {
Exact,
/// Approximate evidence: `fingerprint.bin` only.
/// `b` — fingerprint bits per slot; false-positive rate ≈ 1/2^b per query.
/// `z` — Findere consecutive-kmer parameter (build-time only; not used at query time).
/// `z` — Findere consecutive-kmer parameter (build-time only; not used at
/// query time, not stored on disk, always `0` once recovered via
/// [`detect`](Self::detect)).
Approx { b: u8, z: u8 },
/// Hybrid: both `fingerprint.bin` and `evidence.bin` + `unitigs.bin.idx`.
/// `find()` uses the fingerprint (O(1), approx); `find_strict()` uses exact evidence.
@@ -31,33 +35,24 @@ impl Default for IndexMode {
fn default() -> Self { Self::Exact }
}
// ── PartitionMeta ─────────────────────────────────────────────────────────────
/// Index-level metadata stored in `meta.json` at the root of a partition index.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PartitionMeta {
pub n_layers: usize,
#[serde(default)]
pub mode: IndexMode,
}
impl PartitionMeta {
pub fn new(mode: IndexMode) -> Self {
Self { n_layers: 0, mode }
}
pub fn load(dir: &Path) -> OLMResult<Self> {
let f = File::open(dir.join(META_FILE))?;
Ok(serde_json::from_reader(f)?)
}
pub fn save(&self, dir: &Path) -> OLMResult<()> {
let f = File::create(dir.join(META_FILE))?;
serde_json::to_writer_pretty(f, self)?;
Ok(())
impl IndexMode {
/// Detect the evidence mode actually on disk, from partition index
/// directory `root`'s layer 0 — the single source of truth every layer
/// in the same index shares. `z` cannot be recovered (not stored in
/// `fingerprint.bin`, and irrelevant past build time), so it comes back
/// as `0`.
pub fn detect(root: &Path) -> OKIResult<Self> {
let layer0 = layer_dir(root, 0);
Ok(match EvidenceKind::detect(&layer0)? {
EvidenceKind::Exact => IndexMode::Exact,
EvidenceKind::Approx => IndexMode::Approx {
b: FingerprintVec::open(&layer0.join(FINGERPRINT_FILE))?.b(),
z: 0,
},
EvidenceKind::Hybrid => IndexMode::Hybrid {
b: FingerprintVec::open(&layer0.join(FINGERPRINT_FILE))?.b(),
z: 0,
},
})
}
}
impl Default for PartitionMeta {
fn default() -> Self { Self::new(IndexMode::Exact) }
}
src/obikindex/src/layer/mod.rs
+4 -11
View File
@@ -1,20 +1,13 @@
pub mod content_layer;
pub mod error;
pub mod evidence;
pub mod fingerprint;
pub mod layered_store;
pub mod map;
pub mod meta;
pub(crate) mod mphf_layer;
pub mod typed_layer;
pub(crate) mod utils;
pub use crate::index::error::{OKIError, OKIResult};
pub use content_layer::KmerLayer;
pub use error::{OLMError, OLMResult};
pub use layered_store::LayeredStore;
pub use map::LayeredMap;
pub use meta::{IndexMode, PartitionMeta};
pub use meta::IndexMode;
pub use mphf_layer::{EvidenceKind, KmerBatchIter, KmerIter, MphfLayer, MphfOnly};
pub use typed_layer::{
HasLayerContent, HasStorageKind, Hit, LayerContent, LayerData, TypedLayer,
dereplicated_superkmers_path, layer_dir, open_data, raw_superkmers_path,
};
pub use typed_layer::{HasLayerContent, HasStorageKind, Hit, LayerContent, LayerData, TypedLayer};
src/obikindex/src/layer/mphf_layer.rs
+38 -42
View File
@@ -10,7 +10,7 @@ use obikseq::CanonicalKmer;
use obiskio::{CanonicalKmerIter, UnitigFileReader, UnitigFileWriter, build_unitig_idx};
use ptr_hash::{PtrHash, PtrHashParams, bucket_fn::CubicEps, hash::Xx64};
use crate::layer::error::{OLMError, OLMResult};
use crate::index::error::{OKIError, OKIResult};
use crate::layer::evidence::{Evidence, EvidenceWriter};
use crate::layer::fingerprint::{FingerprintVec, FingerprintVecWriter};
use crate::layer::meta::IndexMode;
@@ -76,7 +76,7 @@ impl LayerEvidence {
/// Identify which evidence files exist at `layer_dir`, without opening
/// any of them — a lightweight disk probe (same signal `EvidenceKind::
/// detect` uses), not a load.
fn at(layer_dir: &Path) -> OLMResult<Self> {
fn at(layer_dir: &Path) -> OKIResult<Self> {
Ok(match EvidenceKind::detect(layer_dir)? {
EvidenceKind::Exact => LayerEvidence::Exact {
evidence_path: layer_dir.join(EVIDENCE_FILE),
@@ -94,7 +94,7 @@ impl LayerEvidence {
}
/// mmap the evidence file(s) this identity points to.
fn open(self) -> OLMResult<OpenLayerEvidence> {
fn open(self) -> OKIResult<OpenLayerEvidence> {
Ok(match self {
LayerEvidence::Exact { evidence_path, unitig_path } => OpenLayerEvidence::Exact {
evidence: Evidence::open(&evidence_path)?,
@@ -127,17 +127,17 @@ impl EvidenceKind {
/// Lightweight disk probe, no MPHF/evidence opened: presence of
/// `evidence.bin`/`fingerprint.bin` alone determines the mode — same
/// signal `MphfLayer::open` uses, just without opening either file.
pub fn detect(layer_dir: &Path) -> OLMResult<EvidenceKind> {
pub fn detect(layer_dir: &Path) -> OKIResult<EvidenceKind> {
let has_evidence = layer_dir.join(EVIDENCE_FILE).exists();
let has_fingerprint = layer_dir.join(FINGERPRINT_FILE).exists();
match (has_evidence, has_fingerprint) {
(true, false) => Ok(EvidenceKind::Exact),
(false, true) => Ok(EvidenceKind::Approx),
(true, true) => Ok(EvidenceKind::Hybrid),
(false, false) => Err(OLMError::InvalidLayer(format!(
(false, false) => Err(OKIError::InvalidData { context: "layer", detail: format!(
"no evidence.bin or fingerprint.bin in {}",
layer_dir.display()
))),
) }),
}
}
}
@@ -151,7 +151,6 @@ impl EvidenceKind {
/// - [`find_strict`](Self::find_strict) — always exact; O(1) on Exact/Hybrid layers,
/// O(n) sequential scan on Approx layers.
pub struct MphfLayer {
layer_dir: PathBuf,
/// Mode-independent copy — used by `iter_kmers` (all modes) and
/// `find_strict`'s `Approx` fallback. `Exact`/`Hybrid` additionally
/// carry their own inside `OpenLayerEvidence` (see its doc comment for
@@ -165,18 +164,15 @@ pub struct MphfLayer {
}
impl MphfLayer {
/// This layer's own directory.
pub(crate) fn dir(&self) -> &Path {
&self.layer_dir
}
/// Open a layer — evidence mode is auto-detected from which files exist
/// on disk (see [`LayerEvidence::at`]), never passed in: the caller
/// can't tell this type anything about its own mode that isn't already
/// derivable from `layer_dir` itself.
pub fn open(dir: &Path) -> OLMResult<Self> {
/// derivable from `dir` itself. `dir` is only used transiently here —
/// not stored: `KmerLayer`, one level up, is what owns a layer's
/// directory as a persistent field, not this type.
pub fn open(dir: &Path) -> OKIResult<Self> {
let mphf: MemCase<MphfEps> = Mphf::mmap(&dir.join(MPHF_FILE), Flags::empty())
.map_err(|e| OLMError::InvalidLayer(e.to_string()))?;
.map_err(|e| OKIError::InvalidData { context: "layer", detail: e.to_string() })?;
let unitigs = Arc::new(UnitigFileReader::open(&dir.join(UNITIGS_FILE))?);
let ev = LayerEvidence::at(dir)?.open()?;
let n = match &ev {
@@ -184,7 +180,7 @@ impl MphfLayer {
OpenLayerEvidence::Hybrid { evidence, .. } => evidence.len(),
OpenLayerEvidence::Approx { fingerprint } => fingerprint.n(),
};
Ok(Self { layer_dir: dir.to_owned(), unitigs, mphf, ev, n })
Ok(Self { unitigs, mphf, ev, n })
}
// ── Query API ─────────────────────────────────────────────────────────────
@@ -405,9 +401,9 @@ impl Iterator for KmerBatchIter {
pub struct MphfOnly(MemCase<MphfEps>);
impl MphfOnly {
pub fn open(dir: &Path) -> OLMResult<Self> {
pub fn open(dir: &Path) -> OKIResult<Self> {
let mphf: MemCase<MphfEps> = Mphf::mmap(&dir.join(MPHF_FILE), Flags::empty())
.map_err(|e| OLMError::InvalidLayer(e.to_string()))?;
.map_err(|e| OKIError::InvalidData { context: "layer", detail: e.to_string() })?;
Ok(Self(mphf))
}
@@ -421,13 +417,13 @@ impl MphfOnly {
impl MphfLayer {
// ── Build helpers ─────────────────────────────────────────────────────────
pub fn unitig_writer(dir: &Path) -> OLMResult<UnitigFileWriter> {
pub fn unitig_writer(dir: &Path) -> OKIResult<UnitigFileWriter> {
fs::create_dir_all(dir)?;
Ok(UnitigFileWriter::create(&dir.join(UNITIGS_FILE))?)
}
/// Build `evidence.bin` + `unitigs.bin.idx` from `unitigs.bin` + `mphf.bin`.
pub fn build_exact_evidence(dir: &Path, block_bits: u8) -> OLMResult<usize> {
pub fn build_exact_evidence(dir: &Path, block_bits: u8) -> OKIResult<usize> {
let unitig_path = dir.join(UNITIGS_FILE);
let unitigs = UnitigFileReader::open_sequential(&unitig_path)?;
let n = unitigs.n_kmers();
@@ -439,7 +435,7 @@ impl MphfLayer {
}
let mphf: Mphf = Mphf::load_full(&dir.join(MPHF_FILE))
.map_err(|e| OLMError::InvalidLayer(e.to_string()))?;
.map_err(|e| OKIError::InvalidData { context: "layer", detail: e.to_string() })?;
let mut ev = EvidenceWriter::new(n);
let mut seen = vec![0u8; (n + 7) / 8];
@@ -447,12 +443,12 @@ impl MphfLayer {
for (kmer, chunk_id, rank) in unitigs.iter_indexed_canonical_kmers() {
let slot = mphf.index(&kmer.raw());
if slot >= n {
return Err(OLMError::Mphf("slot out of bounds".into()));
return Err(OKIError::InvalidData { context: "mphf", detail: "slot out of bounds".into() });
}
let byte = slot / 8;
let bit = 1u8 << (slot % 8);
if seen[byte] & bit != 0 {
return Err(OLMError::Mphf("duplicate slot".into()));
return Err(OKIError::InvalidData { context: "mphf", detail: "duplicate slot".into() });
}
seen[byte] |= bit;
ev.set(slot, chunk_id as u32, rank as u8);
@@ -464,14 +460,14 @@ impl MphfLayer {
}
/// Build `fingerprint.bin` from `unitigs.bin` + `mphf.bin`.
pub fn build_approx_evidence(dir: &Path, b: u8, z: u8) -> OLMResult<usize> {
pub fn build_approx_evidence(dir: &Path, b: u8, z: u8) -> OKIResult<usize> {
if b == 0 || b > 64 {
return Err(OLMError::InvalidLayer(
return Err(OKIError::InvalidData { context: "layer", detail:
"fingerprint width must be 1..=64".into(),
));
});
}
if z == 0 {
return Err(OLMError::InvalidLayer("z must be ≥ 1".into()));
return Err(OKIError::InvalidData { context: "layer", detail: "z must be ≥ 1".into() });
}
let unitig_path = dir.join(UNITIGS_FILE);
let unitigs = UnitigFileReader::open_sequential(&unitig_path)?;
@@ -483,14 +479,14 @@ impl MphfLayer {
}
let mphf: Mphf = Mphf::load_full(&dir.join(MPHF_FILE))
.map_err(|e| OLMError::InvalidLayer(e.to_string()))?;
.map_err(|e| OKIError::InvalidData { context: "layer", detail: e.to_string() })?;
let mut fw = FingerprintVecWriter::new(n, b);
for (kmer, _, _) in unitigs.iter_indexed_canonical_kmers() {
let slot = mphf.index(&kmer.raw());
if slot >= n {
return Err(OLMError::Mphf("slot out of bounds".into()));
return Err(OKIError::InvalidData { context: "mphf", detail: "slot out of bounds".into() });
}
fw.set(slot, kmer.seq_hash());
}
@@ -508,24 +504,24 @@ impl MphfLayer {
dir: &Path,
block_bits: u8,
mode: &IndexMode,
fill_slot: &mut impl FnMut(usize, CanonicalKmer) -> OLMResult<()>,
) -> OLMResult<usize> {
fill_slot: &mut impl FnMut(usize, CanonicalKmer) -> OKIResult<()>,
) -> OKIResult<usize> {
use rayon::prelude::*;
let unitig_path = dir.join(UNITIGS_FILE);
let n = UnitigFileReader::open_sequential(&unitig_path)?.n_kmers();
let sk_to_olm = |e: obiskio::SKError| match e {
obiskio::SKError::Io(io) => OLMError::Io(io),
e => OLMError::InvalidLayer(e.to_string()),
obiskio::SKError::Io(io) => OKIError::Io(io),
e => OKIError::InvalidData { context: "layer", detail: e.to_string() },
};
// ── Empty layer ───────────────────────────────────────────────────────
if n == 0 {
let mphf: Mphf = Mphf::try_new(&[] as &[u64], PtrHashParams::<CubicEps>::default())
.ok_or_else(|| OLMError::Mphf("construction failed".into()))?;
.ok_or_else(|| OKIError::InvalidData { context: "mphf", detail: "construction failed".into() })?;
mphf.store(&dir.join(MPHF_FILE))
.map_err(|e| OLMError::InvalidLayer(e.to_string()))?;
.map_err(|e| OKIError::InvalidData { context: "layer", detail: e.to_string() })?;
match mode {
IndexMode::Exact | IndexMode::Hybrid { .. } => {
fs::File::create(dir.join(EVIDENCE_FILE))?;
@@ -549,7 +545,7 @@ impl MphfLayer {
PtrHashParams::<CubicEps>::default(),
);
mphf.store(&dir.join(MPHF_FILE))
.map_err(|e| OLMError::InvalidLayer(e.to_string()))?;
.map_err(|e| OKIError::InvalidData { context: "layer", detail: e.to_string() })?;
// ── Pass 2: fill evidence files + callback ────────────────────────────
let unitigs2 = UnitigFileReader::open_sequential(&unitig_path)?;
@@ -561,12 +557,12 @@ impl MphfLayer {
for (kmer, chunk_id, rank) in unitigs2.iter_indexed_canonical_kmers() {
let slot = mphf.index(&kmer.raw());
if slot >= n {
return Err(OLMError::Mphf("slot out of bounds".into()));
return Err(OKIError::InvalidData { context: "mphf", detail: "slot out of bounds".into() });
}
let byte = slot / 8;
let bit = 1u8 << (slot % 8);
if seen[byte] & bit != 0 {
return Err(OLMError::Mphf("duplicate slot".into()));
return Err(OKIError::InvalidData { context: "mphf", detail: "duplicate slot".into() });
}
seen[byte] |= bit;
ev.set(slot, chunk_id as u32, rank as u8);
@@ -581,12 +577,12 @@ impl MphfLayer {
for (kmer, _, _) in unitigs2.iter_indexed_canonical_kmers() {
let slot = mphf.index(&kmer.raw());
if slot >= n {
return Err(OLMError::Mphf("slot out of bounds".into()));
return Err(OKIError::InvalidData { context: "mphf", detail: "slot out of bounds".into() });
}
let byte = slot / 8;
let bit = 1u8 << (slot % 8);
if seen[byte] & bit != 0 {
return Err(OLMError::Mphf("duplicate slot".into()));
return Err(OKIError::InvalidData { context: "mphf", detail: "duplicate slot".into() });
}
seen[byte] |= bit;
fw.set(slot, kmer.seq_hash());
@@ -601,12 +597,12 @@ impl MphfLayer {
for (kmer, chunk_id, rank) in unitigs2.iter_indexed_canonical_kmers() {
let slot = mphf.index(&kmer.raw());
if slot >= n {
return Err(OLMError::Mphf("slot out of bounds".into()));
return Err(OKIError::InvalidData { context: "mphf", detail: "slot out of bounds".into() });
}
let byte = slot / 8;
let bit = 1u8 << (slot % 8);
if seen[byte] & bit != 0 {
return Err(OLMError::Mphf("duplicate slot".into()));
return Err(OKIError::InvalidData { context: "mphf", detail: "duplicate slot".into() });
}
seen[byte] |= bit;
ev.set(slot, chunk_id as u32, rank as u8);
src/obikindex/src/layer/tests/layered_store.rs
-381
View File
@@ -1,381 +0,0 @@
use super::*;
use obicompactvec::{
PersistentBitMatrix, PersistentBitMatrixBuilder,
PersistentCompactIntMatrix, PersistentCompactIntMatrixBuilder,
};
use tempfile::tempdir;
fn make_int_matrix(cols: &[&[u32]]) -> (tempfile::TempDir, PersistentCompactIntMatrix) {
let n = cols.first().map_or(0, |c| c.len());
let dir = tempdir().unwrap();
let mut b = PersistentCompactIntMatrixBuilder::new(n, &dir.path().join("counts")).unwrap();
for &col in cols {
let mut cb = b.add_col().unwrap();
for (slot, &v) in col.iter().enumerate() { cb.set(slot, v); }
cb.close().unwrap();
}
b.close().unwrap();
let m = PersistentCompactIntMatrix::open(dir.path()).unwrap();
(dir, m)
}
fn make_bit_matrix(cols: &[&[bool]]) -> (tempfile::TempDir, PersistentBitMatrix) {
let n = cols.first().map_or(0, |c| c.len());
let dir = tempdir().unwrap();
let mut b = PersistentBitMatrixBuilder::new(n, &dir.path().join("presence")).unwrap();
for &col in cols {
let mut cb = b.add_col().unwrap();
for (slot, &v) in col.iter().enumerate() { cb.set(slot, v); }
cb.close().unwrap();
}
b.close().unwrap();
let m = PersistentBitMatrix::open(dir.path()).unwrap();
(dir, m)
}
// ── ColumnWeights ─────────────────────────────────────────────────────────
#[test]
fn col_weights_sums_across_layers() {
// layer 0: col0=[1,2], col1=[3,4] → weights [3, 7]
// layer 1: col0=[10,0], col1=[0,10] → weights [10, 10]
// combined: [13, 17]
let (_d0, m0) = make_int_matrix(&[&[1, 2], &[3, 4]]);
let (_d1, m1) = make_int_matrix(&[&[10, 0], &[0, 10]]);
let store = LayeredStore::new(vec![m0, m1]);
let w = store.col_weights();
assert_eq!(w[0], 13);
assert_eq!(w[1], 17);
}
#[test]
fn col_weights_bit_sums_across_layers() {
// layer 0: col0=[T,F,T], col1=[F,T,T] → counts [2, 2]
// layer 1: col0=[F,F,T], col1=[T,T,F] → counts [1, 2]
// combined: [3, 4]
let (_d0, m0) = make_bit_matrix(&[&[true, false, true], &[false, true, true]]);
let (_d1, m1) = make_bit_matrix(&[&[false, false, true], &[true, true, false]]);
let store = LayeredStore::new(vec![m0, m1]);
let w = store.col_weights();
assert_eq!(w[0], 3);
assert_eq!(w[1], 4);
}
// ── CountPartials — layered (one partition) ───────────────────────────────
#[test]
fn layered_bray_matches_combined() {
// Split [1,2,3,4,5] across two layers; bray dist should equal direct computation
// on [1,2,3,4,5] for each column pair.
// col0=[1,2,3,4,5], col1=[5,4,3,2,1]
let (_d0, m0) = make_int_matrix(&[&[1, 2], &[5, 4]]); // slots 0-1
let (_d1, m1) = make_int_matrix(&[&[3, 4, 5], &[3, 2, 1]]); // slots 2-4
let store = LayeredStore::new(vec![m0, m1]);
// direct on full data
let (_df, mf) = make_int_matrix(&[&[1, 2, 3, 4, 5], &[5, 4, 3, 2, 1]]);
let expected = CountPartials::bray_dist_matrix(&mf);
let got = CountPartials::bray_dist_matrix(&store);
assert!((got[[0, 1]] - expected[[0, 1]]).abs() < 1e-12, "bray [0,1]");
assert!((got[[1, 0]] - expected[[1, 0]]).abs() < 1e-12, "bray [1,0]");
}
#[test]
fn layered_relfreq_bray_matches_combined() {
let (_d0, m0) = make_int_matrix(&[&[1, 2], &[5, 4]]);
let (_d1, m1) = make_int_matrix(&[&[3, 4, 5], &[3, 2, 1]]);
let store = LayeredStore::new(vec![m0, m1]);
let (_df, mf) = make_int_matrix(&[&[1, 2, 3, 4, 5], &[5, 4, 3, 2, 1]]);
let expected = CountPartials::relfreq_bray_dist_matrix(&mf);
let got = CountPartials::relfreq_bray_dist_matrix(&store);
assert!((got[[0, 1]] - expected[[0, 1]]).abs() < 1e-12, "relfreq_bray [0,1]");
}
#[test]
fn layered_euclidean_matches_combined() {
let (_d0, m0) = make_int_matrix(&[&[3, 0], &[0, 4]]);
let (_d1, m1) = make_int_matrix(&[&[1, 1], &[2, 2]]);
let store = LayeredStore::new(vec![m0, m1]);
let (_df, mf) = make_int_matrix(&[&[3, 0, 1, 1], &[0, 4, 2, 2]]);
let expected = CountPartials::euclidean_dist_matrix(&mf);
let got = CountPartials::euclidean_dist_matrix(&store);
assert!((got[[0, 1]] - expected[[0, 1]]).abs() < 1e-12, "euclidean [0,1]");
}
// ── CountPartials — partitioned (LayeredStore<LayeredStore<_>>) ───────────
#[test]
fn partitioned_bray_matches_combined() {
// partition 0: slots [1,2,3,4,5] col0 vs col1
// partition 1: slots [10,20] col0 vs col1
let (_d0, p0) = make_int_matrix(&[&[1, 2, 3, 4, 5], &[5, 4, 3, 2, 1]]);
let (_d1, p1) = make_int_matrix(&[&[10, 20], &[20, 10]]);
let partitioned = LayeredStore::new(vec![
LayeredStore::new(vec![p0]),
LayeredStore::new(vec![p1]),
]);
let (_df, mf) = make_int_matrix(&[&[1, 2, 3, 4, 5, 10, 20], &[5, 4, 3, 2, 1, 20, 10]]);
let expected = CountPartials::bray_dist_matrix(&mf);
let got = CountPartials::bray_dist_matrix(&partitioned);
assert!((got[[0, 1]] - expected[[0, 1]]).abs() < 1e-12, "partitioned bray [0,1]");
}
#[test]
fn partitioned_threshold_jaccard_off_diagonal_is_pairwise() {
// 3 genomes, 2 partitions, 1 layer each — mirrors distance.rs's
// LayeredStore<LayeredStore<PersistentCompactIntMatrix>> shape.
// partition 0: col0=[3,0], col1=[0,3], col2=[3,3]
// partition 1: col0=[1,1], col1=[1,0], col2=[0,1]
let (_d0, p0) = make_int_matrix(&[&[3, 0], &[0, 3], &[3, 3]]);
let (_d1, p1) = make_int_matrix(&[&[1, 1], &[1, 0], &[0, 1]]);
let partitioned = LayeredStore::new(vec![
LayeredStore::new(vec![p0]),
LayeredStore::new(vec![p1]),
]);
let (_df, mf) = make_int_matrix(&[&[3, 0, 1, 1], &[0, 3, 1, 0], &[3, 3, 0, 1]]);
let threshold = 1u32;
let (inter_p, union_p) = CountPartials::partial_threshold_jaccard(&partitioned, threshold);
let (inter_f, union_f) = CountPartials::partial_threshold_jaccard(&mf, threshold);
let n = 3;
for i in 0..n {
for j in 0..n {
assert_eq!(inter_p[[i, j]], inter_f[[i, j]], "inter[{i},{j}]");
assert_eq!(union_p[[i, j]], union_f[[i, j]], "union[{i},{j}]");
}
}
}
#[test]
fn partitioned_threshold_jaccard_packed_off_diagonal_is_pairwise() {
// Same as `partitioned_threshold_jaccard_off_diagonal_is_pairwise` but
// each partition matrix is packed into a single .pcmx file first —
// the on-disk format actually used in production after `pack_matrices`.
use obicompactvec::pack_compact_int_matrix;
let (d0, _p0) = make_int_matrix(&[&[3, 0], &[0, 3], &[3, 3]]);
pack_compact_int_matrix(&d0.path().join("counts")).unwrap();
let p0 = PersistentCompactIntMatrix::open(d0.path()).unwrap();
let (d1, _p1) = make_int_matrix(&[&[1, 1], &[1, 0], &[0, 1]]);
pack_compact_int_matrix(&d1.path().join("counts")).unwrap();
let p1 = PersistentCompactIntMatrix::open(d1.path()).unwrap();
let partitioned = LayeredStore::new(vec![
LayeredStore::new(vec![p0]),
LayeredStore::new(vec![p1]),
]);
let (_df, mf) = make_int_matrix(&[&[3, 0, 1, 1], &[0, 3, 1, 0], &[3, 3, 0, 1]]);
let threshold = 1u32;
let (inter_p, union_p) = CountPartials::partial_threshold_jaccard(&partitioned, threshold);
let (inter_f, union_f) = CountPartials::partial_threshold_jaccard(&mf, threshold);
let n = 3;
for i in 0..n {
for j in 0..n {
assert_eq!(inter_p[[i, j]], inter_f[[i, j]], "inter[{i},{j}]");
assert_eq!(union_p[[i, j]], union_f[[i, j]], "union[{i},{j}]");
}
}
}
#[test]
fn partitioned_multilayer_threshold_jaccard_off_diagonal_is_pairwise() {
// 2 partitions, 2 layers each — the shape production indexes actually
// have (MPHF collision layers within a partition).
// partition 0, layer 0: col0=[3,0], col1=[0,3], col2=[3,3]
// partition 0, layer 1: col0=[2,0], col1=[0,0], col2=[2,0]
// partition 1, layer 0: col0=[1,1], col1=[1,0], col2=[0,1]
// partition 1, layer 1: col0=[0,5], col1=[5,5], col2=[0,0]
let (_d0a, p0a) = make_int_matrix(&[&[3, 0], &[0, 3], &[3, 3]]);
let (_d0b, p0b) = make_int_matrix(&[&[2, 0], &[0, 0], &[2, 0]]);
let (_d1a, p1a) = make_int_matrix(&[&[1, 1], &[1, 0], &[0, 1]]);
let (_d1b, p1b) = make_int_matrix(&[&[0, 5], &[5, 5], &[0, 0]]);
let partitioned = LayeredStore::new(vec![
LayeredStore::new(vec![p0a, p0b]),
LayeredStore::new(vec![p1a, p1b]),
]);
// Flattened equivalent: concatenate every layer's slots into one matrix.
let (_df, mf) = make_int_matrix(&[
&[3, 0, 2, 0, 1, 1, 0, 5],
&[0, 3, 0, 0, 1, 0, 5, 5],
&[3, 3, 2, 0, 0, 1, 0, 0],
]);
let threshold = 1u32;
let (inter_p, union_p) = CountPartials::partial_threshold_jaccard(&partitioned, threshold);
let (inter_f, union_f) = CountPartials::partial_threshold_jaccard(&mf, threshold);
let n = 3;
for i in 0..n {
for j in 0..n {
assert_eq!(inter_p[[i, j]], inter_f[[i, j]], "inter[{i},{j}]");
assert_eq!(union_p[[i, j]], union_f[[i, j]], "union[{i},{j}]");
}
}
}
// ── BitPartials ───────────────────────────────────────────────────────────
#[test]
fn layered_jaccard_matches_combined() {
let (_d0, m0) = make_bit_matrix(&[&[true, false], &[false, true]]);
let (_d1, m1) = make_bit_matrix(&[&[true, true], &[true, false]]);
let store = LayeredStore::new(vec![m0, m1]);
let (_df, mf) = make_bit_matrix(&[
&[true, false, true, true],
&[false, true, true, false],
]);
let expected = BitPartials::jaccard_dist_matrix(&mf);
let got = BitPartials::jaccard_dist_matrix(&store);
assert!((got[[0, 1]] - expected[[0, 1]]).abs() < 1e-12, "jaccard [0,1]");
}
#[test]
fn layered_hamming_matches_combined() {
let (_d0, m0) = make_bit_matrix(&[&[true, false], &[false, true]]);
let (_d1, m1) = make_bit_matrix(&[&[true, true], &[false, false]]);
let store = LayeredStore::new(vec![m0, m1]);
let (_df, mf) = make_bit_matrix(&[
&[true, false, true, true],
&[false, true, false, false],
]);
let expected = BitPartials::hamming_dist_matrix(&mf);
let got = BitPartials::hamming_dist_matrix(&store);
assert_eq!(got[[0, 1]], expected[[0, 1]], "hamming [0,1]");
}
#[test]
fn partitioned_bit_jaccard_off_diagonal_is_pairwise() {
// Same shape as the count-based `partitioned_multilayer_threshold_jaccard_*`
// tests, but for the presence/bit path (`with_counts = false` — what
// `all_specifics` actually uses in production).
// 4 genomes, 3 partitions, 2 layers in the last one.
let (_d0, p0) = make_bit_matrix(&[
&[true, false, true],
&[false, true, true],
&[true, true, false],
&[false, false, true],
]);
let (_d1, p1) = make_bit_matrix(&[
&[true, true],
&[false, true],
&[true, false],
&[true, true],
]);
let (_d2a, p2a) = make_bit_matrix(&[
&[false, true],
&[true, true],
&[false, false],
&[true, false],
]);
let (_d2b, p2b) = make_bit_matrix(&[
&[true],
&[false],
&[true],
&[true],
]);
let partitioned = LayeredStore::new(vec![
LayeredStore::new(vec![p0]),
LayeredStore::new(vec![p1]),
LayeredStore::new(vec![p2a, p2b]),
]);
// Flattened equivalent: concatenate every partition/layer's slots.
let (_df, mf) = make_bit_matrix(&[
&[true, false, true, true, true, false, true, true],
&[false, true, true, false, true, true, true, false],
&[true, true, false, true, false, false, false, true],
&[false, false, true, true, true, true, false, true],
]);
let (inter_p, union_p) = BitPartials::partial_jaccard(&partitioned);
let (inter_f, union_f) = BitPartials::partial_jaccard(&mf);
let n = 4;
for i in 0..n {
for j in 0..n {
assert_eq!(inter_p[[i, j]], inter_f[[i, j]], "inter[{i},{j}]");
assert_eq!(union_p[[i, j]], union_f[[i, j]], "union[{i},{j}]");
}
}
}
#[test]
fn partitioned_bit_jaccard_packed_off_diagonal_is_pairwise() {
// Same as `partitioned_bit_jaccard_off_diagonal_is_pairwise` but every
// partition's presence matrix is packed into a single .pbmx file —
// the on-disk format actually used in production after `pack_matrices`.
use obicompactvec::pack_bit_matrix;
let (d0, _p0) = make_bit_matrix(&[
&[true, false, true],
&[false, true, true],
&[true, true, false],
&[false, false, true],
]);
pack_bit_matrix(&d0.path().join("presence")).unwrap();
let p0 = PersistentBitMatrix::open(d0.path()).unwrap();
let (d1, _p1) = make_bit_matrix(&[
&[true, true],
&[false, true],
&[true, false],
&[true, true],
]);
pack_bit_matrix(&d1.path().join("presence")).unwrap();
let p1 = PersistentBitMatrix::open(d1.path()).unwrap();
let (d2a, _p2a) = make_bit_matrix(&[
&[false, true],
&[true, true],
&[false, false],
&[true, false],
]);
pack_bit_matrix(&d2a.path().join("presence")).unwrap();
let p2a = PersistentBitMatrix::open(d2a.path()).unwrap();
let (d2b, _p2b) = make_bit_matrix(&[
&[true],
&[false],
&[true],
&[true],
]);
pack_bit_matrix(&d2b.path().join("presence")).unwrap();
let p2b = PersistentBitMatrix::open(d2b.path()).unwrap();
let partitioned = LayeredStore::new(vec![
LayeredStore::new(vec![p0]),
LayeredStore::new(vec![p1]),
LayeredStore::new(vec![p2a, p2b]),
]);
let (_df, mf) = make_bit_matrix(&[
&[true, false, true, true, true, false, true, true],
&[false, true, true, false, true, true, true, false],
&[true, true, false, true, false, false, false, true],
&[false, false, true, true, true, true, false, true],
]);
let (inter_p, union_p) = BitPartials::partial_jaccard(&partitioned);
let (inter_f, union_f) = BitPartials::partial_jaccard(&mf);
let n = 4;
for i in 0..n {
for j in 0..n {
assert_eq!(inter_p[[i, j]], inter_f[[i, j]], "inter[{i},{j}]");
assert_eq!(union_p[[i, j]], union_f[[i, j]], "union[{i},{j}]");
}
}
}
src/obikindex/src/layer/tests/map.rs
-219
View File
@@ -1,219 +0,0 @@
use super::*;
use obicompactvec::{pack_bit_matrix, pack_sparse_bit_matrix, PersistentBitMatrix, PersistentCompactIntMatrix, StorageKind};
use obikseq::{set_k, Sequence as _, Unitig};
use obiskio::DEFAULT_BLOCK_BITS;
use crate::layer::typed_layer::LayerContent;
use crate::layer::meta::IndexMode;
use crate::layer::mphf_layer::EvidenceKind;
use tempfile::tempdir;
fn push_unitigs_and_layer(
map: &mut LayeredMap<PersistentCompactIntMatrix>,
seqs: &[&[u8]],
count: u32,
) {
let mut w = map.next_layer_writer().unwrap();
for s in seqs {
w.write(&Unitig::from_ascii(s)).unwrap();
}
w.close().unwrap();
map.push_layer(|_| count).unwrap();
}
fn canonical(ascii: &[u8]) -> CanonicalKmer {
obikseq::Kmer::from_ascii(ascii).unwrap().canonical()
}
#[test]
fn create_empty_map() {
set_k(4);
let dir = tempdir().unwrap();
let map = LayeredMap::<()>::create(dir.path(), IndexMode::Exact).unwrap();
assert_eq!(map.n_layers(), 0);
}
#[test]
fn open_reloads_layer_count() {
set_k(4);
let dir = tempdir().unwrap();
{
let mut map = LayeredMap::<PersistentCompactIntMatrix>::create(dir.path(), IndexMode::Exact).unwrap();
push_unitigs_and_layer(&mut map, &[b"AAAACGT"], 1);
}
let map = LayeredMap::<PersistentCompactIntMatrix>::open(dir.path()).unwrap();
assert_eq!(map.n_layers(), 1);
}
#[test]
fn query_finds_kmer_in_layer_zero() {
set_k(4);
let dir = tempdir().unwrap();
let mut map = LayeredMap::<PersistentCompactIntMatrix>::create(dir.path(), IndexMode::Exact).unwrap();
push_unitigs_and_layer(&mut map, &[b"AAAACGT"], 3);
let kmer = canonical(b"AAAC");
let (layer_idx, hit) = map.query(kmer).expect("kmer must be found");
assert_eq!(layer_idx, 0);
assert_eq!(hit.data[0], 3);
}
#[test]
fn query_finds_kmer_in_correct_layer() {
set_k(4);
let dir = tempdir().unwrap();
let mut map = LayeredMap::<PersistentCompactIntMatrix>::create(dir.path(), IndexMode::Exact).unwrap();
push_unitigs_and_layer(&mut map, &[b"AAAACGT"], 1);
push_unitigs_and_layer(&mut map, &[b"GGGACGT"], 2);
assert_eq!(map.n_layers(), 2);
let (li, hit) = map.query(canonical(b"AAAA")).expect("AAAA must be found");
assert_eq!(li, 0);
assert_eq!(hit.data[0], 1);
let (li, hit) = map.query(canonical(b"GGGA")).expect("GGGA must be found");
assert_eq!(li, 1);
assert_eq!(hit.data[0], 2);
}
#[test]
fn query_absent_returns_none() {
set_k(4);
let dir = tempdir().unwrap();
let mut map = LayeredMap::<PersistentCompactIntMatrix>::create(dir.path(), IndexMode::Exact).unwrap();
push_unitigs_and_layer(&mut map, &[b"AAAACGT"], 1);
let absent = canonical(b"CCCC");
assert!(map.query(absent).is_none());
}
#[test]
fn push_layer_from_map_convenience() {
set_k(4);
let dir = tempdir().unwrap();
let mut map = LayeredMap::<PersistentCompactIntMatrix>::create(dir.path(), IndexMode::Exact).unwrap();
let mut w = map.next_layer_writer().unwrap();
w.write(&Unitig::from_ascii(b"AAAACGT")).unwrap();
w.close().unwrap();
let counts: HashMap<CanonicalKmer, u32> = vec![
(canonical(b"AAAA"), 10u32),
].into_iter().collect();
map.push_layer_from_map(&counts).unwrap();
let (_, hit) = map.query(canonical(b"AAAA")).unwrap();
assert_eq!(hit.data[0], 10);
}
// ── detect_layer_content / detect_layer_storage / detect_layer_evidence ────
//
// All built directly on disk (bypassing `LayeredMap<D>::push_layer`, which
// only exists for modes 1 and 2 — mode 3/presence has no `push_layer`),
// then reopened as `LayeredMap<()>` — `()`'s `LayerData::open` never reads
// the matrix, so it's the right handle for exercising the disk-probing
// `detect_layer_*` methods regardless of what content/storage the layer
// underneath actually is.
fn write_unitigs_at(dir: &Path, seqs: &[&[u8]]) {
fs::create_dir_all(dir).unwrap();
let mut w = UnitigFileWriter::create(&dir.join(crate::layer::typed_layer::UNITIGS_FILE)).unwrap();
for s in seqs {
w.write(&Unitig::from_ascii(s)).unwrap();
}
w.close().unwrap();
}
/// Build a one-layer partition root with a presence layer at `layer_0`,
/// save `meta.json`, and reopen it as `LayeredMap<()>`.
fn presence_partition(seqs: &[&[u8]], n_genomes: usize) -> (tempfile::TempDir, LayeredMap<()>) {
let dir = tempdir().unwrap();
let mode = IndexMode::Exact;
let layer0 = layer_dir(dir.path(), 0);
write_unitigs_at(&layer0, seqs);
TypedLayer::<PersistentBitMatrix>::build_presence(&layer0, DEFAULT_BLOCK_BITS, &mode, n_genomes, |kmer, g| {
(kmer.raw().wrapping_add(g as u64)) % 2 == 0
}).unwrap();
PartitionMeta { n_layers: 1, mode }.save(dir.path()).unwrap();
let map = LayeredMap::<()>::open(dir.path()).unwrap();
(dir, map)
}
#[test]
fn detect_layer_content_count() {
set_k(4);
let dir = tempdir().unwrap();
let mut map = LayeredMap::<PersistentCompactIntMatrix>::create(dir.path(), IndexMode::Exact).unwrap();
push_unitigs_and_layer(&mut map, &[b"AAAACGT"], 3);
assert_eq!(map.detect_layer_content(0), LayerContent::Count);
}
#[test]
fn detect_layer_content_presence() {
set_k(4);
let (_dir, map) = presence_partition(&[b"AAATCTA", b"CTTCGCC", b"TGATACG"], 3);
assert_eq!(map.detect_layer_content(0), LayerContent::Presence);
}
#[test]
fn detect_layer_storage_columnar_then_packed_then_sparse() {
set_k(4);
let (dir, map) = presence_partition(&[b"AAATCTA", b"CTTCGCC", b"TGATACG"], 3);
let presence_dir = layer_dir(dir.path(), 0).join("presence");
assert_eq!(map.detect_layer_storage(0).unwrap(), StorageKind::Columnar);
pack_bit_matrix(&presence_dir).unwrap();
assert_eq!(map.detect_layer_storage(0).unwrap(), StorageKind::Packed);
pack_sparse_bit_matrix(&presence_dir).unwrap();
assert_eq!(map.detect_layer_storage(0).unwrap(), StorageKind::Sparse);
}
#[test]
fn detect_layer_storage_implicit() {
set_k(4);
let dir = tempdir().unwrap();
let mode = IndexMode::Exact;
let layer0 = layer_dir(dir.path(), 0);
write_unitigs_at(&layer0, &[b"AAAACGT"]);
// Mode-1 build: MPHF + evidence + `layer_meta.json`, no matrix at all —
// must read back as `Presence`/`Implicit`, never a third "empty" state.
TypedLayer::<()>::build(&layer0, DEFAULT_BLOCK_BITS, &mode).unwrap();
PartitionMeta { n_layers: 1, mode }.save(dir.path()).unwrap();
let map = LayeredMap::<()>::open(dir.path()).unwrap();
assert_eq!(map.detect_layer_content(0), LayerContent::Presence);
assert_eq!(map.detect_layer_storage(0).unwrap(), StorageKind::Implicit);
}
#[test]
fn detect_layer_storage_count_columnar_then_packed() {
set_k(4);
let dir = tempdir().unwrap();
let mut map = LayeredMap::<PersistentCompactIntMatrix>::create(dir.path(), IndexMode::Exact).unwrap();
push_unitigs_and_layer(&mut map, &[b"AAAACGT"], 3);
drop(map);
let map = LayeredMap::<()>::open(dir.path()).unwrap();
assert_eq!(map.detect_layer_storage(0).unwrap(), StorageKind::Columnar);
obicompactvec::pack_compact_int_matrix(&layer_dir(dir.path(), 0).join("counts")).unwrap();
assert_eq!(map.detect_layer_storage(0).unwrap(), StorageKind::Packed);
}
#[test]
fn detect_layer_evidence_exact_vs_approx() {
set_k(4);
let dir = tempdir().unwrap();
let layer0 = layer_dir(dir.path(), 0);
write_unitigs_at(&layer0, &[b"AAAACGT"]);
TypedLayer::<()>::build(&layer0, DEFAULT_BLOCK_BITS, &IndexMode::Exact).unwrap();
PartitionMeta { n_layers: 1, mode: IndexMode::Exact }.save(dir.path()).unwrap();
let map = LayeredMap::<()>::open(dir.path()).unwrap();
assert_eq!(map.detect_layer_evidence(0).unwrap(), EvidenceKind::Exact);
let approx = IndexMode::Approx { b: 8, z: 1 };
let dir2 = tempdir().unwrap();
let layer0b = layer_dir(dir2.path(), 0);
write_unitigs_at(&layer0b, &[b"AAAACGT"]);
TypedLayer::<()>::build(&layer0b, DEFAULT_BLOCK_BITS, &approx).unwrap();
PartitionMeta { n_layers: 1, mode: approx }.save(dir2.path()).unwrap();
let map2 = LayeredMap::<()>::open(dir2.path()).unwrap();
assert_eq!(map2.detect_layer_evidence(0).unwrap(), EvidenceKind::Approx);
}
src/obikindex/src/layer/typed_layer.rs
+41 -84
View File
@@ -1,19 +1,18 @@
use std::collections::HashMap;
use std::fs;
use std::path::{Path, PathBuf};
use crate::layer::utils::layer_dir;
use obicompactvec::{
BinaryMatrix, PersistentBitMatrix, PersistentBitMatrixBuilder, PersistentCompactIntMatrix,
PersistentCompactIntMatrixBuilder, PersistentSparseBitMatrix,
};
use obikseq::CanonicalKmer;
use obiskio::{UnitigFileReader, UnitigFileWriter};
use std::collections::HashMap;
use std::fs;
use std::path::Path;
use crate::layer::error::{OLMError, OLMResult};
use crate::index::error::{OKIError, OKIResult};
use crate::layer::meta::IndexMode;
use crate::layer::mphf_layer::MphfLayer;
pub(crate) use crate::layer::mphf_layer::UNITIGS_FILE;
pub(crate) const COUNTS_DIR: &str = "counts";
pub(crate) const PRESENCE_DIR: &str = "presence";
@@ -21,48 +20,10 @@ pub(crate) const PRESENCE_DIR: &str = "presence";
pub trait LayerData: Sized {
type Item;
fn open(layer_dir: &Path) -> OLMResult<Self>;
fn open(layer_dir: &Path) -> OKIResult<Self>;
fn read(&self, slot: usize) -> Self::Item;
}
/// TypedLayer directory path for layer `i` within a partition's index root — the
/// single source of truth for the on-disk `layer_N` naming convention
/// (mirrors what `LayeredMap::open`/`push_layer` already use internally).
///
/// `obikindex::layer` operates within a single partition's index root; it has
/// no notion of "partition" at all. Turning a partition number into that
/// root is `obikindex::KmerIndex::index_dir`'s job, one layer up — callers
/// here only ever name a layer *number*, never build the path themselves.
pub fn layer_dir(root: &Path, i: usize) -> PathBuf {
root.join(format!("layer_{i}"))
}
/// Superkmer-file extension used by every pre-layer-construction artifact
/// below (`raw`/`dereplicated`) — an implementation detail of the SK file
/// format, never meant to leak as a literal string past this module.
const SK_EXT: &str = "skmer.zst";
/// Path of a layer's raw, not-yet-dereplicated superkmer file — written by
/// whichever algorithm routes superkmers into this layer (today:
/// `obikindexer::algorithms::partitionner::PartitionRouter`), read by whichever
/// algorithm dereplicates it (today:
/// `obikindexer::algorithms::dereplicator::Dereplicator`). Naming this once
/// here, rather than in either algorithm submodule, is what lets the two
/// agree on the filename without depending on each other directly — see
/// `DevDocMD/implementation/partition_layer_cache.md`.
pub fn raw_superkmers_path(layer_dir: &Path) -> PathBuf {
layer_dir.join(format!("raw.{SK_EXT}"))
}
/// Path of a layer's dereplicated superkmer file — written by
/// `obikindexer::algorithms::dereplicator::Dereplicator`, read by whichever
/// algorithm counts kmer abundances from it (today:
/// `obikindexer::algorithms::partitionner::PartitionRouter::count_kmer`) and,
/// later, by `obikindex::build_index_layer` to build the real layer.
pub fn dereplicated_superkmers_path(layer_dir: &Path) -> PathBuf {
layer_dir.join(format!("dereplicated.{SK_EXT}"))
}
/// Opens layer `i`'s data only, skipping the MPHF — for callers that only
/// need matrix-level operations (distance traits, column weights, group
/// filters, sub-matrix extraction) and never look up a kmer for this layer.
@@ -74,13 +35,13 @@ pub fn dereplicated_superkmers_path(layer_dir: &Path) -> PathBuf {
/// *number* within a partition it already knows the root of, the same
/// vocabulary as [`layer_dir`] and `LayeredMap` — never the `layer_N`
/// naming convention itself, which stays private to this crate.
pub fn open_data<D: LayerData>(root: &Path, i: usize) -> OLMResult<D> {
pub(crate) fn open_data<D: LayerData>(root: &Path, i: usize) -> OKIResult<D> {
D::open(&layer_dir(root, i))
}
impl LayerData for () {
type Item = ();
fn open(_layer_dir: &Path) -> OLMResult<Self> {
fn open(_layer_dir: &Path) -> OKIResult<Self> {
Ok(())
}
fn read(&self, _slot: usize) {}
@@ -88,8 +49,8 @@ impl LayerData for () {
impl LayerData for PersistentCompactIntMatrix {
type Item = Box<[u32]>;
fn open(layer_dir: &Path) -> OLMResult<Self> {
PersistentCompactIntMatrix::open(layer_dir).map_err(OLMError::Io)
fn open(layer_dir: &Path) -> OKIResult<Self> {
PersistentCompactIntMatrix::open(layer_dir).map_err(OKIError::Io)
}
fn read(&self, slot: usize) -> Box<[u32]> {
self.row(slot)
@@ -98,8 +59,8 @@ impl LayerData for PersistentCompactIntMatrix {
impl LayerData for PersistentBitMatrix {
type Item = Box<[bool]>;
fn open(layer_dir: &Path) -> OLMResult<Self> {
PersistentBitMatrix::open(layer_dir).map_err(OLMError::Io)
fn open(layer_dir: &Path) -> OKIResult<Self> {
PersistentBitMatrix::open(layer_dir).map_err(OKIError::Io)
}
fn read(&self, slot: usize) -> Box<[bool]> {
self.row(slot)
@@ -108,8 +69,8 @@ impl LayerData for PersistentBitMatrix {
impl LayerData for PersistentSparseBitMatrix {
type Item = Box<[bool]>;
fn open(layer_dir: &Path) -> OLMResult<Self> {
PersistentSparseBitMatrix::open(&layer_dir.join(PRESENCE_DIR)).map_err(OLMError::Io)
fn open(layer_dir: &Path) -> OKIResult<Self> {
PersistentSparseBitMatrix::open(&layer_dir.join(PRESENCE_DIR)).map_err(OKIError::Io)
}
fn read(&self, slot: usize) -> Box<[bool]> {
self.row(slot)
@@ -201,11 +162,7 @@ pub struct Hit<T = ()> {
// ── Common read path ──────────────────────────────────────────────────────────
impl<D: LayerData> TypedLayer<D> {
// `mode` no longer forwarded to `MphfLayer::open` (auto-detected from
// disk now, see mphf_layer.rs) — kept as a parameter for the moment so
// this signature (and `KmerLayer::open`'s call to it) doesn't have to
// change in the same step; dropped when content_layer.rs is done.
pub fn open(path: &Path, _mode: &IndexMode) -> OLMResult<Self> {
pub fn open(path: &Path) -> OKIResult<Self> {
let mphf = MphfLayer::open(path)?;
let data = D::open(path)?;
Ok(Self { mphf, data })
@@ -266,21 +223,21 @@ impl<D: LayerData> TypedLayer<D> {
self.mphf.enumerate_kmers_batch(n)
}
pub fn unitig_writer(out_dir: &Path) -> OLMResult<UnitigFileWriter> {
pub fn unitig_writer(out_dir: &Path) -> OKIResult<UnitigFileWriter> {
MphfLayer::unitig_writer(out_dir)
}
/// Build `unitigs.bin.idx` and `evidence.bin` from `unitigs.bin` and
/// `mphf.bin` already present in `layer_dir`.
/// `block_bits` controls the `.idx` block size (2^block_bits chunks/block).
pub fn build_exact_evidence(layer_dir: &Path, block_bits: u8) -> OLMResult<usize> {
pub fn build_exact_evidence(layer_dir: &Path, block_bits: u8) -> OKIResult<usize> {
MphfLayer::build_exact_evidence(layer_dir, block_bits)
}
/// Build `fingerprint.bin` from `unitigs.bin` and `mphf.bin` already
/// present in `layer_dir`. `b` — fingerprint bits (1..=64); `z` — Findere
/// consecutive k-mer parameter (≥1).
pub fn build_approx_evidence(layer_dir: &Path, b: u8, z: u8) -> OLMResult<usize> {
pub fn build_approx_evidence(layer_dir: &Path, b: u8, z: u8) -> OKIResult<usize> {
MphfLayer::build_approx_evidence(layer_dir, b, z)
}
@@ -311,21 +268,21 @@ impl<D: LayerData + HasStorageKind> TypedLayer<D> {
// ── Mode 1 — set membership ───────────────────────────────────────────────────
impl TypedLayer<()> {
pub fn build(out_dir: &Path, block_bits: u8, mode: &IndexMode) -> OLMResult<usize> {
pub fn build(out_dir: &Path, block_bits: u8, mode: &IndexMode) -> OKIResult<usize> {
MphfLayer::build(out_dir, block_bits, mode, &mut |_, _| Ok(()))
}
/// Create a presence matrix for a set-membership layer (first merge).
pub fn init_presence_matrix(layer_dir: &Path, n_kmers: usize) -> OLMResult<()> {
pub fn init_presence_matrix(layer_dir: &Path, n_kmers: usize) -> OKIResult<()> {
let presence_dir = layer_dir.join(PRESENCE_DIR);
fs::create_dir_all(&presence_dir).map_err(OLMError::Io)?;
fs::create_dir_all(&presence_dir).map_err(OKIError::Io)?;
let mut mb =
PersistentBitMatrixBuilder::new(n_kmers, &presence_dir).map_err(OLMError::Io)?;
PersistentBitMatrixBuilder::new(n_kmers, &presence_dir).map_err(OKIError::Io)?;
mb.add_col_ones()
.map_err(OLMError::Io)?
.map_err(OKIError::Io)?
.close()
.map_err(OLMError::Io)?;
mb.close().map_err(OLMError::Io)
.map_err(OKIError::Io)?;
mb.close().map_err(OKIError::Io)
}
}
@@ -337,18 +294,18 @@ impl TypedLayer<PersistentCompactIntMatrix> {
block_bits: u8,
mode: &IndexMode,
count_of: impl Fn(CanonicalKmer) -> u32,
) -> OLMResult<usize> {
) -> OKIResult<usize> {
let n = UnitigFileReader::open_sequential(&out_dir.join(UNITIGS_FILE))?.n_kmers();
let counts_dir = out_dir.join(COUNTS_DIR);
let mut mb =
PersistentCompactIntMatrixBuilder::new(n, &counts_dir).map_err(OLMError::Io)?;
let mut col = mb.add_col().map_err(OLMError::Io)?;
PersistentCompactIntMatrixBuilder::new(n, &counts_dir).map_err(OKIError::Io)?;
let mut col = mb.add_col().map_err(OKIError::Io)?;
let n_built = MphfLayer::build(out_dir, block_bits, mode, &mut |slot, kmer| {
col.set(slot, count_of(kmer));
Ok(())
})?;
col.close().map_err(OLMError::Io)?;
mb.close().map_err(OLMError::Io)?;
col.close().map_err(OKIError::Io)?;
mb.close().map_err(OKIError::Io)?;
Ok(n_built)
}
@@ -357,7 +314,7 @@ impl TypedLayer<PersistentCompactIntMatrix> {
block_bits: u8,
mode: &IndexMode,
counts: &HashMap<CanonicalKmer, u32>,
) -> OLMResult<usize> {
) -> OKIResult<usize> {
Self::build(out_dir, block_bits, mode, |kmer| {
counts.get(&kmer).copied().unwrap_or(0)
})
@@ -370,9 +327,9 @@ impl TypedLayer<PersistentCompactIntMatrix> {
pub fn append_genome_column(
layer_dir: &Path,
value_of: impl Fn(usize) -> u32,
) -> OLMResult<()> {
) -> OKIResult<()> {
PersistentCompactIntMatrix::append_column(&layer_dir.join(COUNTS_DIR), value_of)
.map_err(OLMError::Io)
.map_err(OKIError::Io)
}
/// Number of genome columns in this layer's count matrix.
@@ -448,9 +405,9 @@ impl TypedLayer<PersistentBitMatrix> {
pub fn append_genome_column(
layer_dir: &Path,
value_of: impl Fn(usize) -> bool,
) -> OLMResult<()> {
) -> OKIResult<()> {
PersistentBitMatrix::append_column(&layer_dir.join(PRESENCE_DIR), value_of)
.map_err(OLMError::Io)
.map_err(OKIError::Io)
}
pub fn build_presence(
@@ -459,13 +416,13 @@ impl TypedLayer<PersistentBitMatrix> {
mode: &IndexMode,
n_genomes: usize,
present_in: impl Fn(CanonicalKmer, usize) -> bool,
) -> OLMResult<usize> {
) -> OKIResult<usize> {
let n = UnitigFileReader::open_sequential(&out_dir.join(UNITIGS_FILE))?.n_kmers();
let presence_dir = out_dir.join(PRESENCE_DIR);
let mut mb = PersistentBitMatrixBuilder::new(n, &presence_dir).map_err(OLMError::Io)?;
let mut mb = PersistentBitMatrixBuilder::new(n, &presence_dir).map_err(OKIError::Io)?;
let mut cols: Vec<_> = (0..n_genomes)
.map(|_| mb.add_col().map_err(OLMError::Io))
.collect::<OLMResult<_>>()?;
.map(|_| mb.add_col().map_err(OKIError::Io))
.collect::<OKIResult<_>>()?;
let n_built = MphfLayer::build(out_dir, block_bits, mode, &mut |slot, kmer| {
for (g, col) in cols.iter_mut().enumerate() {
col.set(slot, present_in(kmer, g));
@@ -473,9 +430,9 @@ impl TypedLayer<PersistentBitMatrix> {
Ok(())
})?;
for col in cols {
col.close().map_err(OLMError::Io)?;
col.close().map_err(OKIError::Io)?;
}
mb.close().map_err(OLMError::Io)?;
mb.close().map_err(OKIError::Io)?;
Ok(n_built)
}
}
src/obikindex/src/layer/utils.rs
+14
View File
@@ -0,0 +1,14 @@
pub(crate) const LAYERNAME_SUFFIX: &str = "layer_";
use std::path::{Path, PathBuf};
/// TypedLayer directory path for layer `i` within a partition's index root — the
/// single source of truth for the on-disk `layer_N` naming convention
/// (mirrors what `LayeredMap::open`/`push_layer` already use internally).
///
/// `obikindex::layer` operates within a single partition's index root; it has
/// no notion of "partition" at all. Turning a partition number into that
/// root is `obikindex::KmerIndex::index_dir`'s job, one layer up — callers
/// here only ever name a layer *number*, never build the path themselves.
pub(crate) fn layer_dir(root: &Path, i: usize) -> PathBuf {
root.join(format!("{LAYERNAME_SUFFIX}{i}"))
}
src/obikindex/src/lib.rs
+6 -6
View File
@@ -13,10 +13,10 @@ pub mod layer;
pub mod partition;
pub use index::{
validate_label, AggOp, DistanceMetric, DistanceOutput, GenomeInfo, GroupFilterParams,
GroupQuorumFilter, IndexBitsPerKmer, IndexBuilder, IndexConfig, IndexMeta, IndexState,
KmerDesc, KmerFilter, KmerIndex, MergeMode, MetaPred, OKIError, OKIResult, OutputCol,
PartitionRunner, QueryHit, QueryStats, META_FILENAME,
materialize_layer, olm_to_sk, write_graph_as_unitigs, passes_all,
GenomeInfo,
IndexBuilder, IndexConfig, IndexMeta, IndexState,
KmerIndex, OKIError, OKIResult,
META_FILENAME,
ColBuilder,
};
pub use index::{filter, meta};
pub use index::meta;
src/obikindex/src/partition/mod.rs
+54 -33
View File
@@ -19,10 +19,13 @@
//! entirely) both still reinvent a fragment of this. Migrating them is a
//! separate, deferred step.
use crate::OKIError;
use std::path::{Path, PathBuf};
use crate::layer::{IndexMode, KmerLayer, OLMResult, layer_dir};
use obikseq::CanonicalKmer;
use crate::{
KmerIndex, OKIResult,
layer::{KmerLayer, utils::layer_dir},
};
/// Partition subdirectory name, under an index's root — the single source
/// of truth for the on-disk `partitions/part_NNNNN` naming convention.
@@ -34,59 +37,77 @@ pub const PARTITIONS_SUBDIR: &str = "partitions";
/// Path of partition `i`'s directory under `root` — `<root>/partitions/part_NNNNN`,
/// zero-padded to 5 digits.
pub fn partition_dir(root: &Path, i: usize) -> PathBuf {
pub(crate) fn partition_dir(root: &Path, i: usize) -> PathBuf {
root.join(PARTITIONS_SUBDIR).join(format!("part_{i:05}"))
}
/// Path of partition `i`'s layered-index directory — `<partition_dir>/index`,
/// the root every [`crate::layer::layer_dir`] call for this partition is
/// relative to.
pub fn index_dir(root: &Path, i: usize) -> PathBuf {
pub(crate) fn index_dir(root: &Path, i: usize) -> PathBuf {
partition_dir(root, i).join("index")
}
/// One partition's open layers, in layer order (layer 0 first).
/// One partition — identified by its own `index_dir`, nothing eagerly
/// opened. `index_dir` is the minimal field `KmerLayer::at` needs to
/// compute a layer's path; added here only because `layer()` below can't
/// be written without it now that there's no `layers: Vec<KmerLayer>` to
/// index into.
pub struct KmerPartition {
layers: Vec<KmerLayer>,
partition_dir: PathBuf,
id: usize,
}
impl KmerPartition {
/// Open every layer under `index_dir` (`index_dir/layer_0`,
/// `index_dir/layer_1`, ... up to `n_layers`), eagerly — not lazily on
/// first access, so the caller pays the mmap cost once, up front,
/// rather than at an unpredictable point during later lookups.
pub fn open(
index_dir: &Path,
mode: &IndexMode,
n_layers: usize,
with_counts: bool,
) -> OLMResult<Self> {
let layers = (0..n_layers)
.map(|l| KmerLayer::open(&layer_dir(index_dir, l), mode, with_counts))
.collect::<OLMResult<Vec<_>>>()?;
Ok(Self { layers })
pub fn new(index: &KmerIndex, i: usize) -> Self {
KmerPartition {
partition_dir: index.dir().join(format!("part_{i:05}")),
id: i,
}
}
pub fn id(&self) -> usize {
self.id
}
/// Number of layers under `index_dir`, found by probing `layer_0`,
/// `layer_1`, ... sequentially until one is missing — a disk scan, not
/// a metadata read (`meta.json`'s own `n_layers` is redundant with this
/// and no longer consulted here). Intrinsically slow: one `exists()`
/// syscall per layer, called fresh every time — not for a hot path.
pub fn n_layers(&self) -> usize {
self.layers.len()
let mut n = 0;
while layer_dir(&self.partition_dir, n).exists() {
n += 1;
}
n
}
pub fn layer(&self, i: usize) -> &KmerLayer {
&self.layers[i]
/// Evidence mode for this partition, detected from layer 0's evidence
/// files on disk (see [`crate::layer::IndexMode::detect`]) — never
/// persisted as metadata, since every layer within a partition shares
/// the same mode by construction.
pub fn mode(&self) -> OKIResult<crate::layer::IndexMode> {
crate::layer::IndexMode::detect(&self.partition_dir)
}
pub fn layers(&self) -> &[KmerLayer] {
&self.layers
pub fn layer(&self, i: usize) -> OKIResult<KmerLayer> {
let n = self.n_layers();
if i >= n {
return Err(OKIError::Layer(obiskio::SKError::InvalidData {
context: "New KmerLayer",
detail: format!("Layer {i} requested, Partition has only {n} layer"),
}));
}
Ok(KmerLayer::new(self, i))
}
/// Existence lookup of `kmer` across this partition's layers: tries
/// each in turn, stopping at the first hit and reporting which layer it
/// was — `find_slot`, not a data read, for a plain existence check.
pub fn find(&self, kmer: CanonicalKmer) -> Option<usize> {
self.layers
.iter()
.enumerate()
.find_map(|(li, layer)| layer.find_slot(kmer).map(|_| li))
pub fn layer_dir(&self, i: usize) -> OKIResult<PathBuf> {
Ok(self.layer(i)?.dir().to_path_buf())
}
pub fn dir(&self) -> &Path {
&self.partition_dir
}
}