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

Refactor indexing pipeline and partition path resolution

Browse Source 
Invert dependencies by moving partition path primitives to a dedicated module and updating the index crate accordingly. Reshape the partition router to accept a mutable index reference, enabling chainable configuration and resolving lifetime issues with explicit drops. Shift orchestration logic from the index crate to the CLI, replacing monolithic scatter calls with discrete dereplication and counting steps. Introduce a generic progress callback API to decouple rate calculation from UI rendering, and correct file I/O paths to route layer-0 artifacts under the partition index directory.
This commit is contained in:
Eric Coissac
2026-08-20 22:15:13 +02:00
parent 164e879585
commit 616cb76af3
19 changed files with 535 additions and 405 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/obikindex/Cargo.toml
+1 -1
View File
@@ -5,7 +5,7 @@ edition = "2024"
[dependencies]
obikseq = { path = "../obikseq" }
obikpartitionner = { path = "../obikpartitionner" }
obikpartition = { path = "../obikpartition" }
obitaxonomy = { path = "../obitaxonomy" }
obiskio = { path = "../obiskio" }
obisys = { path = "../obisys" }
src/obikindex/src/index.rs
+19 -41
View File
@@ -2,7 +2,6 @@ use std::collections::BTreeMap;
use std::fs;
use std::path::{Path, PathBuf};
use obikpartitionner::{KmerSpectrum, PartitionRouter};
use obilayeredmap::meta::PartitionMeta;
use obisys::{Reporter, Stage, progress_bar};
use rayon::prelude::*;
@@ -29,10 +28,8 @@ impl KmerIndex {
path: P,
config: IndexConfig,
genome_info: Option<GenomeInfo>,
force: bool,
) -> OKIResult<Self> {
let root_path = path.as_ref().to_owned();
PartitionRouter::create(&root_path, config.n_bits, force)?;
set_k(config.kmer_size);
set_m(config.minimizer_size);
let mut meta = IndexMeta::new(config);
@@ -86,7 +83,6 @@ impl KmerIndex {
let output = output.as_ref();
fs::create_dir_all(output).map_err(OKIError::Io)?;
meta.write(output).map_err(OKIError::Io)?;
PartitionRouter::create(output, meta.config.n_bits, false)?;
Ok(KmerIndex { root_path: output.to_owned(), meta: meta.clone() })
}
@@ -160,16 +156,19 @@ impl KmerIndex {
}
/// Path of partition `i`'s raw directory (`partitions/part_{i:05}`) —
/// the on-disk naming convention `obikpartitionner::PartitionRouter`
/// also writes to (raw/dereplicated superkmer files, `mphf1.bin`,
/// `counts1.bin`), the single point of agreement between the two.
/// delegates to `obikpartition`, the Partition tier's own naming
/// primitive (mirrors `layer_dir` delegating to `obilayeredmap`).
/// `obikpartitionner::PartitionRouter` reaches this same directory
/// only indirectly, through this method (it depends on `KmerIndex`,
/// not the other way around — see
/// `DevDocMD/implementation/partition_layer_cache.md`).
pub fn partition_dir(&self, i: usize) -> PathBuf {
obikpartitionner::partition_dir(&self.root_path, i)
obikpartition::partition_dir(&self.root_path, i)
}
/// Path of partition `i`'s layered-index directory (`<partition>/index`).
pub fn index_dir(&self, i: usize) -> PathBuf {
self.partition_dir(i).join("index")
obikpartition::index_dir(&self.root_path, i)
}
/// Path of layer `l` within partition `i`'s layered index.
@@ -223,39 +222,19 @@ impl KmerIndex {
Ok(())
}
/// Open a fresh [`PartitionRouter`] onto this index's partition layout
/// — the write-side handle for `scatter`, or for `dereplicate_and_count`
/// below. Transient: no state is kept in `KmerIndex` itself between
/// calls, only on disk.
pub fn partition_router(&self) -> OKIResult<PartitionRouter> {
Ok(PartitionRouter::open(&self.root_path, self.meta.config.n_bits)?)
}
/// Dereplicate all partitions then compute kmer counts.
///
/// Writes `spectrums/{label}.json` and touches `count.done` upon completion.
/// Per-partition spectrum files are removed unless `keep_intermediate` is true.
pub fn dereplicate_and_count(
&self,
keep_intermediate: bool,
rep: &mut Reporter,
) -> OKIResult<()> {
let router = self.partition_router()?;
let t = Stage::start("dereplicate");
router.dereplicate()?;
rep.push(t.stop());
let t = Stage::start("count_kmer");
let spectrum = router.count_kmer(keep_intermediate)?;
rep.push(t.stop());
self.write_spectrum(&spectrum)?;
/// Mark dereplicate+count as complete and write `count.done`.
pub fn mark_counted(&self) -> OKIResult<()> {
touch(&self.root_path.join(SENTINEL_COUNTED))?;
Ok(())
}
fn write_spectrum(&self, sp: &KmerSpectrum) -> OKIResult<()> {
/// Write `spectrums/{label}.json` from an already-computed kmer
/// spectrum (`f0`/`f1`/abundance histogram). Takes plain values rather
/// than `obikpartitionner::KmerSpectrum` — `KmerIndex` cannot depend on
/// `obikpartitionner` (that dependency runs the other way, see
/// `DevDocMD/implementation/partition_layer_cache.md`), and doesn't
/// need to: 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
@@ -265,15 +244,14 @@ impl KmerIndex {
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> = sp
.counts
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": sp.f0, "f1": sp.f1, "spectrum": spectrum_map }),
&serde_json::json!({ "f0": f0, "f1": f1, "spectrum": spectrum_map }),
)
.map_err(OKIError::Json)?;
Ok(())
src/obikindex/src/index_layer.rs
+13 -14
View File
@@ -42,14 +42,13 @@ impl KmerIndex {
mode: &IndexMode,
block_bits: u8,
) -> Result<usize, SKError> {
let partition_dir = self.partition_dir(i);
let dedup_path = partition_dir.join("dereplicated.skmer.zst");
let layer0_dir = self.layer_dir(i, 0);
let dedup_path = layer0_dir.join("dereplicated.skmer.zst");
if !dedup_path.exists() {
return Ok(0);
}
let layer_dir = self.layer_dir(i, 0);
if layer_dir.join("mphf.bin").exists() {
if layer0_dir.join("mphf.bin").exists() {
return Ok(0);
}
@@ -57,14 +56,14 @@ impl KmerIndex {
let need_counts = filter_active || with_counts;
let mphf1_opt: Option<Mphf> = if need_counts {
let p = partition_dir.join("mphf1.bin");
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 = partition_dir.join("counts1.bin");
let p = layer0_dir.join("counts1.bin");
p.exists()
.then(|| PersistentCompactIntVec::open(&p).ok())
.flatten()
@@ -95,8 +94,8 @@ impl KmerIndex {
let n_kmers =
if with_counts {
let n = write_graph_as_unitigs(g, &layer_dir)?;
TypedLayer::<PersistentCompactIntMatrix>::build(&layer_dir, block_bits, mode, |kmer| {
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,
@@ -105,10 +104,10 @@ impl KmerIndex {
.map_err(|e| olm_to_sk(e, "layer build"))?;
n
} else {
materialize_layer(g, &layer_dir, block_bits, mode)?
materialize_layer(g, &layer0_dir, block_bits, mode)?
};
let index_dir = layer_dir.parent().expect("layer_dir has a parent");
let index_dir = layer0_dir.parent().expect("layer_dir has a parent");
PartitionMeta {
n_layers: 1,
mode: mode.clone(),
@@ -123,11 +122,11 @@ impl KmerIndex {
///
/// Deletes `dereplicated.skmer.zst` (+ sidecar), `mphf1.bin`, `counts1.bin`.
pub fn remove_build_artifacts(&self, i: usize) {
let partition_dir = self.partition_dir(i);
let dedup = partition_dir.join("dereplicated.skmer.zst");
let layer0_dir = self.layer_dir(i, 0);
let dedup = layer0_dir.join("dereplicated.skmer.zst");
remove_if_exists(&SKFileMeta::sidecar_path(&dedup));
remove_if_exists(&dedup);
remove_if_exists(&partition_dir.join("mphf1.bin"));
remove_if_exists(&partition_dir.join("counts1.bin"));
remove_if_exists(&layer0_dir.join("mphf1.bin"));
remove_if_exists(&layer0_dir.join("counts1.bin"));
}
}
src/obikindex/src/tests/query_layer.rs
+2 -2
View File
@@ -53,7 +53,7 @@ fn query_partition_with_missing_index_dir_returns_default_stats() {
evidence: obilayeredmap::IndexMode::Exact,
block_bits: 0,
};
let index = KmerIndex::create(tmp.path().join("idx"), config, None, false).expect("create index");
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
@@ -81,7 +81,7 @@ fn query_partition_with_empty_kmers_is_a_noop() {
evidence: obilayeredmap::IndexMode::Exact,
block_bits: 0,
};
let index = KmerIndex::create(tmp.path().join("idx"), config, None, false).expect("create index");
let index = KmerIndex::create(tmp.path().join("idx"), config, None).expect("create index");
let kmers: HashMap<CanonicalKmer, Vec<KmerDesc>> = HashMap::new();
let stats = index