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

refactor: extract index construction state tracking into extension trait

Browse Source 
Moves pipeline state bookkeeping, including sentinel file marking and spectrum persistence, into the algorithms' run and close methods. Introduces a crate-private extension trait to satisfy Rust's orphan rule while implementing construction-only operations on KmerIndex. Updates helper function visibility for cross-crate access and removes the now-empty index_layer module.
This commit is contained in:
Eric Coissac
2026-08-21 10:38:51 +02:00
parent 31bb324752
commit da3aa5a2cb
16 changed files with 527 additions and 254 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/obikindex/src/index/common.rs
+1 -1
View File
@@ -7,7 +7,7 @@ use obiskio::{SKError, SKResult};
// ── olm_to_sk ────────────────────────────────────────────────────────────────
pub(crate) fn olm_to_sk(e: OLMError, context: &'static str) -> SKError {
pub fn olm_to_sk(e: OLMError, context: &'static str) -> SKError {
match e {
OLMError::Io(e) => SKError::Io(e),
other => SKError::InvalidData {
src/obikindex/src/index/graph_pipeline.rs
+2 -2
View File
@@ -121,7 +121,7 @@ where
/// Phase 2 (write unitigs only): compute degrees, write unitigs to `layer_dir`, drop graph.
///
/// Returns n_kmers. Does NOT build the MPHF — caller does it.
pub(crate) fn write_graph_as_unitigs(g: GraphDeBruijn, layer_dir: &Path) -> SKResult<usize> {
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)?;
@@ -137,7 +137,7 @@ pub(crate) fn write_graph_as_unitigs(g: GraphDeBruijn, layer_dir: &Path) -> SKRe
/// Phase 2 (full): write_graph_as_unitigs + `TypedLayer::<()>::build`.
///
/// Returns n_kmers.
pub(crate) fn materialize_layer(
pub fn materialize_layer(
g: GraphDeBruijn,
layer_dir: &Path,
block_bits: u8,
src/obikindex/src/index/index_layer.rs
-132
View File
@@ -1,132 +0,0 @@
use std::fs;
use std::io;
use cacheline_ef::{CachelineEf, CachelineEfVec};
use epserde::prelude::*;
use obicompactvec::{PersistentCompactIntMatrix, PersistentCompactIntVec};
use obidebruinj::GraphDeBruijn;
use crate::layer::meta::PartitionMeta;
use crate::layer::{IndexMode, TypedLayer};
use obiskio::{SKError, SKFileMeta, SKFileReader};
use ptr_hash::{PtrHash, bucket_fn::CubicEps, hash::Xx64};
use crate::index::common::olm_to_sk;
use crate::index::graph_pipeline::{materialize_layer, write_graph_as_unitigs};
use crate::index::kmer_index::KmerIndex;
type Mphf = PtrHash<u64, CubicEps, CachelineEfVec<Vec<CachelineEf>>, Xx64, Vec<u8>>;
fn remove_if_exists(path: &std::path::Path) {
if let Err(e) = fs::remove_file(path) {
if e.kind() != io::ErrorKind::NotFound {
eprintln!("warning: could not remove {}: {e}", path.display());
}
}
}
impl KmerIndex {
/// Build the layered MPHF index for partition `i`.
///
/// Returns the number of canonical k-mers indexed, or 0 if the partition
/// has no data or its layer was already built (resume-safe).
///
/// Abundance filtering is applied when `min_ab > 1` or `max_ab.is_some()`,
/// using `mphf1.bin` + `counts1.bin` if they exist.
/// Count payload is stored iff `with_counts` is true.
pub fn build_index_layer(
&self,
i: usize,
min_ab: u32,
max_ab: Option<u32>,
with_counts: bool,
mode: &IndexMode,
block_bits: u8,
) -> Result<usize, SKError> {
let layer0_dir = self.layer_dir(i, 0);
let dedup_path = crate::layer::dereplicated_superkmers_path(&layer0_dir);
if !dedup_path.exists() {
return Ok(0);
}
if layer0_dir.join("mphf.bin").exists() {
return Ok(0);
}
let filter_active = min_ab > 1 || max_ab.is_some();
let need_counts = filter_active || with_counts;
let mphf1_opt: Option<Mphf> = if need_counts {
let p = layer0_dir.join("mphf1.bin");
p.exists().then(|| Mphf::load_full(&p).ok()).flatten()
} else {
None
};
let counts1_opt: Option<PersistentCompactIntVec> = if need_counts {
let p = layer0_dir.join("counts1.bin");
p.exists()
.then(|| PersistentCompactIntVec::open(&p).ok())
.flatten()
} else {
None
};
let mut g = GraphDeBruijn::new();
let mut reader = SKFileReader::open(&dedup_path)?;
for sk in reader.iter() {
for kmer in sk.iter_canonical_kmers() {
let accept = if filter_active {
match (&mphf1_opt, &counts1_opt) {
(Some(mphf), Some(counts)) => {
let ab = counts.get(mphf.index(&kmer.raw()));
ab >= min_ab && max_ab.map_or(true, |max| ab <= max)
}
_ => true,
}
} else {
true
};
if accept {
g.push(kmer);
}
}
}
let n_kmers =
if with_counts {
let n = write_graph_as_unitigs(g, &layer0_dir)?;
TypedLayer::<PersistentCompactIntMatrix>::build(&layer0_dir, block_bits, mode, |kmer| {
match (&mphf1_opt, &counts1_opt) {
(Some(mphf), Some(counts)) => counts.get(mphf.index(&kmer.raw())),
_ => 1,
}
})
.map_err(|e| olm_to_sk(e, "layer build"))?;
n
} else {
materialize_layer(g, &layer0_dir, block_bits, mode)?
};
let index_dir = layer0_dir.parent().expect("layer_dir has a parent");
PartitionMeta {
n_layers: 1,
mode: mode.clone(),
}
.save(index_dir)
.map_err(|e| olm_to_sk(e, "layer build"))?;
Ok(n_kmers)
}
/// Remove intermediate build artifacts for partition `i`.
///
/// Deletes `dereplicated.skmer.zst` (+ sidecar), `mphf1.bin`, `counts1.bin`.
pub fn remove_build_artifacts(&self, i: usize) {
let layer0_dir = self.layer_dir(i, 0);
let dedup = crate::layer::dereplicated_superkmers_path(&layer0_dir);
remove_if_exists(&SKFileMeta::sidecar_path(&dedup));
remove_if_exists(&dedup);
remove_if_exists(&layer0_dir.join("mphf1.bin"));
remove_if_exists(&layer0_dir.join("counts1.bin"));
}
}
src/obikindex/src/index/kmer_index.rs
+1 -77
View File
@@ -1,4 +1,3 @@
use std::collections::BTreeMap;
use std::fs;
use std::path::{Path, PathBuf};
@@ -11,7 +10,7 @@ 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};
use crate::index::state::{IndexState, SENTINEL_COUNTED, SENTINEL_INDEXED, SENTINEL_SCATTERED};
use crate::index::state::{IndexState, SENTINEL_INDEXED};
pub struct KmerIndex {
pub(crate) root_path: PathBuf,
@@ -209,62 +208,6 @@ impl KmerIndex {
Ok(self.n_layers(0)?)
}
/// Mark scatter as complete and write `scatter.done`.
///
/// If no genome label was set at creation time, one is derived from
/// the index root directory name (stripped of all extensions).
pub fn mark_scattered(&mut self) -> OKIResult<()> {
if self.meta.genomes.is_empty() {
let label = label_from_path(&self.root_path);
self.meta.genomes.push(GenomeInfo::new(label));
self.meta.write(&self.root_path)?;
}
touch(&self.root_path.join(SENTINEL_SCATTERED))?;
Ok(())
}
/// Mark dereplicate+count as complete and write `count.done`.
pub fn mark_counted(&self) -> OKIResult<()> {
touch(&self.root_path.join(SENTINEL_COUNTED))?;
Ok(())
}
/// Mark layer construction as complete and write `index.done`.
pub fn mark_indexed(&self) -> OKIResult<()> {
touch(&self.root_path.join(SENTINEL_INDEXED))?;
Ok(())
}
/// Write `spectrums/{label}.json` from an already-computed kmer
/// spectrum (`f0`/`f1`/abundance histogram). Takes plain values rather
/// than `obikindexer::algorithms::partitionner::KmerSpectrum` — `KmerIndex`
/// is the data model, `PartitionRouter` the algorithm that depends on
/// it (see `DevDocMD/implementation/partition_layer_cache.md`), not
/// the other way around, and this is the only field of that type it
/// actually uses.
pub fn write_spectrum(&self, f0: u64, f1: u64, counts: &BTreeMap<u32, u64>) -> OKIResult<()> {
let label = self
.meta
.genomes
.first()
.map(|g| g.label.as_str())
.unwrap_or("unknown");
let spectrums_dir = self.root_path.join("spectrums");
fs::create_dir_all(&spectrums_dir)?;
let path = spectrums_dir.join(format!("{label}.json"));
let spectrum_map: BTreeMap<String, u64> = counts
.iter()
.map(|(&c, &f)| (format!("{c:010}"), f))
.collect();
let f = fs::File::create(&path)?;
serde_json::to_writer_pretty(
f,
&serde_json::json!({ "f0": f0, "f1": f1, "spectrum": spectrum_map }),
)
.map_err(OKIError::Json)?;
Ok(())
}
/// 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")
@@ -370,23 +313,4 @@ impl KmerIndex {
}
}
fn label_from_path(path: &Path) -> String {
let name = path
.file_name()
.unwrap_or(path.as_os_str())
.to_string_lossy()
.into_owned();
let mut s = name;
while let Some(pos) = s.rfind('.') {
s.truncate(pos);
}
if s.is_empty() {
"unknown".to_string()
} else {
s
}
}
fn touch(path: &Path) -> Result<(), std::io::Error> {
fs::File::create(path).map(|_| ())
}
src/obikindex/src/index/mod.rs
+2 -1
View File
@@ -9,7 +9,6 @@ mod dump_layer;
pub mod filter;
mod graph_pipeline;
mod kmer_index;
mod index_layer;
mod matrix_store;
mod merge;
mod merge_layer;
@@ -23,6 +22,8 @@ mod select_layer;
mod stats;
pub use error::{OKIError, OKIResult};
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 kmer_index::KmerIndex;
src/obikindex/src/lib.rs
+1 -1
View File
@@ -17,6 +17,6 @@ pub use index::{
GroupQuorumFilter, IndexBitsPerKmer, IndexConfig, IndexMeta, IndexState, KmerDesc, KmerFilter,
KmerIndex, MergeMode, MetaPred, OKIError, OKIResult, OutputCol, PartitionRunner, QueryHit,
QueryStats, META_FILENAME, SENTINEL_COUNTED, SENTINEL_INDEXED, SENTINEL_SCATTERED,
passes_all,
materialize_layer, olm_to_sk, write_graph_as_unitigs, passes_all,
};
pub use index::{filter, meta};