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
@@ -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
|
state, added (panics on every read method). (2a) done — the
|
||||||
`obikpartition` crate and `KmerPartition` itself exist (`open`/`n_layers`/
|
`obikpartition` crate and `KmerPartition` itself exist (`open`/`n_layers`/
|
||||||
`layer`/`layers`/`find`). (2b) — migrating `PartitionCache`/`QueryLayer`
|
`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
|
context: an earlier
|
||||||
version of this doc used the name `KmerPartition` (singular) for what was
|
version of this doc used the name `KmerPartition` (singular) for what was
|
||||||
actually the *collection* type (later renamed `KmerPartitions`, later
|
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
|
`PartitionMeta::load` (see "Remaining instance…" below) — all explicitly
|
||||||
deferred to whenever wiring is tackled next.
|
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
|
## The problem
|
||||||
|
|
||||||
Reading a layer's data (MPHF + matrix) is not free: `MphfLayer::open` mmaps
|
Reading a layer's data (MPHF + matrix) is not free: `MphfLayer::open` mmaps
|
||||||
|
|||||||
Generated
+6
-95
@@ -39,15 +39,6 @@ dependencies = [
|
|||||||
"memchr",
|
"memchr",
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
|
||||||
name = "aligned"
|
|
||||||
version = "0.4.3"
|
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
||||||
checksum = "ee4508988c62edf04abd8d92897fca0c2995d907ce1dfeaf369dac3716a40685"
|
|
||||||
dependencies = [
|
|
||||||
"as-slice",
|
|
||||||
]
|
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "aligned-vec"
|
name = "aligned-vec"
|
||||||
version = "0.6.4"
|
version = "0.6.4"
|
||||||
@@ -137,15 +128,6 @@ version = "0.7.6"
|
|||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50"
|
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]]
|
[[package]]
|
||||||
name = "autocfg"
|
name = "autocfg"
|
||||||
version = "1.5.0"
|
version = "1.5.0"
|
||||||
@@ -340,12 +322,6 @@ version = "1.0.4"
|
|||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801"
|
checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801"
|
||||||
|
|
||||||
[[package]]
|
|
||||||
name = "cfg_aliases"
|
|
||||||
version = "0.2.1"
|
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
||||||
checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724"
|
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "chacha20"
|
name = "chacha20"
|
||||||
version = "0.10.1"
|
version = "0.10.1"
|
||||||
@@ -684,15 +660,6 @@ dependencies = [
|
|||||||
"memchr",
|
"memchr",
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
|
||||||
name = "cvt"
|
|
||||||
version = "0.1.2"
|
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
||||||
checksum = "d2ae9bf77fbf2d39ef573205d554d87e86c12f1994e9ea335b0651b9b278bcf1"
|
|
||||||
dependencies = [
|
|
||||||
"cfg-if",
|
|
||||||
]
|
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "debugid"
|
name = "debugid"
|
||||||
version = "0.8.0"
|
version = "0.8.0"
|
||||||
@@ -913,20 +880,6 @@ version = "0.2.0"
|
|||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb"
|
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]]
|
[[package]]
|
||||||
name = "funty"
|
name = "funty"
|
||||||
version = "2.0.0"
|
version = "2.0.0"
|
||||||
@@ -1422,7 +1375,7 @@ dependencies = [
|
|||||||
"combine",
|
"combine",
|
||||||
"libc",
|
"libc",
|
||||||
"mach2",
|
"mach2",
|
||||||
"nix 0.26.4",
|
"nix",
|
||||||
"sysctl",
|
"sysctl",
|
||||||
"thiserror 1.0.69",
|
"thiserror 1.0.69",
|
||||||
"widestring",
|
"widestring",
|
||||||
@@ -1484,27 +1437,6 @@ dependencies = [
|
|||||||
"pin-utils",
|
"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]]
|
[[package]]
|
||||||
name = "ntapi"
|
name = "ntapi"
|
||||||
version = "0.4.3"
|
version = "0.4.3"
|
||||||
@@ -1606,7 +1538,7 @@ dependencies = [
|
|||||||
"obicompactvec",
|
"obicompactvec",
|
||||||
"obidebruinj",
|
"obidebruinj",
|
||||||
"obikentropy",
|
"obikentropy",
|
||||||
"obikpartitionner",
|
"obikpartition",
|
||||||
"obikseq",
|
"obikseq",
|
||||||
"obilayeredmap",
|
"obilayeredmap",
|
||||||
"obipipeline",
|
"obipipeline",
|
||||||
@@ -1675,15 +1607,17 @@ dependencies = [
|
|||||||
"memmap2",
|
"memmap2",
|
||||||
"niffler",
|
"niffler",
|
||||||
"obicompactvec",
|
"obicompactvec",
|
||||||
|
"obikindex",
|
||||||
"obikrope",
|
"obikrope",
|
||||||
"obikseq",
|
"obikseq",
|
||||||
|
"obilayeredmap",
|
||||||
|
"obipipeline",
|
||||||
"obiread",
|
"obiread",
|
||||||
"obiskbuilder",
|
"obiskbuilder",
|
||||||
"obiskio",
|
"obiskio",
|
||||||
"obisys",
|
"obisys",
|
||||||
"ptr_hash",
|
"ptr_hash",
|
||||||
"rayon",
|
"rayon",
|
||||||
"remove_dir_all",
|
|
||||||
"serde",
|
"serde",
|
||||||
"serde_json",
|
"serde_json",
|
||||||
"sysinfo",
|
"sysinfo",
|
||||||
@@ -2016,7 +1950,7 @@ dependencies = [
|
|||||||
"findshlibs",
|
"findshlibs",
|
||||||
"libc",
|
"libc",
|
||||||
"log",
|
"log",
|
||||||
"nix 0.26.4",
|
"nix",
|
||||||
"once_cell",
|
"once_cell",
|
||||||
"prost",
|
"prost",
|
||||||
"prost-build",
|
"prost-build",
|
||||||
@@ -2371,20 +2305,6 @@ version = "0.8.11"
|
|||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4"
|
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]]
|
[[package]]
|
||||||
name = "ring"
|
name = "ring"
|
||||||
version = "0.17.14"
|
version = "0.17.14"
|
||||||
@@ -3312,15 +3232,6 @@ dependencies = [
|
|||||||
"windows-targets 0.52.6",
|
"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]]
|
[[package]]
|
||||||
name = "windows-sys"
|
name = "windows-sys"
|
||||||
version = "0.61.2"
|
version = "0.61.2"
|
||||||
|
|||||||
@@ -5,7 +5,7 @@ edition = "2024"
|
|||||||
|
|
||||||
[dependencies]
|
[dependencies]
|
||||||
obikseq = { path = "../obikseq" }
|
obikseq = { path = "../obikseq" }
|
||||||
obikpartitionner = { path = "../obikpartitionner" }
|
obikpartition = { path = "../obikpartition" }
|
||||||
obitaxonomy = { path = "../obitaxonomy" }
|
obitaxonomy = { path = "../obitaxonomy" }
|
||||||
obiskio = { path = "../obiskio" }
|
obiskio = { path = "../obiskio" }
|
||||||
obisys = { path = "../obisys" }
|
obisys = { path = "../obisys" }
|
||||||
|
|||||||
+19
-41
@@ -2,7 +2,6 @@ use std::collections::BTreeMap;
|
|||||||
use std::fs;
|
use std::fs;
|
||||||
use std::path::{Path, PathBuf};
|
use std::path::{Path, PathBuf};
|
||||||
|
|
||||||
use obikpartitionner::{KmerSpectrum, PartitionRouter};
|
|
||||||
use obilayeredmap::meta::PartitionMeta;
|
use obilayeredmap::meta::PartitionMeta;
|
||||||
use obisys::{Reporter, Stage, progress_bar};
|
use obisys::{Reporter, Stage, progress_bar};
|
||||||
use rayon::prelude::*;
|
use rayon::prelude::*;
|
||||||
@@ -29,10 +28,8 @@ impl KmerIndex {
|
|||||||
path: P,
|
path: P,
|
||||||
config: IndexConfig,
|
config: IndexConfig,
|
||||||
genome_info: Option<GenomeInfo>,
|
genome_info: Option<GenomeInfo>,
|
||||||
force: bool,
|
|
||||||
) -> OKIResult<Self> {
|
) -> OKIResult<Self> {
|
||||||
let root_path = path.as_ref().to_owned();
|
let root_path = path.as_ref().to_owned();
|
||||||
PartitionRouter::create(&root_path, config.n_bits, force)?;
|
|
||||||
set_k(config.kmer_size);
|
set_k(config.kmer_size);
|
||||||
set_m(config.minimizer_size);
|
set_m(config.minimizer_size);
|
||||||
let mut meta = IndexMeta::new(config);
|
let mut meta = IndexMeta::new(config);
|
||||||
@@ -86,7 +83,6 @@ impl KmerIndex {
|
|||||||
let output = output.as_ref();
|
let output = output.as_ref();
|
||||||
fs::create_dir_all(output).map_err(OKIError::Io)?;
|
fs::create_dir_all(output).map_err(OKIError::Io)?;
|
||||||
meta.write(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() })
|
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}`) —
|
/// Path of partition `i`'s raw directory (`partitions/part_{i:05}`) —
|
||||||
/// the on-disk naming convention `obikpartitionner::PartitionRouter`
|
/// delegates to `obikpartition`, the Partition tier's own naming
|
||||||
/// also writes to (raw/dereplicated superkmer files, `mphf1.bin`,
|
/// primitive (mirrors `layer_dir` delegating to `obilayeredmap`).
|
||||||
/// `counts1.bin`), the single point of agreement between the two.
|
/// `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 {
|
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`).
|
/// Path of partition `i`'s layered-index directory (`<partition>/index`).
|
||||||
pub fn index_dir(&self, i: usize) -> PathBuf {
|
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.
|
/// Path of layer `l` within partition `i`'s layered index.
|
||||||
@@ -223,39 +222,19 @@ impl KmerIndex {
|
|||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Open a fresh [`PartitionRouter`] onto this index's partition layout
|
/// Mark dereplicate+count as complete and write `count.done`.
|
||||||
/// — the write-side handle for `scatter`, or for `dereplicate_and_count`
|
pub fn mark_counted(&self) -> OKIResult<()> {
|
||||||
/// 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)?;
|
|
||||||
touch(&self.root_path.join(SENTINEL_COUNTED))?;
|
touch(&self.root_path.join(SENTINEL_COUNTED))?;
|
||||||
Ok(())
|
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
|
let label = self
|
||||||
.meta
|
.meta
|
||||||
.genomes
|
.genomes
|
||||||
@@ -265,15 +244,14 @@ impl KmerIndex {
|
|||||||
let spectrums_dir = self.root_path.join("spectrums");
|
let spectrums_dir = self.root_path.join("spectrums");
|
||||||
fs::create_dir_all(&spectrums_dir)?;
|
fs::create_dir_all(&spectrums_dir)?;
|
||||||
let path = spectrums_dir.join(format!("{label}.json"));
|
let path = spectrums_dir.join(format!("{label}.json"));
|
||||||
let spectrum_map: BTreeMap<String, u64> = sp
|
let spectrum_map: BTreeMap<String, u64> = counts
|
||||||
.counts
|
|
||||||
.iter()
|
.iter()
|
||||||
.map(|(&c, &f)| (format!("{c:010}"), f))
|
.map(|(&c, &f)| (format!("{c:010}"), f))
|
||||||
.collect();
|
.collect();
|
||||||
let f = fs::File::create(&path)?;
|
let f = fs::File::create(&path)?;
|
||||||
serde_json::to_writer_pretty(
|
serde_json::to_writer_pretty(
|
||||||
f,
|
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)?;
|
.map_err(OKIError::Json)?;
|
||||||
Ok(())
|
Ok(())
|
||||||
|
|||||||
@@ -42,14 +42,13 @@ impl KmerIndex {
|
|||||||
mode: &IndexMode,
|
mode: &IndexMode,
|
||||||
block_bits: u8,
|
block_bits: u8,
|
||||||
) -> Result<usize, SKError> {
|
) -> Result<usize, SKError> {
|
||||||
let partition_dir = self.partition_dir(i);
|
let layer0_dir = self.layer_dir(i, 0);
|
||||||
let dedup_path = partition_dir.join("dereplicated.skmer.zst");
|
let dedup_path = layer0_dir.join("dereplicated.skmer.zst");
|
||||||
if !dedup_path.exists() {
|
if !dedup_path.exists() {
|
||||||
return Ok(0);
|
return Ok(0);
|
||||||
}
|
}
|
||||||
|
|
||||||
let layer_dir = self.layer_dir(i, 0);
|
if layer0_dir.join("mphf.bin").exists() {
|
||||||
if layer_dir.join("mphf.bin").exists() {
|
|
||||||
return Ok(0);
|
return Ok(0);
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -57,14 +56,14 @@ impl KmerIndex {
|
|||||||
let need_counts = filter_active || with_counts;
|
let need_counts = filter_active || with_counts;
|
||||||
|
|
||||||
let mphf1_opt: Option<Mphf> = if need_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()
|
p.exists().then(|| Mphf::load_full(&p).ok()).flatten()
|
||||||
} else {
|
} else {
|
||||||
None
|
None
|
||||||
};
|
};
|
||||||
|
|
||||||
let counts1_opt: Option<PersistentCompactIntVec> = if need_counts {
|
let counts1_opt: Option<PersistentCompactIntVec> = if need_counts {
|
||||||
let p = partition_dir.join("counts1.bin");
|
let p = layer0_dir.join("counts1.bin");
|
||||||
p.exists()
|
p.exists()
|
||||||
.then(|| PersistentCompactIntVec::open(&p).ok())
|
.then(|| PersistentCompactIntVec::open(&p).ok())
|
||||||
.flatten()
|
.flatten()
|
||||||
@@ -95,8 +94,8 @@ impl KmerIndex {
|
|||||||
|
|
||||||
let n_kmers =
|
let n_kmers =
|
||||||
if with_counts {
|
if with_counts {
|
||||||
let n = write_graph_as_unitigs(g, &layer_dir)?;
|
let n = write_graph_as_unitigs(g, &layer0_dir)?;
|
||||||
TypedLayer::<PersistentCompactIntMatrix>::build(&layer_dir, block_bits, mode, |kmer| {
|
TypedLayer::<PersistentCompactIntMatrix>::build(&layer0_dir, block_bits, mode, |kmer| {
|
||||||
match (&mphf1_opt, &counts1_opt) {
|
match (&mphf1_opt, &counts1_opt) {
|
||||||
(Some(mphf), Some(counts)) => counts.get(mphf.index(&kmer.raw())),
|
(Some(mphf), Some(counts)) => counts.get(mphf.index(&kmer.raw())),
|
||||||
_ => 1,
|
_ => 1,
|
||||||
@@ -105,10 +104,10 @@ impl KmerIndex {
|
|||||||
.map_err(|e| olm_to_sk(e, "layer build"))?;
|
.map_err(|e| olm_to_sk(e, "layer build"))?;
|
||||||
n
|
n
|
||||||
} else {
|
} 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 {
|
PartitionMeta {
|
||||||
n_layers: 1,
|
n_layers: 1,
|
||||||
mode: mode.clone(),
|
mode: mode.clone(),
|
||||||
@@ -123,11 +122,11 @@ impl KmerIndex {
|
|||||||
///
|
///
|
||||||
/// Deletes `dereplicated.skmer.zst` (+ sidecar), `mphf1.bin`, `counts1.bin`.
|
/// Deletes `dereplicated.skmer.zst` (+ sidecar), `mphf1.bin`, `counts1.bin`.
|
||||||
pub fn remove_build_artifacts(&self, i: usize) {
|
pub fn remove_build_artifacts(&self, i: usize) {
|
||||||
let partition_dir = self.partition_dir(i);
|
let layer0_dir = self.layer_dir(i, 0);
|
||||||
let dedup = partition_dir.join("dereplicated.skmer.zst");
|
let dedup = layer0_dir.join("dereplicated.skmer.zst");
|
||||||
remove_if_exists(&SKFileMeta::sidecar_path(&dedup));
|
remove_if_exists(&SKFileMeta::sidecar_path(&dedup));
|
||||||
remove_if_exists(&dedup);
|
remove_if_exists(&dedup);
|
||||||
remove_if_exists(&partition_dir.join("mphf1.bin"));
|
remove_if_exists(&layer0_dir.join("mphf1.bin"));
|
||||||
remove_if_exists(&partition_dir.join("counts1.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,
|
evidence: obilayeredmap::IndexMode::Exact,
|
||||||
block_bits: 0,
|
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();
|
let mut kmers: HashMap<CanonicalKmer, Vec<KmerDesc>> = HashMap::new();
|
||||||
// Any well-formed canonical k-mer works here — the call must return
|
// 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,
|
evidence: obilayeredmap::IndexMode::Exact,
|
||||||
block_bits: 0,
|
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 kmers: HashMap<CanonicalKmer, Vec<KmerDesc>> = HashMap::new();
|
||||||
let stats = index
|
let stats = index
|
||||||
|
|||||||
@@ -1,18 +1,19 @@
|
|||||||
use std::path::PathBuf;
|
use std::path::PathBuf;
|
||||||
|
use std::time::Instant;
|
||||||
|
|
||||||
use clap::Args;
|
use clap::Args;
|
||||||
use obikindex::{validate_label, GenomeInfo, IndexConfig, IndexState, KmerIndex};
|
use obikindex::{validate_label, GenomeInfo, IndexConfig, IndexState, KmerIndex};
|
||||||
|
use obikpartitionner::PartitionRouter;
|
||||||
use obilayeredmap::IndexMode;
|
use obilayeredmap::IndexMode;
|
||||||
|
|
||||||
fn parse_key_value(s: &str) -> Result<(String, String), String> {
|
fn parse_key_value(s: &str) -> Result<(String, String), String> {
|
||||||
let pos = s.find('=').ok_or_else(|| format!("invalid key=value: no '=' in '{s}'"))?;
|
let pos = s.find('=').ok_or_else(|| format!("invalid key=value: no '=' in '{s}'"))?;
|
||||||
Ok((s[..pos].to_string(), s[pos + 1..].to_string()))
|
Ok((s[..pos].to_string(), s[pos + 1..].to_string()))
|
||||||
}
|
}
|
||||||
use obisys::Reporter;
|
use obisys::{spinner, Progress, Reporter, Stage};
|
||||||
use tracing::info;
|
use tracing::info;
|
||||||
|
|
||||||
use crate::cli::{CommonArgs, block_size_to_bits, partitions_to_bits};
|
use crate::cli::{CommonArgs, block_size_to_bits, partitions_to_bits};
|
||||||
use crate::steps::scatter;
|
|
||||||
|
|
||||||
#[derive(Args)]
|
#[derive(Args)]
|
||||||
pub struct IndexArgs {
|
pub struct IndexArgs {
|
||||||
@@ -225,7 +226,7 @@ pub fn run(args: IndexArgs) {
|
|||||||
}
|
}
|
||||||
info
|
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}");
|
eprintln!("error creating index: {e}");
|
||||||
std::process::exit(1);
|
std::process::exit(1);
|
||||||
})
|
})
|
||||||
@@ -234,17 +235,50 @@ pub fn run(args: IndexArgs) {
|
|||||||
|
|
||||||
// ── Stage 1: scatter ─────────────────────────────────────────────────────
|
// ── Stage 1: scatter ─────────────────────────────────────────────────────
|
||||||
if idx.state() < IndexState::Scattered {
|
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 n_workers = args.common.threads.max(1);
|
||||||
|
|
||||||
let max_open = args.common.effective_max_open();
|
let max_open = args.common.effective_max_open();
|
||||||
let mut router = idx.partition_router().unwrap_or_else(|e| {
|
|
||||||
eprintln!("error opening partition router: {e}");
|
let t = Stage::start("scatter");
|
||||||
std::process::exit(1);
|
let pb = spinner("scatter");
|
||||||
});
|
let mut ema_rate: f64 = 0.0;
|
||||||
scatter(&mut router, args.common.seqfile_paths(), k, level_max, theta, n_workers, max_open, &mut rep);
|
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| {
|
idx.mark_scattered().unwrap_or_else(|e| {
|
||||||
eprintln!("error marking scatter done: {e}");
|
eprintln!("error marking scatter done: {e}");
|
||||||
@@ -256,10 +290,31 @@ pub fn run(args: IndexArgs) {
|
|||||||
|
|
||||||
// ── Stage 2: dereplicate + count ─────────────────────────────────────────
|
// ── Stage 2: dereplicate + count ─────────────────────────────────────────
|
||||||
if idx.state() < IndexState::Counted {
|
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}");
|
eprintln!("error: {e}");
|
||||||
std::process::exit(1);
|
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 {
|
} else {
|
||||||
info!("dereplicate+count already done, skipping");
|
info!("dereplicate+count already done, skipping");
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,6 +1,5 @@
|
|||||||
mod cli;
|
mod cli;
|
||||||
mod cmd;
|
mod cmd;
|
||||||
mod steps;
|
|
||||||
|
|
||||||
use clap::{Parser, Subcommand};
|
use clap::{Parser, Subcommand};
|
||||||
use tracing_subscriber::{EnvFilter, fmt};
|
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
|
//! entirely) both still reinvent a fragment of this. Migrating them is a
|
||||||
//! separate, deferred step.
|
//! separate, deferred step.
|
||||||
|
|
||||||
use std::path::Path;
|
use std::path::{Path, PathBuf};
|
||||||
|
|
||||||
use obikseq::CanonicalKmer;
|
use obikseq::CanonicalKmer;
|
||||||
use obilayeredmap::{layer_dir, IndexMode, Layer, OLMResult};
|
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).
|
/// One partition's open layers, in layer order (layer 0 first).
|
||||||
pub struct KmerPartition {
|
pub struct KmerPartition {
|
||||||
layers: Vec<Layer>,
|
layers: Vec<Layer>,
|
||||||
|
|||||||
@@ -6,15 +6,17 @@ edition = "2024"
|
|||||||
[dev-dependencies]
|
[dev-dependencies]
|
||||||
tempfile = "3"
|
tempfile = "3"
|
||||||
obikseq = { path = "../obikseq", features = ["test-utils"] }
|
obikseq = { path = "../obikseq", features = ["test-utils"] }
|
||||||
obiread = { path = "../obiread" }
|
|
||||||
obikrope = { path = "../obikrope" }
|
obikrope = { path = "../obikrope" }
|
||||||
|
|
||||||
[dependencies]
|
[dependencies]
|
||||||
niffler = "3.0.0"
|
niffler = "3.0.0"
|
||||||
remove_dir_all = "1.0"
|
obikseq = { path = "../obikseq" }
|
||||||
obikseq = { path = "../obikseq" }
|
obikindex = { path = "../obikindex" }
|
||||||
obiskbuilder = { path = "../obiskbuilder" }
|
obilayeredmap = { path = "../obilayeredmap" }
|
||||||
obiskio = { path = "../obiskio" }
|
obipipeline = { path = "../obipipeline" }
|
||||||
|
obiread = { path = "../obiread" }
|
||||||
|
obiskbuilder = { path = "../obiskbuilder" }
|
||||||
|
obiskio = { path = "../obiskio" }
|
||||||
rayon = "1"
|
rayon = "1"
|
||||||
sysinfo = "0.39"
|
sysinfo = "0.39"
|
||||||
serde = { version = "1", features = ["derive"] }
|
serde = { version = "1", features = ["derive"] }
|
||||||
|
|||||||
@@ -1,4 +1,4 @@
|
|||||||
mod kmer_sort;
|
mod kmer_sort;
|
||||||
mod partition;
|
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,
|
//! K-mer partitioning: routing super-kmers into per-partition, layer-0
|
||||||
//! deduplicating them, and counting unique canonical k-mers.
|
//! files, deduplicating them, and counting unique canonical k-mers.
|
||||||
//!
|
//!
|
||||||
//! Submodules: [`router`] (`PartitionRouter`, `KmerSpectrum`, the
|
//! Submodules: [`router`] (`PartitionRouter`, `KmerSpectrum`, the
|
||||||
//! `partition_dir` naming convention, and the routing/lifecycle API),
|
//! routing/lifecycle API — partition/layer path naming itself lives on
|
||||||
//! [`dereplicate`] (two-phase split+merge deduplication), [`count`]
|
//! `obikindex::KmerIndex`, not here, see
|
||||||
//! (unique-kmer enumeration, MPHF, abundance counting).
|
//! `DevDocMD/implementation/partition_layer_cache.md`), [`dereplicate`]
|
||||||
|
//! (two-phase split+merge deduplication), [`count`] (unique-kmer
|
||||||
|
//! enumeration, MPHF, abundance counting).
|
||||||
|
|
||||||
mod count;
|
mod count;
|
||||||
mod dereplicate;
|
mod dereplicate;
|
||||||
@@ -13,7 +15,6 @@ mod router;
|
|||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests;
|
mod tests;
|
||||||
|
|
||||||
pub use router::{partition_dir, KmerSpectrum, PartitionRouter};
|
pub use router::{KmerSpectrum, PartitionRouter};
|
||||||
|
|
||||||
const SK_EXT: &str = "skmer.zst";
|
const SK_EXT: &str = "skmer.zst";
|
||||||
pub const PARTITIONS_SUBDIR: &str = "partitions";
|
|
||||||
|
|||||||
@@ -1,37 +1,32 @@
|
|||||||
use std::collections::BTreeMap;
|
use std::collections::BTreeMap;
|
||||||
use std::fs;
|
use std::fs;
|
||||||
use std::io;
|
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 std::time::Instant;
|
||||||
|
|
||||||
use obisys::progress_bar;
|
use obikindex::KmerIndex;
|
||||||
|
|
||||||
use obikseq::RoutableSuperKmer;
|
use obikseq::RoutableSuperKmer;
|
||||||
|
use obilayeredmap::Layer;
|
||||||
use obiskio::SKResult;
|
use obiskio::SKResult;
|
||||||
|
use obisys::{progress_bar, Progress};
|
||||||
use rayon::prelude::*;
|
use rayon::prelude::*;
|
||||||
use remove_dir_all::remove_dir_all;
|
|
||||||
use sysinfo::System;
|
use sysinfo::System;
|
||||||
|
use tracing::info;
|
||||||
|
|
||||||
use niffler::Level;
|
use niffler::Level;
|
||||||
use niffler::send::compression::Format;
|
use niffler::send::compression::Format;
|
||||||
use obiskio::SKFileWriter;
|
use obiskio::SKFileWriter;
|
||||||
|
|
||||||
|
use obipipeline::{throttle, ThrottleGuard, Throttled};
|
||||||
|
use obiread::NucPage;
|
||||||
|
|
||||||
use crate::kmer_sort::chunk_size_from_ram;
|
use crate::kmer_sort::chunk_size_from_ram;
|
||||||
|
|
||||||
use super::count::count_partition;
|
use super::count::count_partition;
|
||||||
use super::dereplicate::{dereplicate_partition, optimal_buckets};
|
use super::dereplicate::{dereplicate_partition, optimal_buckets};
|
||||||
use super::{PARTITIONS_SUBDIR, SK_EXT};
|
use super::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}"))
|
|
||||||
}
|
|
||||||
|
|
||||||
pub struct KmerSpectrum {
|
pub struct KmerSpectrum {
|
||||||
pub f0: u64,
|
pub f0: u64,
|
||||||
@@ -39,75 +34,124 @@ pub struct KmerSpectrum {
|
|||||||
pub counts: BTreeMap<u32, u64>,
|
pub counts: BTreeMap<u32, u64>,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Routes raw super-kmers into per-partition files, then dereplicates and
|
// ── Pipeline plumbing, private to `run` ─────────────────────────────────────
|
||||||
/// counts them — this crate's entire job now that layer/query/merge/select/
|
|
||||||
/// rebuild/dump/distance concerns have moved to `obikindex` (they operate on
|
/// Carrier enum for `obipipeline::make_pipe!`'s two-stage transform — local
|
||||||
/// built layers, which don't exist yet at this stage — see
|
/// to this crate, not `obikmer`'s own `PipelineData` (which stays scoped to
|
||||||
/// `DevDocMD/implementation/partition_layer_cache.md`). Transient: it makes
|
/// its other CLI commands): a library crate can't depend on the binary
|
||||||
/// sense only while raw partition files are being written or processed, a
|
/// that consumes it, so this is a self-contained duplicate of the same
|
||||||
/// phase that always precedes any `Layer`.
|
/// shape, not a shared type.
|
||||||
pub struct PartitionRouter {
|
enum PipelineData {
|
||||||
root_path: PathBuf,
|
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,
|
n_partitions: usize,
|
||||||
partitions_mask: u64,
|
partitions_mask: u64,
|
||||||
writers: Vec<Option<SKFileWriter>>,
|
writers: Vec<Option<SKFileWriter>>,
|
||||||
level: Level,
|
level: Level,
|
||||||
closed: bool,
|
closed: bool,
|
||||||
|
level_max: usize,
|
||||||
|
theta: f64,
|
||||||
|
workers: usize,
|
||||||
|
max_open: usize,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl PartitionRouter {
|
impl<'a> PartitionRouter<'a> {
|
||||||
/// Create a fresh partition layout at `root_path` for `n_partitions =
|
/// Configure a router for `index`'s partition layout. Doesn't touch
|
||||||
/// 2^n_bits` partitions.
|
/// disk by itself — partitions and their layer-0 shells are created
|
||||||
pub fn create(root_path: &Path, n_bits: usize, force: bool) -> SKResult<Self> {
|
/// lazily, on the first super-kmer routed to each of them.
|
||||||
// `root_path` itself may already exist as a bare directory: callers
|
pub fn new(index: &'a mut KmerIndex) -> Self {
|
||||||
// typically hold an index-level lock file there before creating the
|
let n_bits = index.n_bits();
|
||||||
// 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> {
|
|
||||||
let n_partitions = 1usize << n_bits;
|
let n_partitions = 1usize << n_bits;
|
||||||
let writers = (0..n_partitions).map(|_| None).collect();
|
let workers = obisys::effective_parallelism();
|
||||||
Ok(Self {
|
Self {
|
||||||
root_path: root_path.to_owned(),
|
index,
|
||||||
n_partitions,
|
n_partitions,
|
||||||
partitions_mask: (1u64 << n_bits) - 1,
|
partitions_mask: (1u64 << n_bits) - 1,
|
||||||
writers,
|
writers: (0..n_partitions).map(|_| None).collect(),
|
||||||
level: Level::One,
|
level: Level::One,
|
||||||
closed: false,
|
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<()> {
|
pub fn write(&mut self, rsk: RoutableSuperKmer) -> SKResult<()> {
|
||||||
self.check_not_closed()?;
|
self.check_not_closed()?;
|
||||||
let partition = (rsk.minimizer().seq_hash() & self.partitions_mask) as usize;
|
let partition = (rsk.minimizer().seq_hash() & self.partitions_mask) as usize;
|
||||||
@@ -149,6 +193,82 @@ impl PartitionRouter {
|
|||||||
!self.closed
|
!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
|
/// Deduplicate all `raw.{ext}` files in parallel, replacing each with a
|
||||||
/// `dereplicated.{ext}` file where identical canonical sequences are merged
|
/// `dereplicated.{ext}` file where identical canonical sequences are merged
|
||||||
/// and their counts summed.
|
/// and their counts summed.
|
||||||
@@ -181,7 +301,7 @@ impl PartitionRouter {
|
|||||||
let results: Vec<SKResult<()>> = (0..self.n_partitions)
|
let results: Vec<SKResult<()>> = (0..self.n_partitions)
|
||||||
.into_par_iter()
|
.into_par_iter()
|
||||||
.map(|i| {
|
.map(|i| {
|
||||||
let dir = partition_dir(&self.root_path, i);
|
let dir = self.layer0_dir(i);
|
||||||
if !dir.exists() {
|
if !dir.exists() {
|
||||||
pb.inc(1);
|
pb.inc(1);
|
||||||
return Ok(());
|
return Ok(());
|
||||||
@@ -226,7 +346,7 @@ impl PartitionRouter {
|
|||||||
let results: Vec<SKResult<()>> = (0..self.n_partitions)
|
let results: Vec<SKResult<()>> = (0..self.n_partitions)
|
||||||
.into_par_iter()
|
.into_par_iter()
|
||||||
.map(|i| {
|
.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}"));
|
let dedup_path = dir.join(format!("dereplicated.{SK_EXT}"));
|
||||||
if !dedup_path.exists() {
|
if !dedup_path.exists() {
|
||||||
pb.inc(1);
|
pb.inc(1);
|
||||||
@@ -251,7 +371,7 @@ impl PartitionRouter {
|
|||||||
let mut f1: u64 = 0;
|
let mut f1: u64 = 0;
|
||||||
|
|
||||||
for i in 0..self.n_partitions {
|
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() {
|
if !path.exists() {
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
@@ -276,6 +396,14 @@ impl PartitionRouter {
|
|||||||
|
|
||||||
// ── private ───────────────────────────────────────────────────────────────
|
// ── 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<()> {
|
fn check_not_closed(&self) -> SKResult<()> {
|
||||||
if self.closed {
|
if self.closed {
|
||||||
Err(io::Error::new(io::ErrorKind::BrokenPipe, "write to closed PartitionRouter").into())
|
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> {
|
fn ensure_writer(&mut self, partition: usize) -> SKResult<&mut SKFileWriter> {
|
||||||
if self.writers[partition].is_none() {
|
if self.writers[partition].is_none() {
|
||||||
let dir = partition_dir(&self.root_path, partition);
|
let dir = self.layer0_dir(partition);
|
||||||
fs::create_dir_all(&dir)?;
|
Layer::create(&dir).map_err(|e| io::Error::other(e.to_string()))?;
|
||||||
let file_path = dir.join(format!("raw.{SK_EXT}"));
|
let file_path = dir.join(format!("raw.{SK_EXT}"));
|
||||||
let writer = SKFileWriter::create_with(file_path, Format::Zstd, self.level)?;
|
let writer = SKFileWriter::create_with(file_path, Format::Zstd, self.level)?;
|
||||||
self.writers[partition] = Some(writer);
|
self.writers[partition] = Some(writer);
|
||||||
@@ -296,7 +424,7 @@ impl PartitionRouter {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Drop for PartitionRouter {
|
impl Drop for PartitionRouter<'_> {
|
||||||
fn drop(&mut self) {
|
fn drop(&mut self) {
|
||||||
let _ = self.close();
|
let _ = self.close();
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,16 +1,30 @@
|
|||||||
use std::collections::HashMap;
|
use std::collections::HashMap;
|
||||||
use std::fs;
|
use std::fs;
|
||||||
|
|
||||||
|
use obikindex::{IndexConfig, KmerIndex};
|
||||||
use obikrope::Rope;
|
use obikrope::Rope;
|
||||||
use obikseq::SuperKmer;
|
use obikseq::SuperKmer;
|
||||||
|
use obilayeredmap::IndexMode;
|
||||||
use obiskbuilder::build_superkmers;
|
use obiskbuilder::build_superkmers;
|
||||||
|
|
||||||
use super::count::count_partition;
|
use super::count::count_partition;
|
||||||
use super::{PartitionRouter, PARTITIONS_SUBDIR};
|
use super::PartitionRouter;
|
||||||
|
|
||||||
const K: usize = 11;
|
const K: usize = 11;
|
||||||
const M: usize = 5;
|
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() {
|
fn setup() {
|
||||||
obikseq::params::set_k(K);
|
obikseq::params::set_k(K);
|
||||||
obikseq::params::set_m(M);
|
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 superkmers: Vec<_> = build_superkmers(rope, K, 1, 0.0);
|
||||||
|
|
||||||
let dir = tempfile::tempdir().unwrap();
|
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.write_batch(superkmers).unwrap();
|
||||||
kp.close().unwrap();
|
kp.close().unwrap();
|
||||||
kp.dereplicate().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");
|
let dedup_path = part_dir.join("dereplicated.skmer.zst");
|
||||||
if !dedup_path.exists() {
|
if !dedup_path.exists() {
|
||||||
return (0, 0);
|
return (0, 0);
|
||||||
|
|||||||
@@ -7,6 +7,7 @@ use obisys::Reporter;
|
|||||||
use tempfile::tempdir;
|
use tempfile::tempdir;
|
||||||
|
|
||||||
use obikindex::{GenomeInfo, IndexConfig, KmerIndex, MergeMode};
|
use obikindex::{GenomeInfo, IndexConfig, KmerIndex, MergeMode};
|
||||||
|
use obikpartitionner::PartitionRouter;
|
||||||
|
|
||||||
use super::alignment::SnpAlignmentExt;
|
use super::alignment::SnpAlignmentExt;
|
||||||
use super::build::SiblingAnnexBuildExt;
|
use super::build::SiblingAnnexBuildExt;
|
||||||
@@ -64,19 +65,24 @@ fn build_single_genome_index(dir: &Path, label: &str, seq: &[u8]) -> KmerIndex {
|
|||||||
evidence: obilayeredmap::IndexMode::Exact,
|
evidence: obilayeredmap::IndexMode::Exact,
|
||||||
block_bits: 0,
|
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");
|
.expect("create");
|
||||||
|
|
||||||
let mut rep = Reporter::new();
|
let mut rep = Reporter::new();
|
||||||
let stream = obiread::open_nuc_stream(fasta_path.to_str().unwrap(), K).expect("open fasta");
|
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 {
|
for page in stream {
|
||||||
let batch = obiskbuilder::build_superkmers_page(page, K, /* level_max */ 1, /* theta */ 0.0);
|
let batch = obiskbuilder::build_superkmers_page(page, K, /* level_max */ 1, /* theta */ 0.0);
|
||||||
router.write_batch(batch).expect("write_batch");
|
router.write_batch(batch).expect("write_batch");
|
||||||
}
|
}
|
||||||
router.close().expect("close partition writers");
|
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.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.build_layers(1, None, false, &mut rep).expect("build_layers");
|
||||||
idx
|
idx
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -9,6 +9,6 @@ mod stage;
|
|||||||
|
|
||||||
pub use budget::MemoryBudget;
|
pub use budget::MemoryBudget;
|
||||||
pub use lock::DirLock;
|
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 resources::{CpuSample, IoSample, available_memory_bytes, effective_parallelism, peak_rss_bytes};
|
||||||
pub use stage::{Reporter, Stage, StageStats};
|
pub use stage::{Reporter, Stage, StageStats};
|
||||||
|
|||||||
@@ -5,6 +5,27 @@ use std::time::{Duration, Instant};
|
|||||||
use indicatif::{ProgressBar, ProgressStyle};
|
use indicatif::{ProgressBar, ProgressStyle};
|
||||||
use tracing::{debug, info};
|
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 BRAILLE: &[&str] = &["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"];
|
||||||
const ETA_REFRESH_MS: u64 = 500;
|
const ETA_REFRESH_MS: u64 = 500;
|
||||||
const ETA_MIN_ELAPSED_MS: u64 = 1000;
|
const ETA_MIN_ELAPSED_MS: u64 = 1000;
|
||||||
|
|||||||
Reference in New Issue
Block a user