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.
This commit is contained in:
Eric Coissac
2026-08-20 22:30:03 +02:00
parent 616cb76af3
commit abc51c2add
17 changed files with 534 additions and 97 deletions
+21
View File
@@ -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",
+1 -1
View File
@@ -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
+21
View File
@@ -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" }
@@ -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<PathBuf> = (0..n_temp)
.map(|j| dir.join(format!("temp_{j:04}.{SK_EXT}")))
.map(|j| dir.join(format!("temp_{j:04}.{TEMP_EXT}")))
.collect();
{
+110
View File
@@ -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<impl Fn(Progress) + Sync>) -> 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<SKResult<()>> = (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(())
}
}
+2 -2
View File
@@ -43,7 +43,7 @@ impl KmerIndex {
block_bits: u8,
) -> Result<usize, SKError> {
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"));
+1
View File
@@ -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 }
+12 -7
View File
@@ -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}");
+1
View File
@@ -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"
+6 -8
View File
@@ -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";
+2 -58
View File
@@ -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<SKResult<()>> = (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);
}
+3 -2
View File
@@ -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::<fn(obisys::Progress)>).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);
}
+2 -1
View File
@@ -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"] }
+5 -1
View File
@@ -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::<fn(obisys::Progress)>).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
+27 -3
View File
@@ -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.
+4 -1
View File
@@ -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};