From abc51c2addb44d8d5823a94c59f7f82b2b8dd855 Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Thu, 20 Aug 2026 19:45:15 +0200 Subject: [PATCH] Extract dereplication logic into obikderep crate and centralize paths Move partition dereplication logic into a dedicated `obikderep` crate that implements a two-phase hash-split/merge strategy with Rayon for parallel processing. Centralize superkmer file path construction in `obilayeredmap` and update dependent crates to use the new pipeline and shared path helpers. Adjust test suites to explicitly invoke the dereplication step. --- .../implementation/partition_layer_cache.md | 297 +++++++++++++++++- src/Cargo.lock | 21 ++ src/Cargo.toml | 2 +- src/obikderep/Cargo.toml | 21 ++ .../src}/dereplicate.rs | 32 +- src/obikderep/src/lib.rs | 110 +++++++ src/obikindex/src/index_layer.rs | 4 +- src/obikmer/Cargo.toml | 1 + src/obikmer/src/cmd/index/mod.rs | 19 +- src/obikpartitionner/Cargo.toml | 1 + src/obikpartitionner/src/partition/mod.rs | 14 +- src/obikpartitionner/src/partition/router.rs | 60 +--- src/obikpartitionner/src/partition/tests.rs | 5 +- src/obikphylo/Cargo.toml | 3 +- src/obikphylo/src/siblings/tests.rs | 6 +- src/obilayeredmap/src/layer.rs | 30 +- src/obilayeredmap/src/lib.rs | 5 +- 17 files changed, 534 insertions(+), 97 deletions(-) create mode 100644 src/obikderep/Cargo.toml rename src/{obikpartitionner/src/partition => obikderep/src}/dereplicate.rs (80%) create mode 100644 src/obikderep/src/lib.rs diff --git a/DevDocMD/implementation/partition_layer_cache.md b/DevDocMD/implementation/partition_layer_cache.md index 90f72b36..5ee7416a 100644 --- a/DevDocMD/implementation/partition_layer_cache.md +++ b/DevDocMD/implementation/partition_layer_cache.md @@ -8,8 +8,18 @@ state, added (panics on every read method). (2a) done — the 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 +the gap `Layer::Empty` was built for in (1b) — see "(3) done" below. (4) +done — dereplication split out into its own crate, `obikderep`, first step +of an incremental "one algorithm at a time" split of `obikpartitionner`'s +remaining bundle (`count_kmer`/`build_layers` not yet moved) — see "(4) +done" below. **(5) — full design agreed, not yet implemented** — +`KmerPartition` was found to be wired into nothing (`KmerIndex` never +calls it; every path is still computed via free functions), and the fix +turned out to be bigger than `KmerPartition` alone: `Layer`'s own +constructors don't self-name either. Full redesign of both, agreed in +detail, session ended (budget) before implementation — see "(5) design +agreed" below; **read it before touching `KmerPartition`/`Layer` +signatures**, the shape is fully specified. 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 @@ -447,6 +457,289 @@ Full workspace suite green (`cargo check --workspace --all-targets` + `index_layer.rs` fix above — the smoke test is what actually caught the regression the suite missed. +## (4) done (2026-08-20): `obikderep` — dereplication split out of `obikpartitionner`, one algorithm at a time + +Follow-on question after (3): the indexing pipeline has 4 stages (scatter, +dereplicate, count_kmer, index-build — see the CLI's own `Reporter` output, +one line per stage), but `obikpartitionner` — a name that says +*partitioning* — owned three of them (routing, dereplication, counting). +Challenged directly, same as (3)'s dependency-direction question: a crate +should hold what its name says, not accumulate unrelated stages just +because they happened to land there first. Two ways to fix it — one crate +renamed to hold all remaining stages, or one crate per stage — decided in +favour of the latter, explicitly **incremental**: build the *second* algo +crate first (`obikderep`, dereplication only), only then look at what it +and `PartitionRouter` actually have in common, and factor a shared +`Algorithm` trait (future `obikalgorithm` crate) from that real overlap — +not guessed at from a single example. `count_kmer` and `build_layers` +(currently `KmerIndex` inherent methods — itself flagged as inconsistent +with "`KmerIndex` is a data structure, not a compute structure") are left +alone this round, on purpose — one stage moves at a time. + +**`obikderep`** (new crate): `Dereplicator<'a> { index: &'a KmerIndex, n_partitions, level }` +— `new(index: &KmerIndex)` (shared borrow, not `&mut`: dereplication never +writes index metadata), no setters yet (nothing to configure), `run(on_progress)` +does the two-phase split+merge dereplication in parallel across partitions, +ported unchanged from `PartitionRouter::dereplicate` (moved wholesale: +`optimal_buckets`/`dereplicate_partition`/`load_bucket`/`flush_map`/ +`remove_skmer_file`, now private to this crate in `dereplicate.rs`). +`obikpartitionner::PartitionRouter::dereplicate` is gone; `count_kmer` +stays. + +**A real signature difference from `PartitionRouter::run`, not an +inconsistency**: `Dereplicator::run` takes `Option`, not `FnMut`. `PartitionRouter::run`'s callback is invoked from one +sequential loop (`FnMut` is fine); `Dereplicator::run`'s work is +`rayon::par_iter`, so the callback can be invoked concurrently from +multiple worker threads — same reason `obisys::TracedBar`'s own methods +take `&self`, not `&mut self`. Progress position is tracked with an +`AtomicU64`, incremented from inside the parallel closure so each +completed partition reports immediately — collecting all results first and +reporting after (the first draft of this) would have delivered every tick +in one burst at the very end, defeating the point of a live progress bar. +`total: Some(n_partitions)` (known up front, unlike scatter's bases count) +— `cmd/index/mod.rs` renders a real `progress_bar`, not a spinner, driven +by the callback exactly like scatter's spinner is. + +**A second, pre-existing instance of the exact bug (3) fixed, caught +before it shipped**: `dereplicated.skmer.zst` was hand-built as a string +literal independently in three places — `obikpartitionner`'s +`dereplicate.rs`/`count.rs` *and* `obikindex`'s `index_layer.rs` (a literal +that already predated this session, never caught until now). Splitting +dereplication into its own crate turns this from "two places, still +matching by luck" into "three independent crates that must agree on a +filename with no shared dependency forcing them to" — no longer +deferrable. Fixed by adding `obilayeredmap::{raw_superkmers_path, +dereplicated_superkmers_path}` (free functions, `layer_dir: &Path -> +PathBuf`, mirroring `layer_dir` itself) — the filename lives in one place, +in the Layer-tier crate every consumer here already depends on +(`obikpartitionner`, `obikderep`, `obikindex` all reach it without a new +edge), and no external crate ever sees the literal `"skmer.zst"` again. +This reverses (3)'s own earlier call to keep `SK_EXT` private to +`obikpartitionner` — that call assumed a single owner; a second owner +appearing (`obikderep`) removed the assumption it rested on, so the +decision changed with it, not out of inconsistency. + +Every `count_kmer` call site that used to run after `router.dereplicate()` +on the same `PartitionRouter` now runs after a separate +`Dereplicator::new(&idx).run(...)` call, on a freshly-constructed +`PartitionRouter` — `PartitionRouter` no longer offers a combined +"dereplicate then count" path. Updated at all three call sites that had +one: `cmd/index/mod.rs`, `obikphylo`'s test harness, `obikpartitionner`'s +own tests. + +Full workspace suite green (`cargo check --workspace --all-targets` + +`cargo test --workspace`, exit code 0), plus an end-to-end CLI smoke test +against real FASTA data (scatter → dereplicate → count → index-build → +query, same numbers as (3)'s smoke test: 870 kmers) — required this time +too, per (3)'s own lesson: no test in the suite exercises `obikmer index`'s +real file-driven path. + +Still not done: `count_kmer`/`build_layers` staying where they are, the +`obikalgorithm` shared-trait extraction (deliberately deferred until a +third data point exists), and everything already listed under (2b). + +## (5) design agreed, not yet implemented (2026-08-20): `KmerPartition` rewritten, `Layer` gains self-naming, a future cache crate over `KmerIndex` + +Session ended (out of budget) before any of this was coded. Everything +below is a **fully specified plan**, agreed sentence by sentence with the +user — not a sketch to re-derive, not a proposal to re-litigate. Implement +it as written; if something here turns out to be wrong once coded, fix it +and update this section, don't restart the design conversation. + +### How this was found + +Direct question from the user: "tu as bien créé une structure +`KmerPartition` ?" — yes (2a), but investigating exposed that it is +**wired into nothing**. `KmerIndex` has no `partition(i)` method at all; +`partition_dir`/`index_dir`/`layer_dir` still call `obikpartition:: +partition_dir`/`index_dir` and `obilayeredmap::layer_dir` as bare free +functions directly, never touching a `KmerPartition`/`Layer` object to get +there. The end result on disk is identical (same paths), which is exactly +why no test caught it — but the *responsibility* is in the wrong place: +one function (on `KmerIndex`) knows the whole three-tier naming +convention, instead of each tier asking the one below it for its own +path. User's framing, verbatim, now saved as [[feedback_no_spaghetti_petits_pois]]: +"spaghetti" (logic untraceable, split across too many unrelated crates) +and "petits pois" (small bits of naming logic dispersed with no owning +object) are **strictly forbidden** — this was a live example of both. + +Pushed further, twice: +1. First correction: `Layer::create(&obilayeredmap::layer_dir(&dir, 0))` — + still a free-function call from *outside* `Layer` to compute where it + should live. "Le layer n'est pas con, c'est lui qui dit où est-ce qu'il + doit être sauvé" (the layer isn't stupid, it says itself where it + should be saved). +2. Second correction, the general principle: **"une partition est juste + identifiée par un numéro, tout se calcule à partir du numéro, et un + layer est identifié à partir d'un numéro et tout se calcule à partir de + ce numéro."** Concretely: each object stores its own local identifying + number *plus* its immediate parent's path (captured once, at + construction) — never a path handed in again later by a caller, and + never a free function outside the object that can compute that path + independently. Explicitly rejected along the way: making users pass + "the partition's path that contains the layer" to open a layer — the + parent path is captured once, at the child's construction, not + re-supplied at every call. + +### The agreed shape + +**`Layer`** (`obilayeredmap`) — identified by `l` + its parent partition's +directory, both captured at construction, never received again: + +```rust +pub enum Layer { + Empty { partition_dir: PathBuf, l: usize }, // pure identification, no disk I/O + Count(TypedLayer), + Presence(TypedLayer), +} + +impl Layer { + pub fn at(partition_dir: &Path, l: usize) -> Self; // identify only + fn dir(&self) -> PathBuf; // private — layer_dir() no longer a public free function, folded in here + pub fn create(self) -> io::Result; // creates the directory if needed; no path parameter anymore + pub fn open(self, mode: &IndexMode, with_counts: bool) -> OLMResult; // no path parameter anymore + // mphf_path()/unitigs_path()/evidence_path()/fingerprint_path()/counts_dir()/presence_dir() + // unchanged in spirit, implemented via self.dir() instead of a stored `dir` field read directly +} +``` + +Note this **replaces** `Layer::Empty { dir: PathBuf }` from (1b) — `dir` +becomes a computed value (`partition_dir.join(format!("layer_{l}"))`), not +a stored field. `obilayeredmap::layer_dir`/`raw_superkmers_path`/ +`dereplicated_superkmers_path` (currently public free functions, +introduced in (3)/(4)) stop being called from outside `obilayeredmap` +entirely once this lands — they were the right fix for their moment (a +second crate, `obikderep`, needed to agree on a filename with no owner), +but the *real* fix, now visible with a third data point, is that `Layer` +itself should be the only thing anyone asks. + +**`KmerPartition`** (`obikpartition`) — same principle, one tier up: + +```rust +pub struct KmerPartition { + index_root: PathBuf, // the parent KmerIndex's root, captured once + i: usize, +} + +impl KmerPartition { + pub fn new(index_root: PathBuf, i: usize) -> Self; // identify only, no disk I/O + pub fn create(index_root: PathBuf, i: usize) -> io::Result; // creates this partition's directory + an empty layer 0 (a partition is never born without one — that knowledge lives here, not in whoever calls create) + pub fn partition_dir(&self) -> PathBuf; // part_{i:05} + pub fn index_dir(&self) -> PathBuf; // part_{i:05}/index + pub fn layer(&self, l: usize) -> Layer; // Layer::at(&self.index_dir(), l) — caller never touches a path + pub fn meta(&self) -> SKResult; // n_layers + mode; must absorb the recovery-on-missing-file logic + // currently private in obikindex::common::load_meta (obikpartition + // can't depend on obikindex to reuse it — this logic moves down) + pub fn n_layers(&self) -> SKResult; // meta()?.n_layers + pub fn mode(&self) -> SKResult; // meta()?.mode — "exact/approximatif" + pub fn is_filled(&self) -> bool; // does this partition's directory exist at all + pub fn n_kmers(&self) -> io::Result; // LayerMeta::load(&self.layer(0).dir()).n — reads layer 0's count as + // a representative figure, same "read the first one" trick + // n_layers_per_partition() already uses at the KmerIndex level +} +``` + +This **replaces** (2a)'s `KmerPartition { layers: Vec }` entirely +— no eagerly-opened `Vec`, no `find()` (both belong to the future +cache, see below, which is the thing that actually holds opened layers +alive across many lookups). (2a)'s version is safe to delete outright: it +was never wired into anything (confirmed above), so nothing depends on +its current shape. New dependencies needed: `obikpartition` gains +`obiskio` (for `SKResult`) and `obicompactvec` (for `LayerMeta`). + +Deliberately **not built this round**: cross-level consistency checks +("verify everything below me is in the same state") — a real idea, raised +by the user, but nothing concrete needs it yet; building it speculatively +would be exactly the premature-abstraction pattern this project avoids. + +**`KmerIndex`** (`obikindex`) — becomes the sole entry point: + +```rust +pub fn partition(&self, i: usize) -> KmerPartition; // KmerPartition::new(self.root_path.clone(), i) +``` + +`partition_dir(i)`/`index_dir(i)`/`layer_dir(i, l)` **stay** as public +methods (≈30 existing call sites across `obikindex`/`obikphylo` — see (3)'s +option A, applied identically here) but become pure delegations: +`self.partition(i).partition_dir()`, `self.partition(i).index_dir()`, +`self.partition(i).layer(l).dir()` (needs `Layer::dir()` to be visible +enough for this — likely `pub(crate)` in `obilayeredmap` plus a thin +public wrapper, or a public accessor on `Layer` itself; not fully nailed +down, decide while implementing). No caller outside `obikindex` changes. + +### Known blast radius (why this wasn't done in the same session) + +- **14 files** call `Layer::open`/`Layer::create` directly today + (`obikphylo/siblings/{cache,build,family_scan,tests}.rs`, + `obikpartitionner/partition/router.rs`, `obikpartition/src/lib.rs`, + `obikindex/{rebuild_layer,dump_layer,index,query_layer}.rs`, + `obilayeredmap/{mphf_layer,layer,map,content_layer}.rs`) — every one + loses its path parameter and gains a `(partition_dir, l)` or an + already-identified `Layer` to call `.create()`/`.open()` on instead. +- **≈30 files** call `KmerIndex::partition_dir`/`index_dir`/`layer_dir` — + unaffected in their own code (same public signatures), but worth + re-checking once (5) lands that none of them were relying on the old + free-function-based implementation in a way the new delegation breaks. +- `obikpartitionner::PartitionRouter::ensure_writer` and `obikderep`'s + `run` both currently call `obilayeredmap::{layer_dir, raw_superkmers_path, + dereplicated_superkmers_path}` directly (from (3)/(4)) — both need to + switch to going through `index.partition(i).layer(0)` instead. + +### Also agreed, separately: `PartitionRouter::new` never needed `&mut KmerIndex` + +Verified by reading the code: every call `PartitionRouter` makes on +`index` is `&self` (`index.kmer_size()`, `index.index_dir(i)`). The `&mut` +in its current signature (from (3)) was inherited from the original +"the router writes to the partitions" reasoning, never actually required +by any method call. This is *exactly* what caused every `drop(router)` +workaround needed throughout (3)/(4) (`cmd/index/mod.rs` ×2, `obikphylo`'s +test harness, `obikpartitionner`'s own tests) — `PartitionRouter` holds a +`Drop` impl, which extends a `&mut` borrow to the end of its scope even +past its last real use. **Fix alongside (5)**: change +`PartitionRouter::new(index: &'a mut KmerIndex)` to `&'a KmerIndex`, and +remove the now-unnecessary `drop(router)` calls at all four sites. + +### Also discussed: a future cache crate, not part of (5), not `obikalgorithm` either + +Separate idea, explicitly **not** part of this design and **not** started: +a new crate whose only job is to cache open `KmerPartition`s (and their +opened `Layer`s) across one run — replacing both `obikphylo::siblings:: +cache::PartitionCache` (today, sibling-specific, holds `Vec>`) +and `obikindex::query_layer::QueryLayer` (today, uncached, bypasses +`Layer` entirely) — the two consumers (2b) already identified as each +reinventing a fragment of the same thing. + +User's framing: this is **not** a third `obikalgorithm` data point — an +algorithm has a `new → run → done` shape; a cache has a fundamentally +different one (open, stay alive for a whole run, serve lookups, maybe +evict) — "on crée un cache sur un index, ça consomme un index." Two +distinct crate *roles* in this ecosystem (data crates: `obikpartition`/ +`obilayeredmap`; algorithm crates: `obikpartitionner`/`obikderep`/future +`obikalgorithm` implementors; and now a cache/service crate), not one +unified shape to force everything into. + +Depends on (5) being done first: the cache crate's whole job is holding +`Vec`/opened `Layer`s alive, built via `KmerIndex:: +partition(i)` as its factory — nothing to build it on top of until (5) +lands. Still open once (5) is done: eviction policy vs. holding everything +open for the process lifetime (the never-measured mmap/VM-mapping-count +question from earlier in this doc), and whether it lives in `obikpartition` +itself or a new crate. + +### Order of remaining work, as currently understood + +1. **(5)** — `Layer`/`KmerPartition`/`KmerIndex` rewrite described above, + plus the `PartitionRouter` `&mut` → `&` fix (same root cause, same + session, do together). +2. The future cache crate (name not chosen), consuming `KmerIndex:: + partition(i)` — unblocks migrating `PartitionCache`/`QueryLayer` (2b). +3. `obikalgorithm` — still deliberately waiting for a third `run()`-shaped + data point (`count_kmer` or `build_layers` migrating out of + `KmerIndex`/`PartitionRouter`) before extracting a shared trait; two + examples were judged not enough to be sure of the shape (`Fn+Sync` vs + `FnMut` callback bound already diverged between the two that exist). + ## The problem Reading a layer's data (MPHF + matrix) is not free: `MphfLayer::open` mmaps diff --git a/src/Cargo.lock b/src/Cargo.lock index c68b5468..3b836687 100644 --- a/src/Cargo.lock +++ b/src/Cargo.lock @@ -1515,6 +1515,24 @@ dependencies = [ "xxhash-rust", ] +[[package]] +name = "obikderep" +version = "0.1.0" +dependencies = [ + "niffler", + "obikindex", + "obikrope", + "obikseq", + "obilayeredmap", + "obiskbuilder", + "obiskio", + "obisys", + "rayon", + "sysinfo", + "tempfile", + "tracing", +] + [[package]] name = "obikentropy" version = "0.1.0" @@ -1565,6 +1583,7 @@ dependencies = [ "kodama", "obidebruinj", "obifastwrite", + "obikderep", "obikindex", "obikpartitionner", "obikphylo", @@ -1607,6 +1626,7 @@ dependencies = [ "memmap2", "niffler", "obicompactvec", + "obikderep", "obikindex", "obikrope", "obikseq", @@ -1632,6 +1652,7 @@ dependencies = [ "memmap2", "ndarray", "obicompactvec", + "obikderep", "obikindex", "obikpartitionner", "obikseq", diff --git a/src/Cargo.toml b/src/Cargo.toml index 3ead8c89..6dd93601 100644 --- a/src/Cargo.toml +++ b/src/Cargo.toml @@ -1,5 +1,5 @@ [workspace] resolver = "3" -members = ["obikseq", "obiread", "obiskbuilder", "obifastwrite", "obikmer","obikrope","obipipeline", "obikpartitionner","obiskio","obidebruinj","obilayeredmap", "obicompactvec", "obisys", "obikindex", "obitaxonomy", "obikentropy", "obikphylo", "obikpartition"] +members = ["obikseq", "obiread", "obiskbuilder", "obifastwrite", "obikmer","obikrope","obipipeline", "obikpartitionner","obiskio","obidebruinj","obilayeredmap", "obicompactvec", "obisys", "obikindex", "obitaxonomy", "obikentropy", "obikphylo", "obikpartition", "obikderep"] [profile.release] debug = 1 diff --git a/src/obikderep/Cargo.toml b/src/obikderep/Cargo.toml new file mode 100644 index 00000000..40a5e7a9 --- /dev/null +++ b/src/obikderep/Cargo.toml @@ -0,0 +1,21 @@ +[package] +name = "obikderep" +version = "0.1.0" +edition = "2024" + +[dependencies] +niffler = "3.0.0" +obikseq = { path = "../obikseq" } +obikindex = { path = "../obikindex" } +obilayeredmap = { path = "../obilayeredmap" } +obiskio = { path = "../obiskio" } +obisys = { path = "../obisys" } +rayon = "1" +sysinfo = "0.39" +tracing = "0.1.44" + +[dev-dependencies] +tempfile = "3" +obikseq = { path = "../obikseq", features = ["test-utils"] } +obikrope = { path = "../obikrope" } +obiskbuilder = { path = "../obiskbuilder" } diff --git a/src/obikpartitionner/src/partition/dereplicate.rs b/src/obikderep/src/dereplicate.rs similarity index 80% rename from src/obikpartitionner/src/partition/dereplicate.rs rename to src/obikderep/src/dereplicate.rs index 912e6424..de02ec29 100644 --- a/src/obikpartitionner/src/partition/dereplicate.rs +++ b/src/obikderep/src/dereplicate.rs @@ -1,3 +1,7 @@ +//! Per-partition dereplication mechanics — private to this crate. +//! [`crate::Dereplicator`] is the public entry point; this module is the +//! two-phase split+merge algorithm it runs once per partition. + use std::collections::HashMap; use std::fs; use std::io; @@ -5,16 +9,21 @@ use std::path::{Path, PathBuf}; use tracing::debug; -use niffler::Level; use niffler::send::compression::Format; -use obikseq::Sequence; +use niffler::Level; use obikseq::superkmer::SuperKmer; +use obikseq::Sequence; +use obilayeredmap::{dereplicated_superkmers_path, raw_superkmers_path}; use obiskio::{SKFileMeta, SKFileReader, SKFileWriter, SKResult}; -use super::SK_EXT; +/// Scratch-file extension for this algorithm's own intermediate split +/// buckets — never read by anything outside [`dereplicate_partition`], +/// unlike `raw`/`dereplicated` (see `obilayeredmap::{raw_superkmers_path, +/// dereplicated_superkmers_path}`, the actual cross-crate contract). +const TEMP_EXT: &str = "skmer.zst"; -/// Estimate the number of in-memory buckets needed to deduplicate the partition -/// file at `raw_path` given `available_bytes` of free RAM. +/// Estimate the number of in-memory buckets needed to deduplicate the +/// partition file at `raw_path` given `available_bytes` of free RAM. /// /// Memory per HashMap entry: /// key Box (1 + avg_seq_bytes) + SuperKmer header (4 B) + avg seq bytes + u64 count (8 B), @@ -22,7 +31,7 @@ use super::SK_EXT; /// /// Returns 1 if the partition fits comfortably in memory (no split needed). /// Always returns a power of two. -pub(super) fn optimal_buckets(raw_path: &Path, available_bytes: u64) -> usize { +pub(crate) fn optimal_buckets(raw_path: &Path, available_bytes: u64) -> usize { // Use 60 % of available RAM to leave headroom for the rest of the process. let budget = (available_bytes as f64 * 0.60) as u64; @@ -62,14 +71,15 @@ fn remove_skmer_file(path: &Path) -> SKResult<()> { /// Maximum value that fits in the 24-bit COUNT field of a SuperKmer header. const MAX_SK_COUNT: u64 = (1 << 24) - 1; -/// Deduplicate one partition directory in place (two-phase split + merge). -pub(super) fn dereplicate_partition(dir: &Path, level: Level, n_temp: usize) -> SKResult<()> { - let raw_path = dir.join(format!("raw.{SK_EXT}")); +/// Deduplicate one partition's layer-0 directory in place (two-phase split +/// + merge): `raw_superkmers_path(dir)` -> `dereplicated_superkmers_path(dir)`. +pub(crate) fn dereplicate_partition(dir: &Path, level: Level, n_temp: usize) -> SKResult<()> { + let raw_path = raw_superkmers_path(dir); if !raw_path.exists() { return Ok(()); } - let out_path = dir.join(format!("dereplicated.{SK_EXT}")); + let out_path = dereplicated_superkmers_path(dir); let mut writer = SKFileWriter::create_with(&out_path, Format::Zstd, level)?; if n_temp == 1 { @@ -81,7 +91,7 @@ pub(super) fn dereplicate_partition(dir: &Path, level: Level, n_temp: usize) -> // ── Phase 1: split raw file into temp buckets ───────────────────────── let temp_mask = (n_temp as u64) - 1; let temp_paths: Vec = (0..n_temp) - .map(|j| dir.join(format!("temp_{j:04}.{SK_EXT}"))) + .map(|j| dir.join(format!("temp_{j:04}.{TEMP_EXT}"))) .collect(); { diff --git a/src/obikderep/src/lib.rs b/src/obikderep/src/lib.rs new file mode 100644 index 00000000..080cfcea --- /dev/null +++ b/src/obikderep/src/lib.rs @@ -0,0 +1,110 @@ +//! Superkmer dereplication — the second stage of the indexing pipeline, +//! after `obikpartitionner::PartitionRouter::run` (scatter) has written +//! each partition's raw superkmer file, before +//! `obikpartitionner::PartitionRouter::count_kmer` (counting) reads the +//! result. One algorithm, one crate — see +//! `DevDocMD/implementation/partition_layer_cache.md`'s "on avance pas à +//! pas" note: `obikpartitionner` used to also own dereplication and +//! counting; this crate is step one of splitting that bundle apart, +//! deliberately one algorithm at a time rather than all at once, so a +//! shared `Algorithm` pattern can be factored out later from real +//! examples instead of guessed at up front. + +mod dereplicate; + +use std::sync::atomic::{AtomicU64, Ordering}; + +use niffler::Level; +use obikindex::KmerIndex; +use obiskio::SKResult; +use obisys::Progress; +use rayon::prelude::*; +use sysinfo::System; + +use dereplicate::{dereplicate_partition, optimal_buckets}; + +/// Deduplicates every partition's raw superkmer file in place, replacing +/// each with a dereplicated one where identical canonical sequences are +/// merged and their counts summed. +/// +/// Two-phase construction, same shape as `PartitionRouter`: `new` (no disk +/// access), then `run`. Nothing to configure yet — no setters, unlike +/// `PartitionRouter`'s `level_max`/`theta`/etc. — added if a real need +/// shows up, not speculatively. +pub struct Dereplicator<'a> { + index: &'a KmerIndex, + n_partitions: usize, + level: Level, +} + +impl<'a> Dereplicator<'a> { + pub fn new(index: &'a KmerIndex) -> Self { + Self { + index, + n_partitions: index.n_partitions(), + level: Level::One, + } + } + + /// Dereplicate every partition in parallel. + /// + /// Each partition file is processed in two phases to bound memory use: + /// + /// 1. **Split** — the raw file is scattered into `2^temp_bits` temporary + /// files routed by `hash(canonical_seq) & temp_mask`. Because duplicates + /// always share the same hash, they always land in the same temp file. + /// 2. **Merge** — each temp file is loaded fully into a `HashMap`, counts + /// are accumulated in `u64` (no 24-bit overflow risk), and the result is + /// appended to the partition's dereplicated file. + /// + /// If a merged count exceeds the 24-bit header limit, the sequence is + /// emitted as multiple records whose counts sum to the true total. + /// + /// `on_progress`, when set, is called once per completed partition, from + /// whichever rayon worker thread finished it — `Fn(...) + Sync`, not + /// `FnMut`, unlike `PartitionRouter::run`'s callback: that one is driven + /// from a single sequential loop, this one from a parallel `par_iter`, + /// so the callback itself must tolerate concurrent calls (same reason + /// `obisys::TracedBar`'s own methods take `&self`, not `&mut self`). + /// `total: Some(n_partitions)` — known up front here, unlike + /// `PartitionRouter::run`'s bases-processed count, so the caller can + /// render an actual progress bar rather than a spinner. This crate + /// never renders anything itself — see the module docs. + pub fn run(&self, on_progress: Option) -> SKResult<()> { + let level = self.level; + let sys = System::new_all(); + // available_memory() can return 0 on macOS when the compressor page count exceeds + // free+inactive+purgeable pages (sysinfo saturating_sub). Fall back to half of total. + let available = match sys.available_memory() { + 0 => sys.total_memory() / 2, + n => n, + }; + let n_threads = rayon::current_num_threads().max(1) as u64; + let available_per_thread = available / n_threads; + let done = AtomicU64::new(0); + + let results: Vec> = (0..self.n_partitions) + .into_par_iter() + .map(|i| { + let dir = self.index.layer_dir(i, 0); + let result = if dir.exists() { + let raw_path = obilayeredmap::raw_superkmers_path(&dir); + let n_buckets = optimal_buckets(&raw_path, available_per_thread); + dereplicate_partition(&dir, level, n_buckets) + } else { + Ok(()) + }; + if let Some(cb) = &on_progress { + let pos = done.fetch_add(1, Ordering::Relaxed) + 1; + cb(Progress { position: pos, total: Some(self.n_partitions as u64) }); + } + result + }) + .collect(); + + for r in results { + r?; + } + Ok(()) + } +} diff --git a/src/obikindex/src/index_layer.rs b/src/obikindex/src/index_layer.rs index b726ff78..f72cbfe0 100644 --- a/src/obikindex/src/index_layer.rs +++ b/src/obikindex/src/index_layer.rs @@ -43,7 +43,7 @@ impl KmerIndex { block_bits: u8, ) -> Result { let layer0_dir = self.layer_dir(i, 0); - let dedup_path = layer0_dir.join("dereplicated.skmer.zst"); + let dedup_path = obilayeredmap::dereplicated_superkmers_path(&layer0_dir); if !dedup_path.exists() { return Ok(0); } @@ -123,7 +123,7 @@ impl KmerIndex { /// Deletes `dereplicated.skmer.zst` (+ sidecar), `mphf1.bin`, `counts1.bin`. pub fn remove_build_artifacts(&self, i: usize) { let layer0_dir = self.layer_dir(i, 0); - let dedup = layer0_dir.join("dereplicated.skmer.zst"); + let dedup = obilayeredmap::dereplicated_superkmers_path(&layer0_dir); remove_if_exists(&SKFileMeta::sidecar_path(&dedup)); remove_if_exists(&dedup); remove_if_exists(&layer0_dir.join("mphf1.bin")); diff --git a/src/obikmer/Cargo.toml b/src/obikmer/Cargo.toml index bec904c2..6ed3c6b3 100644 --- a/src/obikmer/Cargo.toml +++ b/src/obikmer/Cargo.toml @@ -16,6 +16,7 @@ obidebruinj = { path = "../obidebruinj" } obipipeline = { path = "../obipipeline" } obikrope = { path = "../obikrope" } obikpartitionner = { path = "../obikpartitionner" } +obikderep = { path = "../obikderep" } obisys = { path = "../obisys" } obiskio = { path = "../obiskio" } obikindex = { path = "../obikindex", default-features = false } diff --git a/src/obikmer/src/cmd/index/mod.rs b/src/obikmer/src/cmd/index/mod.rs index b3bb94e1..21cbee65 100644 --- a/src/obikmer/src/cmd/index/mod.rs +++ b/src/obikmer/src/cmd/index/mod.rs @@ -2,6 +2,7 @@ use std::path::PathBuf; use std::time::Instant; use clap::Args; +use obikderep::Dereplicator; use obikindex::{validate_label, GenomeInfo, IndexConfig, IndexState, KmerIndex}; use obikpartitionner::PartitionRouter; use obilayeredmap::IndexMode; @@ -10,7 +11,7 @@ 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::{spinner, Progress, Reporter, Stage}; +use obisys::{progress_bar, spinner, Progress, Reporter, Stage}; use tracing::info; use crate::cli::{CommonArgs, block_size_to_bits, partitions_to_bits}; @@ -290,15 +291,19 @@ pub fn run(args: IndexArgs) { // ── Stage 2: dereplicate + count ───────────────────────────────────────── if idx.state() < IndexState::Counted { - let router = PartitionRouter::new(&mut idx); - let t = Stage::start("dereplicate"); - router.dereplicate().unwrap_or_else(|e| { - eprintln!("error: {e}"); - std::process::exit(1); - }); + let pb = progress_bar("dereplication", idx.n_partitions() as u64, "partitions"); + Dereplicator::new(&idx) + .run(Some(|_: Progress| pb.inc(1))) + .unwrap_or_else(|e| { + eprintln!("error: {e}"); + std::process::exit(1); + }); + pb.finish_and_clear(); rep.push(t.stop()); + let router = PartitionRouter::new(&mut idx); + let t = Stage::start("count_kmer"); let spectrum = router.count_kmer(args.keep_intermediate).unwrap_or_else(|e| { eprintln!("error: {e}"); diff --git a/src/obikpartitionner/Cargo.toml b/src/obikpartitionner/Cargo.toml index fffd0890..565efce5 100644 --- a/src/obikpartitionner/Cargo.toml +++ b/src/obikpartitionner/Cargo.toml @@ -7,6 +7,7 @@ edition = "2024" tempfile = "3" obikseq = { path = "../obikseq", features = ["test-utils"] } obikrope = { path = "../obikrope" } +obikderep = { path = "../obikderep" } [dependencies] niffler = "3.0.0" diff --git a/src/obikpartitionner/src/partition/mod.rs b/src/obikpartitionner/src/partition/mod.rs index a4cd8b36..489c2074 100644 --- a/src/obikpartitionner/src/partition/mod.rs +++ b/src/obikpartitionner/src/partition/mod.rs @@ -1,20 +1,18 @@ //! K-mer partitioning: routing super-kmers into per-partition, layer-0 -//! files, deduplicating them, and counting unique canonical k-mers. +//! files, and counting unique canonical k-mers. Dereplication itself moved +//! to `obikderep` (2026-08-20) — see +//! `DevDocMD/implementation/partition_layer_cache.md`: one algorithm, one +//! crate, split off one at a time rather than all at once. //! //! Submodules: [`router`] (`PartitionRouter`, `KmerSpectrum`, the -//! 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 +//! routing/counting lifecycle API — partition/layer path naming itself +//! lives on `obikindex::KmerIndex`, not here), [`count`] (unique-kmer //! enumeration, MPHF, abundance counting). mod count; -mod dereplicate; mod router; #[cfg(test)] mod tests; pub use router::{KmerSpectrum, PartitionRouter}; - -const SK_EXT: &str = "skmer.zst"; diff --git a/src/obikpartitionner/src/partition/router.rs b/src/obikpartitionner/src/partition/router.rs index 29272ea0..e8db0bff 100644 --- a/src/obikpartitionner/src/partition/router.rs +++ b/src/obikpartitionner/src/partition/router.rs @@ -25,8 +25,6 @@ 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::SK_EXT; pub struct KmerSpectrum { pub f0: u64, @@ -269,60 +267,6 @@ impl<'a> PartitionRouter<'a> { 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. - /// - /// Each partition file is processed in two phases to bound memory use: - /// - /// 1. **Split** — the raw file is scattered into `2^temp_bits` temporary - /// files routed by `hash(canonical_seq) & temp_mask`. Because duplicates - /// always share the same hash, they always land in the same temp file. - /// 2. **Merge** — each temp file is loaded fully into a `HashMap`, counts - /// are accumulated in `u64` (no 24-bit overflow risk), and the result is - /// appended to `dereplicated.{ext}`. - /// - /// If a merged count exceeds the 24-bit header limit, the sequence is - /// emitted as multiple records whose counts sum to the true total. - pub fn dereplicate(&self) -> SKResult<()> { - let level = self.level; - let sys = System::new_all(); - // available_memory() can return 0 on macOS when the compressor page count exceeds - // free+inactive+purgeable pages (sysinfo saturating_sub). Fall back to half of total. - let available = match sys.available_memory() { - 0 => sys.total_memory() / 2, - n => n, - }; - let n_threads = rayon::current_num_threads().max(1) as u64; - let available_per_thread = available / n_threads; - - let pb = progress_bar("dereplication", self.n_partitions as u64, "partitions"); - - let results: Vec> = (0..self.n_partitions) - .into_par_iter() - .map(|i| { - let dir = self.layer0_dir(i); - if !dir.exists() { - pb.inc(1); - return Ok(()); - } - let raw_path = dir.join(format!("raw.{SK_EXT}")); - let t = Instant::now(); - let n_buckets = optimal_buckets(&raw_path, available_per_thread); - let result = dereplicate_partition(&dir, level, n_buckets); - pb.set_message(format!("last {:.0}ms", t.elapsed().as_millis())); - pb.inc(1); - result - }) - .collect(); - - pb.finish_and_clear(); - for r in results { - r?; - } - Ok(()) - } - /// For each partition that has a `dereplicated.{ext}` file: /// 1. Enumerates all unique canonical kmers (two passes over the file). /// 2. Builds a provisional MPHF (FMPHGO) over those kmers. @@ -347,7 +291,7 @@ impl<'a> PartitionRouter<'a> { .into_par_iter() .map(|i| { let dir = self.layer0_dir(i); - let dedup_path = dir.join(format!("dereplicated.{SK_EXT}")); + let dedup_path = obilayeredmap::dereplicated_superkmers_path(&dir); if !dedup_path.exists() { pb.inc(1); return Ok(()); @@ -416,7 +360,7 @@ impl<'a> PartitionRouter<'a> { if self.writers[partition].is_none() { 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 file_path = obilayeredmap::raw_superkmers_path(&dir); let writer = SKFileWriter::create_with(file_path, Format::Zstd, self.level)?; self.writers[partition] = Some(writer); } diff --git a/src/obikpartitionner/src/partition/tests.rs b/src/obikpartitionner/src/partition/tests.rs index ed730aba..f4dcd646 100644 --- a/src/obikpartitionner/src/partition/tests.rs +++ b/src/obikpartitionner/src/partition/tests.rs @@ -1,6 +1,7 @@ use std::collections::HashMap; use std::fs; +use obikderep::Dereplicator; use obikindex::{IndexConfig, KmerIndex}; use obikrope::Rope; use obikseq::SuperKmer; @@ -64,11 +65,11 @@ fn pipeline_counts(seqs: &[&[u8]]) -> (u64, u64) { 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 + Dereplicator::new(&index).run(None::).unwrap(); let part_dir = index.layer_dir(0, 0); - let dedup_path = part_dir.join("dereplicated.skmer.zst"); + let dedup_path = obilayeredmap::dereplicated_superkmers_path(&part_dir); if !dedup_path.exists() { return (0, 0); } diff --git a/src/obikphylo/Cargo.toml b/src/obikphylo/Cargo.toml index 35b1f7f0..2e986665 100644 --- a/src/obikphylo/Cargo.toml +++ b/src/obikphylo/Cargo.toml @@ -20,6 +20,7 @@ rayon = "1" tracing = "0.1.44" [dev-dependencies] -obiread = { path = "../obiread" } +obiread = { path = "../obiread" } +obikderep = { path = "../obikderep" } tempfile = "3" tracing-subscriber = { version = "0.3", features = ["fmt", "env-filter"] } diff --git a/src/obikphylo/src/siblings/tests.rs b/src/obikphylo/src/siblings/tests.rs index c3937006..db332f64 100644 --- a/src/obikphylo/src/siblings/tests.rs +++ b/src/obikphylo/src/siblings/tests.rs @@ -6,6 +6,7 @@ use obilayeredmap::MphfLayer; use obisys::Reporter; use tempfile::tempdir; +use obikderep::Dereplicator; use obikindex::{GenomeInfo, IndexConfig, KmerIndex, MergeMode}; use obikpartitionner::PartitionRouter; @@ -76,7 +77,10 @@ fn build_single_genome_index(dir: &Path, label: &str, seq: &[u8]) -> KmerIndex { router.write_batch(batch).expect("write_batch"); } router.close().expect("close partition writers"); - router.dereplicate().expect("dereplicate"); + drop(router); // ends the borrow of `idx` early — `PartitionRouter`'s `Drop` impl would otherwise extend it to the end of scope + + Dereplicator::new(&idx).run(None::).expect("dereplicate"); + let router = PartitionRouter::new(&mut idx); 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 diff --git a/src/obilayeredmap/src/layer.rs b/src/obilayeredmap/src/layer.rs index 3c5428b6..9656931c 100644 --- a/src/obilayeredmap/src/layer.rs +++ b/src/obilayeredmap/src/layer.rs @@ -33,13 +33,37 @@ pub trait LayerData: Sized { /// /// `obilayeredmap` operates within a single partition's index root; it has /// no notion of "partition" at all. Turning a partition number into that -/// root is `obikpartitionner::KmerPartition::part_dir`'s job, one layer up -/// — callers here only ever name a layer *number*, never build the path -/// themselves. +/// root is `obikindex::KmerIndex::index_dir`'s job, one layer up — callers +/// here only ever name a layer *number*, never build the path themselves. pub fn layer_dir(root: &Path, i: usize) -> PathBuf { root.join(format!("layer_{i}")) } +/// Superkmer-file extension used by every pre-layer-construction artifact +/// below (`raw`/`dereplicated`) — an implementation detail of the SK file +/// format, never meant to leak as a literal string past this module. +const SK_EXT: &str = "skmer.zst"; + +/// Path of a layer's raw, not-yet-dereplicated superkmer file — written by +/// whichever algorithm routes superkmers into this layer (today: +/// `obikpartitionner::PartitionRouter`), read by whichever algorithm +/// dereplicates it (today: `obikderep::Dereplicator`). Naming this once +/// here, rather than in either algorithm crate, is what lets two +/// independent crates agree on the filename without depending on each +/// other — see `DevDocMD/implementation/partition_layer_cache.md`. +pub fn raw_superkmers_path(layer_dir: &Path) -> PathBuf { + layer_dir.join(format!("raw.{SK_EXT}")) +} + +/// Path of a layer's dereplicated superkmer file — written by +/// `obikderep::Dereplicator`, read by whichever algorithm counts kmer +/// abundances from it (today: `obikpartitionner::PartitionRouter:: +/// count_kmer`) and, later, by `obikindex::build_index_layer` to build the +/// real layer. +pub fn dereplicated_superkmers_path(layer_dir: &Path) -> PathBuf { + layer_dir.join(format!("dereplicated.{SK_EXT}")) +} + /// Opens layer `i`'s data only, skipping the MPHF — for callers that only /// need matrix-level operations (distance traits, column weights, group /// filters, sub-matrix extraction) and never look up a kmer for this layer. diff --git a/src/obilayeredmap/src/lib.rs b/src/obilayeredmap/src/lib.rs index 3df4fa9d..5ef4e24d 100644 --- a/src/obilayeredmap/src/lib.rs +++ b/src/obilayeredmap/src/lib.rs @@ -10,7 +10,10 @@ pub(crate) mod mphf_layer; pub use content_layer::Layer; pub use error::{OLMError, OLMResult}; -pub use layer::{layer_dir, open_data, HasLayerContent, HasStorageKind, Hit, LayerContent, LayerData, TypedLayer}; +pub use layer::{ + dereplicated_superkmers_path, layer_dir, open_data, raw_superkmers_path, HasLayerContent, + HasStorageKind, Hit, LayerContent, LayerData, TypedLayer, +}; pub use layered_store::LayeredStore; pub use map::LayeredMap; pub use meta::{IndexMode, PartitionMeta};