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:
@@ -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
|
||||
|
||||
Generated
+6
-95
@@ -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"
|
||||
|
||||
@@ -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
@@ -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(())
|
||||
|
||||
@@ -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"));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,6 +1,5 @@
|
||||
mod cli;
|
||||
mod cmd;
|
||||
mod steps;
|
||||
|
||||
use clap::{Parser, Subcommand};
|
||||
use tracing_subscriber::{EnvFilter, fmt};
|
||||
|
||||
@@ -1,3 +0,0 @@
|
||||
mod scatter;
|
||||
|
||||
pub use scatter::scatter;
|
||||
@@ -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());
|
||||
}
|
||||
@@ -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>,
|
||||
|
||||
@@ -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,4 +1,4 @@
|
||||
mod kmer_sort;
|
||||
mod partition;
|
||||
|
||||
pub use partition::{partition_dir, KmerSpectrum, PartitionRouter, PARTITIONS_SUBDIR};
|
||||
pub use partition::{KmerSpectrum, PartitionRouter};
|
||||
|
||||
@@ -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";
|
||||
|
||||
@@ -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();
|
||||
}
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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};
|
||||
|
||||
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user