Refactor indexing pipeline and partition path resolution

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
@@ -5,7 +5,11 @@ exists, `Mat` is gone. (1b) done — `Layer::Empty`, the first non-ready
state, added (panics on every read method). (2a) done — the
`obikpartition` crate and `KmerPartition` itself exist (`open`/`n_layers`/
`layer`/`layers`/`find`). (2b) — migrating `PartitionCache`/`QueryLayer`
onto it — **not started**, deliberately deferred. Earlier mix-up, for
onto it — **not started**, deliberately deferred. (3) done — the
`obikindex ↔ obikpartitionner` dependency inverted: `PartitionRouter` now
takes `&mut KmerIndex` and produces `Layer::Empty` shells directly, closing
the gap `Layer::Empty` was built for in (1b) — see "(3) done" below.
Earlier mix-up, for
context: an earlier
version of this doc used the name `KmerPartition` (singular) for what was
actually the *collection* type (later renamed `KmerPartitions`, later
@@ -315,6 +319,134 @@ tuple moved in wholesale); `scan_layer_families`'s still-independent
`PartitionMeta::load` (see "Remaining instance…" below) — all explicitly
deferred to whenever wiring is tackled next.
## (3) done (2026-08-20): `obikindex ↔ obikpartitionner` dependency inverted, `PartitionRouter` now fills `Layer::Empty` shells
Resolved a question left implicit since "Major restructuring": that pass
set the direction `obikindex → obikpartitionner` (so `KmerIndex` could
delegate `partition_dir` to it) without questioning whether that was the
right direction at all. Challenged directly: `obikpartitionner` is an
*algorithm* (superkmer routing/dereplication/counting) operating on an
*index* (`KmerIndex`, the data structure) — algorithms depend on the data
types they need, not the other way around. [[feedback_no_precedent_defense]]
applied here: "that's the direction we already picked" was not treated as
a justification for keeping it.
**New direction**: `obikpartitionner → obikindex` (+ `obilayeredmap`,
`obipipeline`, `obiread` directly, for what `run`'s pipeline itself needs).
`obikindex → obikpartitionner` is gone entirely — `KmerIndex` no longer
imports `PartitionRouter`/`KmerSpectrum` in any form. Two path-naming
primitives that used to make this edge necessary moved down a tier instead
of staying put:
- `partition_dir`/`PARTITIONS_SUBDIR` moved from `obikpartitionner` into
`obikpartition` (the Partition-tier crate `KmerPartition` already lives
in), alongside a new `index_dir(root, i)` — both free functions,
mirroring `obilayeredmap::layer_dir` one tier down. `KmerIndex::
partition_dir`/`index_dir` now delegate here instead of to
`obikpartitionner`/an inline `.join("index")`.
- `KmerIndex::create`/`create_skeleton` no longer call
`PartitionRouter::create` to lay out an empty `partitions/` skeleton
upfront — turned out to be dead weight once traced: `select_layer.rs`/
`rebuild_layer.rs` already `create_dir_all` their own partition/layer
directories on demand, and `Layer::create`'s directory-creation covers
the scatter path the same way. Partitions and their layer-0 shells now
come into existence lazily, on first write, with nothing to pre-create.
`KmerIndex::create`'s now-unused `force: bool` parameter was dropped
(4 call sites updated) rather than left as a dead parameter.
**`PartitionRouter` reshaped** (`obikpartitionner/src/partition/router.rs`)
around the "création, paramétrage, run()" shape agreed on: `new(index:
&mut KmerIndex) -> Self` (no disk access), chainable setters
(`level_max`/`theta`/`workers`/`max_open`, defaults matching the CLI's old
hardcoded values), then `run(path_source, on_progress)`. `write`/
`write_batch`/`flush`/`close`/`dereplicate`/`count_kmer` stay public,
unconsumed (`&self`/`&mut self`, not `self`) — callers needing fine-grained
control (tests, `obikphylo`'s test harness) still get it, `run` is a
convenience layered on top, not the only way in.
`run` absorbs the entire body of what used to be the free function
`obikmer::steps::scatter` (now deleted, along with the `steps` module
entirely) — the `obipipeline::make_pipe!` two-stage pipeline
(file→pages→superkmers), throttling, per-file logging. What changed:
- Every `ensure_writer(partition)` call now does `Layer::create(&layer0_dir)`
(`layer0_dir = obilayeredmap::layer_dir(&index.index_dir(i), 0)`) before
opening `raw.{ext}` inside it — raw/dereplicated superkmer files and the
provisional `mphf1.bin`/`counts1.bin`/`kmer_spectrum_raw.json` now live
under `<partition>/index/layer_0/`, not flat under `<partition>/` as
before. This is `Layer::Empty` actually being used as the "builder code
holding an `Empty` layer" its own (1b) docs anticipated, not just a shell
with no consumer.
- **Caught by an end-to-end smoke test, not by `cargo test`**: this path
move broke `obikindex::index_layer::build_index_layer` and
`remove_build_artifacts`, both of which still read/deleted
`dereplicated.skmer.zst`/`mphf1.bin`/`counts1.bin` from
`self.partition_dir(i)` (the old flat location) — no test in the
workspace suite exercises the real CLI's file-reading `scatter` path
end-to-end (`obikphylo`'s test harness and `obikpartitionner`'s own
tests both call `write_batch` directly, bypassing `run`/file discovery
entirely), so the whole suite stayed green while `obikmer index` on
real FASTA silently indexed 0 kmers. Found by running the actual CLI
against a small FASTA and noticing `count.json`'s `f0` (870, correct)
didn't match "0 total kmers indexed" at the final stage. Fixed by
retargeting both functions to `self.layer_dir(i, 0)`. **Lesson,
consistent with the retracted-claim lesson above**: a green test suite
is not proof a refactor is correct when no test in it exercises the
specific path that changed — for anything touching the CLI's own
file-driven entry point, running the CLI for real is not optional
verification.
- The internal `obisys::spinner("scatter")` + hand-rolled EMA-rate display
is gone from the library entirely, replaced by an `Option<impl
FnMut(obisys::Progress)>` parameter — a new, deliberately generic
progress-reporting type (`obisys::Progress { position: u64, total:
Option<u64> }`, alongside the existing `TracedBar`/`spinner`/
`progress_bar`) added specifically so every future algo crate's `run()`
reports progress the same shape, once, rather than each inventing its
own. `total: None` here (bases processed isn't knowable without
pre-scanning every input file) — deliberately simpler than the old
in-library rate/file-count/thread-count message; the caller can
recompute a Mbp/s rate from consecutive `position` values +
wall-clock time itself, which is exactly what `cmd/index/mod.rs` now
does to reproduce the old spinner message. This is a real, intentional
restriction of the library's job: it reports raw ticks, the CLI decides
what a human sees — same "generic vs. domain-specific" split applied
again, this time to progress reporting rather than to Layer content.
Explicitly **not** the same mechanism as `Stage`/`Reporter` (per
[[feedback_stage_reporter_in_cmd_layer]]): `Stage`/`Reporter` measures a
whole call's wall time from outside it; a progress callback has to fire
*from inside* a loop mid-call, which wrapping from outside cannot
express — two different needs, not the same rule reapplied under a new
name. `Stage::start("scatter")`/`rep.push(...)` stayed in
`cmd/index/mod.rs`, wrapping the whole `run()` call, unchanged in kind.
- `dereplicate`/`count_kmer` keep their existing internal
`obisys::progress_bar(...)` calls as-is (unconverted to the callback) —
explicitly out of scope for this pass, by agreement.
**Forced, not optional, consequence of the dependency inversion**:
`KmerIndex::dereplicate_and_count`/`partition_router`/`write_spectrum(&
KmerSpectrum)` could not stay on `KmerIndex` at all once `obikindex` can no
longer name `obikpartitionner::{PartitionRouter, KmerSpectrum}` in any
position — not a design choice, a mechanical requirement of severing the
edge. Replaced by: `KmerIndex::write_spectrum(f0: u64, f1: u64, counts:
&BTreeMap<u32, u64>)` (plain values, no `KmerSpectrum` dependency) and a
new `KmerIndex::mark_counted()` (symmetric to the already-existing
`mark_scattered`), with the orchestration itself (`router.dereplicate()`
`router.count_kmer()``write_spectrum``mark_counted()`) now living in
`cmd/index/mod.rs`, not `obikindex`.
Every `PartitionRouter::new(&mut index)` call in this codebase runs into
the same NLL trap once: `PartitionRouter` has a `Drop` impl (auto-`close`
on scope exit), which extends its `&mut KmerIndex` borrow to the end of
the enclosing scope even after its last real use — `idx.mark_scattered()`
right after `router.run(...)` (or `idx.write_spectrum(...)` right after
`router.count_kmer(...)`) fails to borrow-check unless the router is
`drop()`-ed explicitly first. Hit and fixed identically at all three call
sites that needed it (`cmd/index/mod.rs` ×2, `obikphylo`'s test harness,
`obikpartitionner`'s own tests).
Full workspace suite green (`cargo check --workspace --all-targets` +
`cargo test --workspace`, exit code 0) both before and after the
`index_layer.rs` fix above — the smoke test is what actually caught the
regression the suite missed.
## The problem
Reading a layer's data (MPHF + matrix) is not free: `MphfLayer::open` mmaps
+6 -95
View File
@@ -39,15 +39,6 @@ dependencies = [
"memchr",
]
[[package]]
name = "aligned"
version = "0.4.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ee4508988c62edf04abd8d92897fca0c2995d907ce1dfeaf369dac3716a40685"
dependencies = [
"as-slice",
]
[[package]]
name = "aligned-vec"
version = "0.6.4"
@@ -137,15 +128,6 @@ version = "0.7.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50"
[[package]]
name = "as-slice"
version = "0.2.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "516b6b4f0e40d50dcda9365d53964ec74560ad4284da2e7fc97122cd83174516"
dependencies = [
"stable_deref_trait",
]
[[package]]
name = "autocfg"
version = "1.5.0"
@@ -340,12 +322,6 @@ version = "1.0.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801"
[[package]]
name = "cfg_aliases"
version = "0.2.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724"
[[package]]
name = "chacha20"
version = "0.10.1"
@@ -684,15 +660,6 @@ dependencies = [
"memchr",
]
[[package]]
name = "cvt"
version = "0.1.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d2ae9bf77fbf2d39ef573205d554d87e86c12f1994e9ea335b0651b9b278bcf1"
dependencies = [
"cfg-if",
]
[[package]]
name = "debugid"
version = "0.8.0"
@@ -913,20 +880,6 @@ version = "0.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb"
[[package]]
name = "fs_at"
version = "0.2.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "14af6c9694ea25db25baa2a1788703b9e7c6648dcaeeebeb98f7561b5384c036"
dependencies = [
"aligned",
"cfg-if",
"cvt",
"libc",
"nix 0.29.0",
"windows-sys 0.52.0",
]
[[package]]
name = "funty"
version = "2.0.0"
@@ -1422,7 +1375,7 @@ dependencies = [
"combine",
"libc",
"mach2",
"nix 0.26.4",
"nix",
"sysctl",
"thiserror 1.0.69",
"widestring",
@@ -1484,27 +1437,6 @@ dependencies = [
"pin-utils",
]
[[package]]
name = "nix"
version = "0.29.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "71e2746dc3a24dd78b3cfcb7be93368c6de9963d30f43a6a73998a9cf4b17b46"
dependencies = [
"bitflags 2.11.1",
"cfg-if",
"cfg_aliases",
"libc",
]
[[package]]
name = "normpath"
version = "1.5.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bf23ab2b905654b4cb177e30b629937b3868311d4e1cba859f899c041046e69b"
dependencies = [
"windows-sys 0.61.2",
]
[[package]]
name = "ntapi"
version = "0.4.3"
@@ -1606,7 +1538,7 @@ dependencies = [
"obicompactvec",
"obidebruinj",
"obikentropy",
"obikpartitionner",
"obikpartition",
"obikseq",
"obilayeredmap",
"obipipeline",
@@ -1675,15 +1607,17 @@ dependencies = [
"memmap2",
"niffler",
"obicompactvec",
"obikindex",
"obikrope",
"obikseq",
"obilayeredmap",
"obipipeline",
"obiread",
"obiskbuilder",
"obiskio",
"obisys",
"ptr_hash",
"rayon",
"remove_dir_all",
"serde",
"serde_json",
"sysinfo",
@@ -2016,7 +1950,7 @@ dependencies = [
"findshlibs",
"libc",
"log",
"nix 0.26.4",
"nix",
"once_cell",
"prost",
"prost-build",
@@ -2371,20 +2305,6 @@ version = "0.8.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4"
[[package]]
name = "remove_dir_all"
version = "1.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "808cc0b475acf76adf36f08ca49429b12aad9f678cb56143d5b3cb49b9a1dd08"
dependencies = [
"cfg-if",
"cvt",
"fs_at",
"libc",
"normpath",
"windows-sys 0.59.0",
]
[[package]]
name = "ring"
version = "0.17.14"
@@ -3312,15 +3232,6 @@ dependencies = [
"windows-targets 0.52.6",
]
[[package]]
name = "windows-sys"
version = "0.59.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b"
dependencies = [
"windows-targets 0.52.6",
]
[[package]]
name = "windows-sys"
version = "0.61.2"
+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" }
+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(())
+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"));
}
}
+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
+68 -13
View File
@@ -1,18 +1,19 @@
use std::path::PathBuf;
use std::time::Instant;
use clap::Args;
use obikindex::{validate_label, GenomeInfo, IndexConfig, IndexState, KmerIndex};
use obikpartitionner::PartitionRouter;
use obilayeredmap::IndexMode;
fn parse_key_value(s: &str) -> Result<(String, String), String> {
let pos = s.find('=').ok_or_else(|| format!("invalid key=value: no '=' in '{s}'"))?;
Ok((s[..pos].to_string(), s[pos + 1..].to_string()))
}
use obisys::Reporter;
use obisys::{spinner, Progress, Reporter, Stage};
use tracing::info;
use crate::cli::{CommonArgs, block_size_to_bits, partitions_to_bits};
use crate::steps::scatter;
#[derive(Args)]
pub struct IndexArgs {
@@ -225,7 +226,7 @@ pub fn run(args: IndexArgs) {
}
info
});
KmerIndex::create(&output, config, genome_info, args.force).unwrap_or_else(|e| {
KmerIndex::create(&output, config, genome_info).unwrap_or_else(|e| {
eprintln!("error creating index: {e}");
std::process::exit(1);
})
@@ -234,17 +235,50 @@ pub fn run(args: IndexArgs) {
// ── Stage 1: scatter ─────────────────────────────────────────────────────
if idx.state() < IndexState::Scattered {
let k = idx.kmer_size();
let level_max = args.common.level_max;
let theta = args.common.theta;
let n_workers = args.common.threads.max(1);
let max_open = args.common.effective_max_open();
let mut router = idx.partition_router().unwrap_or_else(|e| {
eprintln!("error opening partition router: {e}");
std::process::exit(1);
});
scatter(&mut router, args.common.seqfile_paths(), k, level_max, theta, n_workers, max_open, &mut rep);
let t = Stage::start("scatter");
let pb = spinner("scatter");
let mut ema_rate: f64 = 0.0;
let mut last_t = Instant::now();
let mut last_bases: u64 = 0;
const ALPHA: f64 = 0.15;
let mut router = PartitionRouter::new(&mut idx)
.level_max(args.common.level_max)
.theta(args.common.theta)
.workers(n_workers)
.max_open(max_open);
router
.run(
args.common.seqfile_paths(),
Some(|p: Progress| {
let now = Instant::now();
let dt = now.duration_since(last_t).as_secs_f64();
if dt > 0.0 {
let instant = (p.position - last_bases) as f64 / dt;
ema_rate = ALPHA * instant + (1.0 - ALPHA) * ema_rate;
}
last_t = now;
last_bases = p.position;
let bp = p.position as f64;
let (count_str, rate_str) = if bp >= 1e9 {
(format!("{:.2} Gbp", bp / 1e9), format!("{:.0} Mbp/s", ema_rate / 1e6))
} else {
(format!("{:.0} Mbp", bp / 1e6), format!("{:.0} Mbp/s", ema_rate / 1e6))
};
pb.set_message(format!("{count_str} {rate_str}"));
}),
)
.unwrap_or_else(|e| {
eprintln!("error: {e}");
std::process::exit(1);
});
pb.finish_and_clear();
rep.push(t.stop());
drop(router); // ends the borrow of `idx` early — `PartitionRouter`'s `Drop` impl would otherwise extend it to the end of scope
idx.mark_scattered().unwrap_or_else(|e| {
eprintln!("error marking scatter done: {e}");
@@ -256,10 +290,31 @@ pub fn run(args: IndexArgs) {
// ── Stage 2: dereplicate + count ─────────────────────────────────────────
if idx.state() < IndexState::Counted {
idx.dereplicate_and_count(args.keep_intermediate, &mut rep).unwrap_or_else(|e| {
let router = PartitionRouter::new(&mut idx);
let t = Stage::start("dereplicate");
router.dereplicate().unwrap_or_else(|e| {
eprintln!("error: {e}");
std::process::exit(1);
});
rep.push(t.stop());
let t = Stage::start("count_kmer");
let spectrum = router.count_kmer(args.keep_intermediate).unwrap_or_else(|e| {
eprintln!("error: {e}");
std::process::exit(1);
});
rep.push(t.stop());
drop(router); // ends the borrow of `idx` early — `PartitionRouter`'s `Drop` impl would otherwise extend it to the end of scope
idx.write_spectrum(spectrum.f0, spectrum.f1, &spectrum.counts).unwrap_or_else(|e| {
eprintln!("error: {e}");
std::process::exit(1);
});
idx.mark_counted().unwrap_or_else(|e| {
eprintln!("error marking count done: {e}");
std::process::exit(1);
});
} else {
info!("dereplicate+count already done, skipping");
}
-1
View File
@@ -1,6 +1,5 @@
mod cli;
mod cmd;
mod steps;
use clap::{Parser, Subcommand};
use tracing_subscriber::{EnvFilter, fmt};
-3
View File
@@ -1,3 +0,0 @@
mod scatter;
pub use scatter::scatter;
-136
View File
@@ -1,136 +0,0 @@
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
use std::time::Instant;
use obikpartitionner::PartitionRouter;
use obipipeline::{ThrottleGuard, Throttled, throttle};
use obiread::NucPage;
use obisys::spinner;
use obisys::{Reporter, Stage};
use tracing::info;
use crate::cli::PipelineData;
// ── Iterator that keeps the slot guard alive until the file is exhausted ──────
struct GuardedIter {
inner: Box<dyn Iterator<Item = NucPage> + Send>,
_guard: ThrottleGuard,
flat_active: Arc<AtomicU32>,
}
impl Iterator for GuardedIter {
type Item = NucPage;
fn next(&mut self) -> Option<NucPage> {
self.inner.next()
}
}
impl Drop for GuardedIter {
fn drop(&mut self) {
self.flat_active.fetch_sub(1, Ordering::Relaxed);
}
}
// ── scatter ───────────────────────────────────────────────────────────────────
/// Run scatter: normalise → build superkmers → route to partition → close.
/// Reports the "scatter" stage to `rep`.
pub fn scatter(
kp: &mut PartitionRouter,
path_source: impl Iterator<Item = PathBuf> + Send + 'static,
k: usize,
level_max: usize,
theta: f64,
n_workers: usize,
max_open: usize,
rep: &mut Reporter,
) {
use obikseq::RoutableSuperKmer;
// Throttle in the source thread — never in a worker — to prevent deadlock.
let throttled = throttle(path_source, max_open);
let file_count = Arc::new(AtomicU64::new(0));
let flat_active = Arc::new(AtomicU32::new(0));
let transform_active = Arc::new(AtomicU32::new(0));
let t = Stage::start("scatter");
let pipe = obipipeline::make_pipe! {
PipelineData : Throttled<PathBuf> => Vec<RoutableSuperKmer>,
||? {
let file_count = Arc::clone(&file_count);
let flat_active = Arc::clone(&flat_active);
let k = k;
move |pw: Throttled<PathBuf>| {
let path = pw.item;
let guard = pw.guard;
let n = file_count.fetch_add(1, Ordering::Relaxed) + 1;
info!("indexing [{}]: {}", n, path.display());
let path_str = path.to_str().unwrap_or("").to_owned();
flat_active.fetch_add(1, Ordering::Relaxed);
obiread::open_nuc_stream(&path_str, k)
.map(|iter| GuardedIter { inner: iter, _guard: guard, flat_active: Arc::clone(&flat_active) })
}
} : Path => NucPage,
| {
let transform_active = Arc::clone(&transform_active);
move |page| {
transform_active.fetch_add(1, Ordering::Relaxed);
let result = obiskbuilder::build_superkmers_page(page, k, level_max, theta);
transform_active.fetch_sub(1, Ordering::Relaxed);
result
}
} : NucPage => Batch,
};
let pb = spinner("scatter");
let mut total_bases: u64 = 0;
let mut ema_rate: f64 = 0.0;
let mut last_t = Instant::now();
let mut last_bases: u64 = 0;
let kmer_overlap = (k - 1) as u64;
const ALPHA: f64 = 0.15;
for batch in pipe.apply(throttled, n_workers, 1) {
total_bases += batch
.iter()
.map(|sk| (sk.seql() as u64).saturating_sub(kmer_overlap))
.sum::<u64>();
let now = Instant::now();
let dt = now.duration_since(last_t).as_secs_f64();
if dt > 0.1 {
let instant = (total_bases - last_bases) as f64 / dt;
ema_rate = ALPHA * instant + (1.0 - ALPHA) * ema_rate;
last_t = now;
last_bases = total_bases;
let bp = total_bases as f64;
let (count_str, rate_str) = if bp >= 1e9 {
(
format!("{:.2} Gbp", bp / 1e9),
format!("{:.0} Mbp/s", ema_rate / 1e6),
)
} else {
(
format!("{:.0} Mbp", bp / 1e6),
format!("{:.0} Mbp/s", ema_rate / 1e6),
)
};
let n_files = file_count.load(Ordering::Relaxed);
let r = flat_active.load(Ordering::Relaxed);
let c = transform_active.load(Ordering::Relaxed);
pb.set_message(format!(
"{count_str} {rate_str} {n_files} files [R:{r} C:{c}]"
));
}
kp.write_batch(batch).unwrap_or_else(|e| {
eprintln!("error: {e}");
std::process::exit(1);
});
}
pb.finish_and_clear();
kp.close().expect("close error");
rep.push(t.stop());
}
+22 -1
View File
@@ -20,11 +20,32 @@
//! entirely) both still reinvent a fragment of this. Migrating them is a
//! separate, deferred step.
use std::path::Path;
use std::path::{Path, PathBuf};
use obikseq::CanonicalKmer;
use obilayeredmap::{layer_dir, IndexMode, Layer, OLMResult};
/// Partition subdirectory name, under an index's root — the single source
/// of truth for the on-disk `partitions/part_NNNNN` naming convention.
/// Moved here from `obikpartitionner` (2026-08-20): naming primitives for
/// the Partition tier belong with the other Partition-tier code, not with
/// the superkmer-routing algorithm that happens to be their first
/// consumer — see `DevDocMD/implementation/partition_layer_cache.md`.
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 {
root.join(PARTITIONS_SUBDIR).join(format!("part_{i:05}"))
}
/// Path of partition `i`'s layered-index directory — `<partition_dir>/index`,
/// the root every [`obilayeredmap::layer_dir`] call for this partition is
/// relative to.
pub 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).
pub struct KmerPartition {
layers: Vec<Layer>,
+7 -5
View File
@@ -6,15 +6,17 @@ edition = "2024"
[dev-dependencies]
tempfile = "3"
obikseq = { path = "../obikseq", features = ["test-utils"] }
obiread = { path = "../obiread" }
obikrope = { path = "../obikrope" }
[dependencies]
niffler = "3.0.0"
remove_dir_all = "1.0"
obikseq = { path = "../obikseq" }
obiskbuilder = { path = "../obiskbuilder" }
obiskio = { path = "../obiskio" }
obikseq = { path = "../obikseq" }
obikindex = { path = "../obikindex" }
obilayeredmap = { path = "../obilayeredmap" }
obipipeline = { path = "../obipipeline" }
obiread = { path = "../obiread" }
obiskbuilder = { path = "../obiskbuilder" }
obiskio = { path = "../obiskio" }
rayon = "1"
sysinfo = "0.39"
serde = { version = "1", features = ["derive"] }
+1 -1
View File
@@ -1,4 +1,4 @@
mod kmer_sort;
mod partition;
pub use partition::{partition_dir, KmerSpectrum, PartitionRouter, PARTITIONS_SUBDIR};
pub use partition::{KmerSpectrum, PartitionRouter};
+8 -7
View File
@@ -1,10 +1,12 @@
//! K-mer partitioning: routing super-kmers into per-partition files,
//! deduplicating them, and counting unique canonical k-mers.
//! K-mer partitioning: routing super-kmers into per-partition, layer-0
//! files, deduplicating them, and counting unique canonical k-mers.
//!
//! Submodules: [`router`] (`PartitionRouter`, `KmerSpectrum`, the
//! `partition_dir` naming convention, and the routing/lifecycle API),
//! [`dereplicate`] (two-phase split+merge deduplication), [`count`]
//! (unique-kmer enumeration, MPHF, abundance counting).
//! routing/lifecycle API — partition/layer path naming itself lives on
//! `obikindex::KmerIndex`, not here, see
//! `DevDocMD/implementation/partition_layer_cache.md`), [`dereplicate`]
//! (two-phase split+merge deduplication), [`count`] (unique-kmer
//! enumeration, MPHF, abundance counting).
mod count;
mod dereplicate;
@@ -13,7 +15,6 @@ mod router;
#[cfg(test)]
mod tests;
pub use router::{partition_dir, KmerSpectrum, PartitionRouter};
pub use router::{KmerSpectrum, PartitionRouter};
const SK_EXT: &str = "skmer.zst";
pub const PARTITIONS_SUBDIR: &str = "partitions";
+205 -77
View File
@@ -1,37 +1,32 @@
use std::collections::BTreeMap;
use std::fs;
use std::io;
use std::path::{Path, PathBuf};
use std::path::PathBuf;
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Instant;
use obisys::progress_bar;
use obikindex::KmerIndex;
use obikseq::RoutableSuperKmer;
use obilayeredmap::Layer;
use obiskio::SKResult;
use obisys::{progress_bar, Progress};
use rayon::prelude::*;
use remove_dir_all::remove_dir_all;
use sysinfo::System;
use tracing::info;
use niffler::Level;
use niffler::send::compression::Format;
use obiskio::SKFileWriter;
use obipipeline::{throttle, ThrottleGuard, Throttled};
use obiread::NucPage;
use crate::kmer_sort::chunk_size_from_ram;
use super::count::count_partition;
use super::dereplicate::{dereplicate_partition, optimal_buckets};
use super::{PARTITIONS_SUBDIR, SK_EXT};
/// Path of partition `i`'s directory under `root` — the single source of
/// truth for the `part_{i:05}` on-disk naming convention. Shared by
/// `PartitionRouter` (which writes here) and, one crate up, `KmerIndex`
/// (which builds `index_dir`/`layer_dir` on top of this same root, once
/// `PartitionRouter` has finished writing) — the only piece of
/// partition-directory knowledge that genuinely needs to cross the
/// crate boundary, since both sides must agree on where a partition lives.
pub fn partition_dir(root: &Path, i: usize) -> PathBuf {
root.join(PARTITIONS_SUBDIR).join(format!("part_{i:05}"))
}
use super::SK_EXT;
pub struct KmerSpectrum {
pub f0: u64,
@@ -39,75 +34,124 @@ pub struct KmerSpectrum {
pub counts: BTreeMap<u32, u64>,
}
/// Routes raw super-kmers into per-partition files, then dereplicates and
/// counts them — this crate's entire job now that layer/query/merge/select/
/// rebuild/dump/distance concerns have moved to `obikindex` (they operate on
/// built layers, which don't exist yet at this stage — see
/// `DevDocMD/implementation/partition_layer_cache.md`). Transient: it makes
/// sense only while raw partition files are being written or processed, a
/// phase that always precedes any `Layer`.
pub struct PartitionRouter {
root_path: PathBuf,
// ── Pipeline plumbing, private to `run` ─────────────────────────────────────
/// Carrier enum for `obipipeline::make_pipe!`'s two-stage transform — local
/// to this crate, not `obikmer`'s own `PipelineData` (which stays scoped to
/// its other CLI commands): a library crate can't depend on the binary
/// that consumes it, so this is a self-contained duplicate of the same
/// shape, not a shared type.
enum PipelineData {
Path(Throttled<PathBuf>),
NucPage(NucPage),
Batch(Vec<RoutableSuperKmer>),
}
unsafe impl Send for PipelineData {}
unsafe impl Sync for PipelineData {}
/// Keeps a file's throttle-slot guard alive until the file's page iterator
/// is exhausted, so the next queued file can only start once this one has
/// actually finished producing pages, not merely been dequeued.
struct GuardedIter {
inner: Box<dyn Iterator<Item = NucPage> + Send>,
_guard: ThrottleGuard,
flat_active: Arc<AtomicU32>,
}
impl Iterator for GuardedIter {
type Item = NucPage;
fn next(&mut self) -> Option<NucPage> {
self.inner.next()
}
}
impl Drop for GuardedIter {
fn drop(&mut self) {
self.flat_active.fetch_sub(1, Ordering::Relaxed);
}
}
// ── PartitionRouter ──────────────────────────────────────────────────────────
/// Routes raw super-kmers into per-partition, layer-0 files, then
/// dereplicates and counts them — this crate's entire job now that
/// layer/query/merge/select/rebuild/dump/distance concerns have moved to
/// `obikindex` (they operate on built layers, which don't exist yet at this
/// stage — see `DevDocMD/implementation/partition_layer_cache.md`).
///
/// Holds `&mut KmerIndex` — this is an algorithm operating on an index, not
/// a data structure of its own; it owns no path-naming knowledge (every
/// path comes from `index.index_dir`/`obilayeredmap::layer_dir`/
/// `Layer::create`), only the transient routing/dereplication/counting
/// state a run needs.
///
/// Two-phase construction: `new` + optional setters (`level_max`/`theta`/
/// `workers`/`max_open`) configure the run, `run` executes it. `run` takes
/// an `Option` progress callback — the router reports raw `(position,
/// total)` ticks via [`obisys::Progress`] and stays unaware of *how* (or
/// whether) the caller displays them; rendering a spinner/progress bar is
/// the caller's decision, not this crate's.
pub struct PartitionRouter<'a> {
index: &'a mut KmerIndex,
n_partitions: usize,
partitions_mask: u64,
writers: Vec<Option<SKFileWriter>>,
level: Level,
closed: bool,
level_max: usize,
theta: f64,
workers: usize,
max_open: usize,
}
impl PartitionRouter {
/// Create a fresh partition layout at `root_path` for `n_partitions =
/// 2^n_bits` partitions.
pub fn create(root_path: &Path, n_bits: usize, force: bool) -> SKResult<Self> {
// `root_path` itself may already exist as a bare directory: callers
// typically hold an index-level lock file there before creating the
// partition layout. What actually signals a pre-existing partition
// set is the `PARTITIONS_SUBDIR` subdirectory, not the root itself.
if root_path.join(PARTITIONS_SUBDIR).exists() {
if force {
remove_dir_all(root_path)?;
} else {
return Err(io::Error::new(
io::ErrorKind::AlreadyExists,
format!("{}: partition directory already exists", root_path.display()),
)
.into());
}
}
fs::create_dir_all(root_path.join(PARTITIONS_SUBDIR))?;
Self::new(root_path, n_bits)
}
/// Reopen the partition layout at `root_path` for further routing (or
/// for `dereplicate`/`count_kmer`, which don't need `writers` but reuse
/// this same handle for consistency). Every caller of `open` wants to
/// write or process, never just read paths (those live on `KmerIndex`
/// directly), so this always starts open too.
pub fn open(root_path: &Path, n_bits: usize) -> SKResult<Self> {
if !root_path.exists() {
return Err(io::Error::new(
io::ErrorKind::NotFound,
format!("{}: partition directory not found", root_path.display()),
)
.into());
}
Self::new(root_path, n_bits)
}
fn new(root_path: &Path, n_bits: usize) -> SKResult<Self> {
impl<'a> PartitionRouter<'a> {
/// Configure a router for `index`'s partition layout. Doesn't touch
/// disk by itself — partitions and their layer-0 shells are created
/// lazily, on the first super-kmer routed to each of them.
pub fn new(index: &'a mut KmerIndex) -> Self {
let n_bits = index.n_bits();
let n_partitions = 1usize << n_bits;
let writers = (0..n_partitions).map(|_| None).collect();
Ok(Self {
root_path: root_path.to_owned(),
let workers = obisys::effective_parallelism();
Self {
index,
n_partitions,
partitions_mask: (1u64 << n_bits) - 1,
writers,
writers: (0..n_partitions).map(|_| None).collect(),
level: Level::One,
closed: false,
})
level_max: 6,
theta: 0.7,
workers,
max_open: (workers / 4).max(1),
}
}
/// Route and write one super-kmer to its partition file.
/// Maximum sub-word size for entropy computation (superkmer building).
pub fn level_max(mut self, v: usize) -> Self {
self.level_max = v;
self
}
/// Entropy threshold (k-mers with score ≤ theta are rejected).
pub fn theta(mut self, v: f64) -> Self {
self.theta = v;
self
}
/// Number of worker threads for `run`'s file-reading/superkmer-building pipeline.
pub fn workers(mut self, v: usize) -> Self {
self.workers = v;
self
}
/// Maximum number of input files `run` opens simultaneously.
pub fn max_open(mut self, v: usize) -> Self {
self.max_open = v;
self
}
/// Route and write one super-kmer to its partition's raw super-kmer file.
pub fn write(&mut self, rsk: RoutableSuperKmer) -> SKResult<()> {
self.check_not_closed()?;
let partition = (rsk.minimizer().seq_hash() & self.partitions_mask) as usize;
@@ -149,6 +193,82 @@ impl PartitionRouter {
!self.closed
}
/// Run the full scatter pipeline: normalise every file in
/// `path_source` -> build super-kmers -> route -> write, then close.
/// `on_progress`, when set, is called periodically (rate-limited, not
/// once per batch) with cumulative bases processed so far —
/// `total: None`, since the total isn't known without pre-scanning
/// every input file.
pub fn run(
&mut self,
path_source: impl Iterator<Item = PathBuf> + Send + 'static,
mut on_progress: Option<impl FnMut(Progress)>,
) -> SKResult<()> {
let k = self.index.kmer_size();
let level_max = self.level_max;
let theta = self.theta;
let n_workers = self.workers;
let max_open = self.max_open;
// Throttle in the source thread — never in a worker — to prevent deadlock.
let throttled = throttle(path_source, max_open);
let file_count = Arc::new(AtomicU64::new(0));
let flat_active = Arc::new(AtomicU32::new(0));
let transform_active = Arc::new(AtomicU32::new(0));
let pipe = obipipeline::make_pipe! {
PipelineData : Throttled<PathBuf> => Vec<RoutableSuperKmer>,
||? {
let file_count = Arc::clone(&file_count);
let flat_active = Arc::clone(&flat_active);
move |pw: Throttled<PathBuf>| {
let path = pw.item;
let guard = pw.guard;
let n = file_count.fetch_add(1, Ordering::Relaxed) + 1;
info!("indexing [{}]: {}", n, path.display());
let path_str = path.to_str().unwrap_or("").to_owned();
flat_active.fetch_add(1, Ordering::Relaxed);
obiread::open_nuc_stream(&path_str, k)
.map(|iter| GuardedIter { inner: iter, _guard: guard, flat_active: Arc::clone(&flat_active) })
}
} : Path => NucPage,
| {
let transform_active = Arc::clone(&transform_active);
move |page| {
transform_active.fetch_add(1, Ordering::Relaxed);
let result = obiskbuilder::build_superkmers_page(page, k, level_max, theta);
transform_active.fetch_sub(1, Ordering::Relaxed);
result
}
} : NucPage => Batch,
};
let mut total_bases: u64 = 0;
let mut last_report = Instant::now();
let kmer_overlap = (k - 1) as u64;
const REPORT_INTERVAL: f64 = 0.1;
for batch in pipe.apply(throttled, n_workers, 1) {
total_bases += batch
.iter()
.map(|sk| (sk.seql() as u64).saturating_sub(kmer_overlap))
.sum::<u64>();
if let Some(cb) = on_progress.as_mut() {
let now = Instant::now();
if now.duration_since(last_report).as_secs_f64() > REPORT_INTERVAL {
last_report = now;
cb(Progress { position: total_bases, total: None });
}
}
self.write_batch(batch)?;
}
if let Some(cb) = on_progress.as_mut() {
cb(Progress { position: total_bases, total: None });
}
self.close()
}
/// Deduplicate all `raw.{ext}` files in parallel, replacing each with a
/// `dereplicated.{ext}` file where identical canonical sequences are merged
/// and their counts summed.
@@ -181,7 +301,7 @@ impl PartitionRouter {
let results: Vec<SKResult<()>> = (0..self.n_partitions)
.into_par_iter()
.map(|i| {
let dir = partition_dir(&self.root_path, i);
let dir = self.layer0_dir(i);
if !dir.exists() {
pb.inc(1);
return Ok(());
@@ -226,7 +346,7 @@ impl PartitionRouter {
let results: Vec<SKResult<()>> = (0..self.n_partitions)
.into_par_iter()
.map(|i| {
let dir = partition_dir(&self.root_path, i);
let dir = self.layer0_dir(i);
let dedup_path = dir.join(format!("dereplicated.{SK_EXT}"));
if !dedup_path.exists() {
pb.inc(1);
@@ -251,7 +371,7 @@ impl PartitionRouter {
let mut f1: u64 = 0;
for i in 0..self.n_partitions {
let path = partition_dir(&self.root_path, i).join("kmer_spectrum_raw.json");
let path = self.layer0_dir(i).join("kmer_spectrum_raw.json");
if !path.exists() {
continue;
}
@@ -276,6 +396,14 @@ impl PartitionRouter {
// ── private ───────────────────────────────────────────────────────────────
/// Directory of partition `i`'s layer 0 — every raw/dereplicated
/// superkmer file and provisional `mphf1.bin`/`counts1.bin` this router
/// produces lives here, alongside where `build_index_layer`
/// (`obikindex`) will later turn it into the real layer 0.
fn layer0_dir(&self, i: usize) -> PathBuf {
obilayeredmap::layer_dir(&self.index.index_dir(i), 0)
}
fn check_not_closed(&self) -> SKResult<()> {
if self.closed {
Err(io::Error::new(io::ErrorKind::BrokenPipe, "write to closed PartitionRouter").into())
@@ -286,8 +414,8 @@ impl PartitionRouter {
fn ensure_writer(&mut self, partition: usize) -> SKResult<&mut SKFileWriter> {
if self.writers[partition].is_none() {
let dir = partition_dir(&self.root_path, partition);
fs::create_dir_all(&dir)?;
let dir = self.layer0_dir(partition);
Layer::create(&dir).map_err(|e| io::Error::other(e.to_string()))?;
let file_path = dir.join(format!("raw.{SK_EXT}"));
let writer = SKFileWriter::create_with(file_path, Format::Zstd, self.level)?;
self.writers[partition] = Some(writer);
@@ -296,7 +424,7 @@ impl PartitionRouter {
}
}
impl Drop for PartitionRouter {
impl Drop for PartitionRouter<'_> {
fn drop(&mut self) {
let _ = self.close();
}
+19 -3
View File
@@ -1,16 +1,30 @@
use std::collections::HashMap;
use std::fs;
use obikindex::{IndexConfig, KmerIndex};
use obikrope::Rope;
use obikseq::SuperKmer;
use obilayeredmap::IndexMode;
use obiskbuilder::build_superkmers;
use super::count::count_partition;
use super::{PartitionRouter, PARTITIONS_SUBDIR};
use super::PartitionRouter;
const K: usize = 11;
const M: usize = 5;
fn test_index(dir: &std::path::Path) -> KmerIndex {
let config = IndexConfig {
kmer_size: K,
minimizer_size: M,
n_bits: 0, // 1 partition — matches this suite's ground-truth setup
with_counts: false,
evidence: IndexMode::Exact,
block_bits: 0,
};
KmerIndex::create(dir, config, None).unwrap()
}
fn setup() {
obikseq::params::set_k(K);
obikseq::params::set_m(M);
@@ -46,12 +60,14 @@ fn pipeline_counts(seqs: &[&[u8]]) -> (u64, u64) {
let superkmers: Vec<_> = build_superkmers(rope, K, 1, 0.0);
let dir = tempfile::tempdir().unwrap();
let mut kp = PartitionRouter::create(dir.path(), 0, true).unwrap();
let mut index = test_index(&dir.path().join("idx"));
let mut kp = PartitionRouter::new(&mut index);
kp.write_batch(superkmers).unwrap();
kp.close().unwrap();
kp.dereplicate().unwrap();
drop(kp); // ends the borrow of `index` early — `PartitionRouter`'s `Drop` impl would otherwise extend it to the end of scope
let part_dir = dir.path().join(PARTITIONS_SUBDIR).join("part_00000");
let part_dir = index.layer_dir(0, 0);
let dedup_path = part_dir.join("dereplicated.skmer.zst");
if !dedup_path.exists() {
return (0, 0);
+9 -3
View File
@@ -7,6 +7,7 @@ use obisys::Reporter;
use tempfile::tempdir;
use obikindex::{GenomeInfo, IndexConfig, KmerIndex, MergeMode};
use obikpartitionner::PartitionRouter;
use super::alignment::SnpAlignmentExt;
use super::build::SiblingAnnexBuildExt;
@@ -64,19 +65,24 @@ fn build_single_genome_index(dir: &Path, label: &str, seq: &[u8]) -> KmerIndex {
evidence: obilayeredmap::IndexMode::Exact,
block_bits: 0,
};
let mut idx = KmerIndex::create(&index_path, config, Some(GenomeInfo::new(label)), false)
let mut idx = KmerIndex::create(&index_path, config, Some(GenomeInfo::new(label)))
.expect("create");
let mut rep = Reporter::new();
let stream = obiread::open_nuc_stream(fasta_path.to_str().unwrap(), K).expect("open fasta");
let mut router = idx.partition_router().expect("partition_router");
let mut router = PartitionRouter::new(&mut idx);
for page in stream {
let batch = obiskbuilder::build_superkmers_page(page, K, /* level_max */ 1, /* theta */ 0.0);
router.write_batch(batch).expect("write_batch");
}
router.close().expect("close partition writers");
router.dereplicate().expect("dereplicate");
let spectrum = router.count_kmer(false).expect("count_kmer");
drop(router); // ends the borrow of `idx` early — `PartitionRouter`'s `Drop` impl would otherwise extend it to the end of scope
idx.mark_scattered().expect("mark_scattered");
idx.dereplicate_and_count(false, &mut rep).expect("dereplicate_and_count");
idx.write_spectrum(spectrum.f0, spectrum.f1, &spectrum.counts).expect("write_spectrum");
idx.mark_counted().expect("mark_counted");
idx.build_layers(1, None, false, &mut rep).expect("build_layers");
idx
}
+1 -1
View File
@@ -9,6 +9,6 @@ mod stage;
pub use budget::MemoryBudget;
pub use lock::DirLock;
pub use progress::{TracedBar, progress_bar, spinner};
pub use progress::{Progress, TracedBar, progress_bar, spinner};
pub use resources::{CpuSample, IoSample, available_memory_bytes, effective_parallelism, peak_rss_bytes};
pub use stage::{Reporter, Stage, StageStats};
+21
View File
@@ -5,6 +5,27 @@ use std::time::{Duration, Instant};
use indicatif::{ProgressBar, ProgressStyle};
use tracing::{debug, info};
// ── Generic progress reporting ──────────────────────────────────────────────
/// One tick of progress from a long-running algorithm, passed to a
/// caller-supplied callback — the standard shape every algo crate's `run()`
/// should use, so a CLI command only has to learn to read this once
/// regardless of which algorithm is reporting.
///
/// `total: None` means the algorithm doesn't know its total work up front
/// (e.g. bases scattered from input files of unknown size until read) — the
/// caller should render a spinner. `total: Some(n)` means it does — the
/// caller can render an actual progress bar. The algorithm itself never
/// decides which; it only reports what it knows, `position` and `total`,
/// and stays unaware of how — or whether — the caller displays it.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Progress {
pub position: u64,
pub total: Option<u64>,
}
// ── Terminal progress bar/spinner ────────────────────────────────────────────
const BRAILLE: &[&str] = &["", "", "", "", "", "", "", "", "", ""];
const ETA_REFRESH_MS: u64 = 500;
const ETA_MIN_ELAPSED_MS: u64 = 1000;