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
Generated
+6
-95
@@ -39,15 +39,6 @@ dependencies = [
|
||||
"memchr",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "aligned"
|
||||
version = "0.4.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ee4508988c62edf04abd8d92897fca0c2995d907ce1dfeaf369dac3716a40685"
|
||||
dependencies = [
|
||||
"as-slice",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "aligned-vec"
|
||||
version = "0.6.4"
|
||||
@@ -137,15 +128,6 @@ version = "0.7.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50"
|
||||
|
||||
[[package]]
|
||||
name = "as-slice"
|
||||
version = "0.2.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "516b6b4f0e40d50dcda9365d53964ec74560ad4284da2e7fc97122cd83174516"
|
||||
dependencies = [
|
||||
"stable_deref_trait",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "autocfg"
|
||||
version = "1.5.0"
|
||||
@@ -340,12 +322,6 @@ version = "1.0.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801"
|
||||
|
||||
[[package]]
|
||||
name = "cfg_aliases"
|
||||
version = "0.2.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724"
|
||||
|
||||
[[package]]
|
||||
name = "chacha20"
|
||||
version = "0.10.1"
|
||||
@@ -684,15 +660,6 @@ dependencies = [
|
||||
"memchr",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "cvt"
|
||||
version = "0.1.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d2ae9bf77fbf2d39ef573205d554d87e86c12f1994e9ea335b0651b9b278bcf1"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "debugid"
|
||||
version = "0.8.0"
|
||||
@@ -913,20 +880,6 @@ version = "0.2.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb"
|
||||
|
||||
[[package]]
|
||||
name = "fs_at"
|
||||
version = "0.2.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "14af6c9694ea25db25baa2a1788703b9e7c6648dcaeeebeb98f7561b5384c036"
|
||||
dependencies = [
|
||||
"aligned",
|
||||
"cfg-if",
|
||||
"cvt",
|
||||
"libc",
|
||||
"nix 0.29.0",
|
||||
"windows-sys 0.52.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "funty"
|
||||
version = "2.0.0"
|
||||
@@ -1422,7 +1375,7 @@ dependencies = [
|
||||
"combine",
|
||||
"libc",
|
||||
"mach2",
|
||||
"nix 0.26.4",
|
||||
"nix",
|
||||
"sysctl",
|
||||
"thiserror 1.0.69",
|
||||
"widestring",
|
||||
@@ -1484,27 +1437,6 @@ dependencies = [
|
||||
"pin-utils",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "nix"
|
||||
version = "0.29.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "71e2746dc3a24dd78b3cfcb7be93368c6de9963d30f43a6a73998a9cf4b17b46"
|
||||
dependencies = [
|
||||
"bitflags 2.11.1",
|
||||
"cfg-if",
|
||||
"cfg_aliases",
|
||||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "normpath"
|
||||
version = "1.5.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "bf23ab2b905654b4cb177e30b629937b3868311d4e1cba859f899c041046e69b"
|
||||
dependencies = [
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ntapi"
|
||||
version = "0.4.3"
|
||||
@@ -1606,7 +1538,7 @@ dependencies = [
|
||||
"obicompactvec",
|
||||
"obidebruinj",
|
||||
"obikentropy",
|
||||
"obikpartitionner",
|
||||
"obikpartition",
|
||||
"obikseq",
|
||||
"obilayeredmap",
|
||||
"obipipeline",
|
||||
@@ -1675,15 +1607,17 @@ dependencies = [
|
||||
"memmap2",
|
||||
"niffler",
|
||||
"obicompactvec",
|
||||
"obikindex",
|
||||
"obikrope",
|
||||
"obikseq",
|
||||
"obilayeredmap",
|
||||
"obipipeline",
|
||||
"obiread",
|
||||
"obiskbuilder",
|
||||
"obiskio",
|
||||
"obisys",
|
||||
"ptr_hash",
|
||||
"rayon",
|
||||
"remove_dir_all",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"sysinfo",
|
||||
@@ -2016,7 +1950,7 @@ dependencies = [
|
||||
"findshlibs",
|
||||
"libc",
|
||||
"log",
|
||||
"nix 0.26.4",
|
||||
"nix",
|
||||
"once_cell",
|
||||
"prost",
|
||||
"prost-build",
|
||||
@@ -2371,20 +2305,6 @@ version = "0.8.11"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4"
|
||||
|
||||
[[package]]
|
||||
name = "remove_dir_all"
|
||||
version = "1.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "808cc0b475acf76adf36f08ca49429b12aad9f678cb56143d5b3cb49b9a1dd08"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"cvt",
|
||||
"fs_at",
|
||||
"libc",
|
||||
"normpath",
|
||||
"windows-sys 0.59.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ring"
|
||||
version = "0.17.14"
|
||||
@@ -3312,15 +3232,6 @@ dependencies = [
|
||||
"windows-targets 0.52.6",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-sys"
|
||||
version = "0.59.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b"
|
||||
dependencies = [
|
||||
"windows-targets 0.52.6",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-sys"
|
||||
version = "0.61.2"
|
||||
|
||||
@@ -5,7 +5,7 @@ edition = "2024"
|
||||
|
||||
[dependencies]
|
||||
obikseq = { path = "../obikseq" }
|
||||
obikpartitionner = { path = "../obikpartitionner" }
|
||||
obikpartition = { path = "../obikpartition" }
|
||||
obitaxonomy = { path = "../obitaxonomy" }
|
||||
obiskio = { path = "../obiskio" }
|
||||
obisys = { path = "../obisys" }
|
||||
|
||||
+19
-41
@@ -2,7 +2,6 @@ use std::collections::BTreeMap;
|
||||
use std::fs;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use obikpartitionner::{KmerSpectrum, PartitionRouter};
|
||||
use obilayeredmap::meta::PartitionMeta;
|
||||
use obisys::{Reporter, Stage, progress_bar};
|
||||
use rayon::prelude::*;
|
||||
@@ -29,10 +28,8 @@ impl KmerIndex {
|
||||
path: P,
|
||||
config: IndexConfig,
|
||||
genome_info: Option<GenomeInfo>,
|
||||
force: bool,
|
||||
) -> OKIResult<Self> {
|
||||
let root_path = path.as_ref().to_owned();
|
||||
PartitionRouter::create(&root_path, config.n_bits, force)?;
|
||||
set_k(config.kmer_size);
|
||||
set_m(config.minimizer_size);
|
||||
let mut meta = IndexMeta::new(config);
|
||||
@@ -86,7 +83,6 @@ impl KmerIndex {
|
||||
let output = output.as_ref();
|
||||
fs::create_dir_all(output).map_err(OKIError::Io)?;
|
||||
meta.write(output).map_err(OKIError::Io)?;
|
||||
PartitionRouter::create(output, meta.config.n_bits, false)?;
|
||||
Ok(KmerIndex { root_path: output.to_owned(), meta: meta.clone() })
|
||||
}
|
||||
|
||||
@@ -160,16 +156,19 @@ impl KmerIndex {
|
||||
}
|
||||
|
||||
/// Path of partition `i`'s raw directory (`partitions/part_{i:05}`) —
|
||||
/// the on-disk naming convention `obikpartitionner::PartitionRouter`
|
||||
/// also writes to (raw/dereplicated superkmer files, `mphf1.bin`,
|
||||
/// `counts1.bin`), the single point of agreement between the two.
|
||||
/// delegates to `obikpartition`, the Partition tier's own naming
|
||||
/// primitive (mirrors `layer_dir` delegating to `obilayeredmap`).
|
||||
/// `obikpartitionner::PartitionRouter` reaches this same directory
|
||||
/// only indirectly, through this method (it depends on `KmerIndex`,
|
||||
/// not the other way around — see
|
||||
/// `DevDocMD/implementation/partition_layer_cache.md`).
|
||||
pub fn partition_dir(&self, i: usize) -> PathBuf {
|
||||
obikpartitionner::partition_dir(&self.root_path, i)
|
||||
obikpartition::partition_dir(&self.root_path, i)
|
||||
}
|
||||
|
||||
/// Path of partition `i`'s layered-index directory (`<partition>/index`).
|
||||
pub fn index_dir(&self, i: usize) -> PathBuf {
|
||||
self.partition_dir(i).join("index")
|
||||
obikpartition::index_dir(&self.root_path, i)
|
||||
}
|
||||
|
||||
/// Path of layer `l` within partition `i`'s layered index.
|
||||
@@ -223,39 +222,19 @@ impl KmerIndex {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Open a fresh [`PartitionRouter`] onto this index's partition layout
|
||||
/// — the write-side handle for `scatter`, or for `dereplicate_and_count`
|
||||
/// below. Transient: no state is kept in `KmerIndex` itself between
|
||||
/// calls, only on disk.
|
||||
pub fn partition_router(&self) -> OKIResult<PartitionRouter> {
|
||||
Ok(PartitionRouter::open(&self.root_path, self.meta.config.n_bits)?)
|
||||
}
|
||||
|
||||
/// Dereplicate all partitions then compute kmer counts.
|
||||
///
|
||||
/// Writes `spectrums/{label}.json` and touches `count.done` upon completion.
|
||||
/// Per-partition spectrum files are removed unless `keep_intermediate` is true.
|
||||
pub fn dereplicate_and_count(
|
||||
&self,
|
||||
keep_intermediate: bool,
|
||||
rep: &mut Reporter,
|
||||
) -> OKIResult<()> {
|
||||
let router = self.partition_router()?;
|
||||
|
||||
let t = Stage::start("dereplicate");
|
||||
router.dereplicate()?;
|
||||
rep.push(t.stop());
|
||||
|
||||
let t = Stage::start("count_kmer");
|
||||
let spectrum = router.count_kmer(keep_intermediate)?;
|
||||
rep.push(t.stop());
|
||||
|
||||
self.write_spectrum(&spectrum)?;
|
||||
/// Mark dereplicate+count as complete and write `count.done`.
|
||||
pub fn mark_counted(&self) -> OKIResult<()> {
|
||||
touch(&self.root_path.join(SENTINEL_COUNTED))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn write_spectrum(&self, sp: &KmerSpectrum) -> OKIResult<()> {
|
||||
/// Write `spectrums/{label}.json` from an already-computed kmer
|
||||
/// spectrum (`f0`/`f1`/abundance histogram). Takes plain values rather
|
||||
/// than `obikpartitionner::KmerSpectrum` — `KmerIndex` cannot depend on
|
||||
/// `obikpartitionner` (that dependency runs the other way, see
|
||||
/// `DevDocMD/implementation/partition_layer_cache.md`), and doesn't
|
||||
/// need to: this is the only field of that type it actually uses.
|
||||
pub fn write_spectrum(&self, f0: u64, f1: u64, counts: &BTreeMap<u32, u64>) -> OKIResult<()> {
|
||||
let label = self
|
||||
.meta
|
||||
.genomes
|
||||
@@ -265,15 +244,14 @@ impl KmerIndex {
|
||||
let spectrums_dir = self.root_path.join("spectrums");
|
||||
fs::create_dir_all(&spectrums_dir)?;
|
||||
let path = spectrums_dir.join(format!("{label}.json"));
|
||||
let spectrum_map: BTreeMap<String, u64> = sp
|
||||
.counts
|
||||
let spectrum_map: BTreeMap<String, u64> = counts
|
||||
.iter()
|
||||
.map(|(&c, &f)| (format!("{c:010}"), f))
|
||||
.collect();
|
||||
let f = fs::File::create(&path)?;
|
||||
serde_json::to_writer_pretty(
|
||||
f,
|
||||
&serde_json::json!({ "f0": sp.f0, "f1": sp.f1, "spectrum": spectrum_map }),
|
||||
&serde_json::json!({ "f0": f0, "f1": f1, "spectrum": spectrum_map }),
|
||||
)
|
||||
.map_err(OKIError::Json)?;
|
||||
Ok(())
|
||||
|
||||
@@ -42,14 +42,13 @@ impl KmerIndex {
|
||||
mode: &IndexMode,
|
||||
block_bits: u8,
|
||||
) -> Result<usize, SKError> {
|
||||
let partition_dir = self.partition_dir(i);
|
||||
let dedup_path = partition_dir.join("dereplicated.skmer.zst");
|
||||
let layer0_dir = self.layer_dir(i, 0);
|
||||
let dedup_path = layer0_dir.join("dereplicated.skmer.zst");
|
||||
if !dedup_path.exists() {
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
let layer_dir = self.layer_dir(i, 0);
|
||||
if layer_dir.join("mphf.bin").exists() {
|
||||
if layer0_dir.join("mphf.bin").exists() {
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
@@ -57,14 +56,14 @@ impl KmerIndex {
|
||||
let need_counts = filter_active || with_counts;
|
||||
|
||||
let mphf1_opt: Option<Mphf> = if need_counts {
|
||||
let p = partition_dir.join("mphf1.bin");
|
||||
let p = layer0_dir.join("mphf1.bin");
|
||||
p.exists().then(|| Mphf::load_full(&p).ok()).flatten()
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let counts1_opt: Option<PersistentCompactIntVec> = if need_counts {
|
||||
let p = partition_dir.join("counts1.bin");
|
||||
let p = layer0_dir.join("counts1.bin");
|
||||
p.exists()
|
||||
.then(|| PersistentCompactIntVec::open(&p).ok())
|
||||
.flatten()
|
||||
@@ -95,8 +94,8 @@ impl KmerIndex {
|
||||
|
||||
let n_kmers =
|
||||
if with_counts {
|
||||
let n = write_graph_as_unitigs(g, &layer_dir)?;
|
||||
TypedLayer::<PersistentCompactIntMatrix>::build(&layer_dir, block_bits, mode, |kmer| {
|
||||
let n = write_graph_as_unitigs(g, &layer0_dir)?;
|
||||
TypedLayer::<PersistentCompactIntMatrix>::build(&layer0_dir, block_bits, mode, |kmer| {
|
||||
match (&mphf1_opt, &counts1_opt) {
|
||||
(Some(mphf), Some(counts)) => counts.get(mphf.index(&kmer.raw())),
|
||||
_ => 1,
|
||||
@@ -105,10 +104,10 @@ impl KmerIndex {
|
||||
.map_err(|e| olm_to_sk(e, "layer build"))?;
|
||||
n
|
||||
} else {
|
||||
materialize_layer(g, &layer_dir, block_bits, mode)?
|
||||
materialize_layer(g, &layer0_dir, block_bits, mode)?
|
||||
};
|
||||
|
||||
let index_dir = layer_dir.parent().expect("layer_dir has a parent");
|
||||
let index_dir = layer0_dir.parent().expect("layer_dir has a parent");
|
||||
PartitionMeta {
|
||||
n_layers: 1,
|
||||
mode: mode.clone(),
|
||||
@@ -123,11 +122,11 @@ impl KmerIndex {
|
||||
///
|
||||
/// Deletes `dereplicated.skmer.zst` (+ sidecar), `mphf1.bin`, `counts1.bin`.
|
||||
pub fn remove_build_artifacts(&self, i: usize) {
|
||||
let partition_dir = self.partition_dir(i);
|
||||
let dedup = partition_dir.join("dereplicated.skmer.zst");
|
||||
let layer0_dir = self.layer_dir(i, 0);
|
||||
let dedup = layer0_dir.join("dereplicated.skmer.zst");
|
||||
remove_if_exists(&SKFileMeta::sidecar_path(&dedup));
|
||||
remove_if_exists(&dedup);
|
||||
remove_if_exists(&partition_dir.join("mphf1.bin"));
|
||||
remove_if_exists(&partition_dir.join("counts1.bin"));
|
||||
remove_if_exists(&layer0_dir.join("mphf1.bin"));
|
||||
remove_if_exists(&layer0_dir.join("counts1.bin"));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -53,7 +53,7 @@ fn query_partition_with_missing_index_dir_returns_default_stats() {
|
||||
evidence: obilayeredmap::IndexMode::Exact,
|
||||
block_bits: 0,
|
||||
};
|
||||
let index = KmerIndex::create(tmp.path().join("idx"), config, None, false).expect("create index");
|
||||
let index = KmerIndex::create(tmp.path().join("idx"), config, None).expect("create index");
|
||||
|
||||
let mut kmers: HashMap<CanonicalKmer, Vec<KmerDesc>> = HashMap::new();
|
||||
// Any well-formed canonical k-mer works here — the call must return
|
||||
@@ -81,7 +81,7 @@ fn query_partition_with_empty_kmers_is_a_noop() {
|
||||
evidence: obilayeredmap::IndexMode::Exact,
|
||||
block_bits: 0,
|
||||
};
|
||||
let index = KmerIndex::create(tmp.path().join("idx"), config, None, false).expect("create index");
|
||||
let index = KmerIndex::create(tmp.path().join("idx"), config, None).expect("create index");
|
||||
|
||||
let kmers: HashMap<CanonicalKmer, Vec<KmerDesc>> = HashMap::new();
|
||||
let stats = index
|
||||
|
||||
@@ -1,18 +1,19 @@
|
||||
use std::path::PathBuf;
|
||||
use std::time::Instant;
|
||||
|
||||
use clap::Args;
|
||||
use obikindex::{validate_label, GenomeInfo, IndexConfig, IndexState, KmerIndex};
|
||||
use obikpartitionner::PartitionRouter;
|
||||
use obilayeredmap::IndexMode;
|
||||
|
||||
fn parse_key_value(s: &str) -> Result<(String, String), String> {
|
||||
let pos = s.find('=').ok_or_else(|| format!("invalid key=value: no '=' in '{s}'"))?;
|
||||
Ok((s[..pos].to_string(), s[pos + 1..].to_string()))
|
||||
}
|
||||
use obisys::Reporter;
|
||||
use obisys::{spinner, Progress, Reporter, Stage};
|
||||
use tracing::info;
|
||||
|
||||
use crate::cli::{CommonArgs, block_size_to_bits, partitions_to_bits};
|
||||
use crate::steps::scatter;
|
||||
|
||||
#[derive(Args)]
|
||||
pub struct IndexArgs {
|
||||
@@ -225,7 +226,7 @@ pub fn run(args: IndexArgs) {
|
||||
}
|
||||
info
|
||||
});
|
||||
KmerIndex::create(&output, config, genome_info, args.force).unwrap_or_else(|e| {
|
||||
KmerIndex::create(&output, config, genome_info).unwrap_or_else(|e| {
|
||||
eprintln!("error creating index: {e}");
|
||||
std::process::exit(1);
|
||||
})
|
||||
@@ -234,17 +235,50 @@ pub fn run(args: IndexArgs) {
|
||||
|
||||
// ── Stage 1: scatter ─────────────────────────────────────────────────────
|
||||
if idx.state() < IndexState::Scattered {
|
||||
let k = idx.kmer_size();
|
||||
let level_max = args.common.level_max;
|
||||
let theta = args.common.theta;
|
||||
let n_workers = args.common.threads.max(1);
|
||||
|
||||
let max_open = args.common.effective_max_open();
|
||||
let mut router = idx.partition_router().unwrap_or_else(|e| {
|
||||
eprintln!("error opening partition router: {e}");
|
||||
std::process::exit(1);
|
||||
});
|
||||
scatter(&mut router, args.common.seqfile_paths(), k, level_max, theta, n_workers, max_open, &mut rep);
|
||||
|
||||
let t = Stage::start("scatter");
|
||||
let pb = spinner("scatter");
|
||||
let mut ema_rate: f64 = 0.0;
|
||||
let mut last_t = Instant::now();
|
||||
let mut last_bases: u64 = 0;
|
||||
const ALPHA: f64 = 0.15;
|
||||
|
||||
let mut router = PartitionRouter::new(&mut idx)
|
||||
.level_max(args.common.level_max)
|
||||
.theta(args.common.theta)
|
||||
.workers(n_workers)
|
||||
.max_open(max_open);
|
||||
|
||||
router
|
||||
.run(
|
||||
args.common.seqfile_paths(),
|
||||
Some(|p: Progress| {
|
||||
let now = Instant::now();
|
||||
let dt = now.duration_since(last_t).as_secs_f64();
|
||||
if dt > 0.0 {
|
||||
let instant = (p.position - last_bases) as f64 / dt;
|
||||
ema_rate = ALPHA * instant + (1.0 - ALPHA) * ema_rate;
|
||||
}
|
||||
last_t = now;
|
||||
last_bases = p.position;
|
||||
let bp = p.position as f64;
|
||||
let (count_str, rate_str) = if bp >= 1e9 {
|
||||
(format!("{:.2} Gbp", bp / 1e9), format!("{:.0} Mbp/s", ema_rate / 1e6))
|
||||
} else {
|
||||
(format!("{:.0} Mbp", bp / 1e6), format!("{:.0} Mbp/s", ema_rate / 1e6))
|
||||
};
|
||||
pb.set_message(format!("{count_str} {rate_str}"));
|
||||
}),
|
||||
)
|
||||
.unwrap_or_else(|e| {
|
||||
eprintln!("error: {e}");
|
||||
std::process::exit(1);
|
||||
});
|
||||
pb.finish_and_clear();
|
||||
rep.push(t.stop());
|
||||
drop(router); // ends the borrow of `idx` early — `PartitionRouter`'s `Drop` impl would otherwise extend it to the end of scope
|
||||
|
||||
idx.mark_scattered().unwrap_or_else(|e| {
|
||||
eprintln!("error marking scatter done: {e}");
|
||||
@@ -256,10 +290,31 @@ pub fn run(args: IndexArgs) {
|
||||
|
||||
// ── Stage 2: dereplicate + count ─────────────────────────────────────────
|
||||
if idx.state() < IndexState::Counted {
|
||||
idx.dereplicate_and_count(args.keep_intermediate, &mut rep).unwrap_or_else(|e| {
|
||||
let router = PartitionRouter::new(&mut idx);
|
||||
|
||||
let t = Stage::start("dereplicate");
|
||||
router.dereplicate().unwrap_or_else(|e| {
|
||||
eprintln!("error: {e}");
|
||||
std::process::exit(1);
|
||||
});
|
||||
rep.push(t.stop());
|
||||
|
||||
let t = Stage::start("count_kmer");
|
||||
let spectrum = router.count_kmer(args.keep_intermediate).unwrap_or_else(|e| {
|
||||
eprintln!("error: {e}");
|
||||
std::process::exit(1);
|
||||
});
|
||||
rep.push(t.stop());
|
||||
drop(router); // ends the borrow of `idx` early — `PartitionRouter`'s `Drop` impl would otherwise extend it to the end of scope
|
||||
|
||||
idx.write_spectrum(spectrum.f0, spectrum.f1, &spectrum.counts).unwrap_or_else(|e| {
|
||||
eprintln!("error: {e}");
|
||||
std::process::exit(1);
|
||||
});
|
||||
idx.mark_counted().unwrap_or_else(|e| {
|
||||
eprintln!("error marking count done: {e}");
|
||||
std::process::exit(1);
|
||||
});
|
||||
} else {
|
||||
info!("dereplicate+count already done, skipping");
|
||||
}
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
mod cli;
|
||||
mod cmd;
|
||||
mod steps;
|
||||
|
||||
use clap::{Parser, Subcommand};
|
||||
use tracing_subscriber::{EnvFilter, fmt};
|
||||
|
||||
@@ -1,3 +0,0 @@
|
||||
mod scatter;
|
||||
|
||||
pub use scatter::scatter;
|
||||
@@ -1,136 +0,0 @@
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
|
||||
use std::time::Instant;
|
||||
|
||||
use obikpartitionner::PartitionRouter;
|
||||
use obipipeline::{ThrottleGuard, Throttled, throttle};
|
||||
use obiread::NucPage;
|
||||
use obisys::spinner;
|
||||
use obisys::{Reporter, Stage};
|
||||
use tracing::info;
|
||||
|
||||
use crate::cli::PipelineData;
|
||||
|
||||
// ── Iterator that keeps the slot guard alive until the file is exhausted ──────
|
||||
|
||||
struct GuardedIter {
|
||||
inner: Box<dyn Iterator<Item = NucPage> + Send>,
|
||||
_guard: ThrottleGuard,
|
||||
flat_active: Arc<AtomicU32>,
|
||||
}
|
||||
|
||||
impl Iterator for GuardedIter {
|
||||
type Item = NucPage;
|
||||
fn next(&mut self) -> Option<NucPage> {
|
||||
self.inner.next()
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for GuardedIter {
|
||||
fn drop(&mut self) {
|
||||
self.flat_active.fetch_sub(1, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
// ── scatter ───────────────────────────────────────────────────────────────────
|
||||
|
||||
/// Run scatter: normalise → build superkmers → route to partition → close.
|
||||
/// Reports the "scatter" stage to `rep`.
|
||||
pub fn scatter(
|
||||
kp: &mut PartitionRouter,
|
||||
path_source: impl Iterator<Item = PathBuf> + Send + 'static,
|
||||
k: usize,
|
||||
level_max: usize,
|
||||
theta: f64,
|
||||
n_workers: usize,
|
||||
max_open: usize,
|
||||
rep: &mut Reporter,
|
||||
) {
|
||||
use obikseq::RoutableSuperKmer;
|
||||
|
||||
// Throttle in the source thread — never in a worker — to prevent deadlock.
|
||||
let throttled = throttle(path_source, max_open);
|
||||
|
||||
let file_count = Arc::new(AtomicU64::new(0));
|
||||
let flat_active = Arc::new(AtomicU32::new(0));
|
||||
let transform_active = Arc::new(AtomicU32::new(0));
|
||||
|
||||
let t = Stage::start("scatter");
|
||||
let pipe = obipipeline::make_pipe! {
|
||||
PipelineData : Throttled<PathBuf> => Vec<RoutableSuperKmer>,
|
||||
||? {
|
||||
let file_count = Arc::clone(&file_count);
|
||||
let flat_active = Arc::clone(&flat_active);
|
||||
let k = k;
|
||||
move |pw: Throttled<PathBuf>| {
|
||||
let path = pw.item;
|
||||
let guard = pw.guard;
|
||||
let n = file_count.fetch_add(1, Ordering::Relaxed) + 1;
|
||||
info!("indexing [{}]: {}", n, path.display());
|
||||
let path_str = path.to_str().unwrap_or("").to_owned();
|
||||
flat_active.fetch_add(1, Ordering::Relaxed);
|
||||
obiread::open_nuc_stream(&path_str, k)
|
||||
.map(|iter| GuardedIter { inner: iter, _guard: guard, flat_active: Arc::clone(&flat_active) })
|
||||
}
|
||||
} : Path => NucPage,
|
||||
| {
|
||||
let transform_active = Arc::clone(&transform_active);
|
||||
move |page| {
|
||||
transform_active.fetch_add(1, Ordering::Relaxed);
|
||||
let result = obiskbuilder::build_superkmers_page(page, k, level_max, theta);
|
||||
transform_active.fetch_sub(1, Ordering::Relaxed);
|
||||
result
|
||||
}
|
||||
} : NucPage => Batch,
|
||||
};
|
||||
|
||||
let pb = spinner("scatter");
|
||||
|
||||
let mut total_bases: u64 = 0;
|
||||
let mut ema_rate: f64 = 0.0;
|
||||
let mut last_t = Instant::now();
|
||||
let mut last_bases: u64 = 0;
|
||||
let kmer_overlap = (k - 1) as u64;
|
||||
const ALPHA: f64 = 0.15;
|
||||
|
||||
for batch in pipe.apply(throttled, n_workers, 1) {
|
||||
total_bases += batch
|
||||
.iter()
|
||||
.map(|sk| (sk.seql() as u64).saturating_sub(kmer_overlap))
|
||||
.sum::<u64>();
|
||||
let now = Instant::now();
|
||||
let dt = now.duration_since(last_t).as_secs_f64();
|
||||
if dt > 0.1 {
|
||||
let instant = (total_bases - last_bases) as f64 / dt;
|
||||
ema_rate = ALPHA * instant + (1.0 - ALPHA) * ema_rate;
|
||||
last_t = now;
|
||||
last_bases = total_bases;
|
||||
let bp = total_bases as f64;
|
||||
let (count_str, rate_str) = if bp >= 1e9 {
|
||||
(
|
||||
format!("{:.2} Gbp", bp / 1e9),
|
||||
format!("{:.0} Mbp/s", ema_rate / 1e6),
|
||||
)
|
||||
} else {
|
||||
(
|
||||
format!("{:.0} Mbp", bp / 1e6),
|
||||
format!("{:.0} Mbp/s", ema_rate / 1e6),
|
||||
)
|
||||
};
|
||||
let n_files = file_count.load(Ordering::Relaxed);
|
||||
let r = flat_active.load(Ordering::Relaxed);
|
||||
let c = transform_active.load(Ordering::Relaxed);
|
||||
pb.set_message(format!(
|
||||
"{count_str} {rate_str} {n_files} files [R:{r} C:{c}]"
|
||||
));
|
||||
}
|
||||
kp.write_batch(batch).unwrap_or_else(|e| {
|
||||
eprintln!("error: {e}");
|
||||
std::process::exit(1);
|
||||
});
|
||||
}
|
||||
pb.finish_and_clear();
|
||||
kp.close().expect("close error");
|
||||
rep.push(t.stop());
|
||||
}
|
||||
@@ -20,11 +20,32 @@
|
||||
//! entirely) both still reinvent a fragment of this. Migrating them is a
|
||||
//! separate, deferred step.
|
||||
|
||||
use std::path::Path;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use obikseq::CanonicalKmer;
|
||||
use obilayeredmap::{layer_dir, IndexMode, Layer, OLMResult};
|
||||
|
||||
/// Partition subdirectory name, under an index's root — the single source
|
||||
/// of truth for the on-disk `partitions/part_NNNNN` naming convention.
|
||||
/// Moved here from `obikpartitionner` (2026-08-20): naming primitives for
|
||||
/// the Partition tier belong with the other Partition-tier code, not with
|
||||
/// the superkmer-routing algorithm that happens to be their first
|
||||
/// consumer — see `DevDocMD/implementation/partition_layer_cache.md`.
|
||||
pub const PARTITIONS_SUBDIR: &str = "partitions";
|
||||
|
||||
/// Path of partition `i`'s directory under `root` — `<root>/partitions/part_NNNNN`,
|
||||
/// zero-padded to 5 digits.
|
||||
pub fn partition_dir(root: &Path, i: usize) -> PathBuf {
|
||||
root.join(PARTITIONS_SUBDIR).join(format!("part_{i:05}"))
|
||||
}
|
||||
|
||||
/// Path of partition `i`'s layered-index directory — `<partition_dir>/index`,
|
||||
/// the root every [`obilayeredmap::layer_dir`] call for this partition is
|
||||
/// relative to.
|
||||
pub fn index_dir(root: &Path, i: usize) -> PathBuf {
|
||||
partition_dir(root, i).join("index")
|
||||
}
|
||||
|
||||
/// One partition's open layers, in layer order (layer 0 first).
|
||||
pub struct KmerPartition {
|
||||
layers: Vec<Layer>,
|
||||
|
||||
@@ -6,15 +6,17 @@ edition = "2024"
|
||||
[dev-dependencies]
|
||||
tempfile = "3"
|
||||
obikseq = { path = "../obikseq", features = ["test-utils"] }
|
||||
obiread = { path = "../obiread" }
|
||||
obikrope = { path = "../obikrope" }
|
||||
|
||||
[dependencies]
|
||||
niffler = "3.0.0"
|
||||
remove_dir_all = "1.0"
|
||||
obikseq = { path = "../obikseq" }
|
||||
obiskbuilder = { path = "../obiskbuilder" }
|
||||
obiskio = { path = "../obiskio" }
|
||||
obikseq = { path = "../obikseq" }
|
||||
obikindex = { path = "../obikindex" }
|
||||
obilayeredmap = { path = "../obilayeredmap" }
|
||||
obipipeline = { path = "../obipipeline" }
|
||||
obiread = { path = "../obiread" }
|
||||
obiskbuilder = { path = "../obiskbuilder" }
|
||||
obiskio = { path = "../obiskio" }
|
||||
rayon = "1"
|
||||
sysinfo = "0.39"
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
mod kmer_sort;
|
||||
mod partition;
|
||||
|
||||
pub use partition::{partition_dir, KmerSpectrum, PartitionRouter, PARTITIONS_SUBDIR};
|
||||
pub use partition::{KmerSpectrum, PartitionRouter};
|
||||
|
||||
@@ -1,10 +1,12 @@
|
||||
//! K-mer partitioning: routing super-kmers into per-partition files,
|
||||
//! deduplicating them, and counting unique canonical k-mers.
|
||||
//! K-mer partitioning: routing super-kmers into per-partition, layer-0
|
||||
//! files, deduplicating them, and counting unique canonical k-mers.
|
||||
//!
|
||||
//! Submodules: [`router`] (`PartitionRouter`, `KmerSpectrum`, the
|
||||
//! `partition_dir` naming convention, and the routing/lifecycle API),
|
||||
//! [`dereplicate`] (two-phase split+merge deduplication), [`count`]
|
||||
//! (unique-kmer enumeration, MPHF, abundance counting).
|
||||
//! routing/lifecycle API — partition/layer path naming itself lives on
|
||||
//! `obikindex::KmerIndex`, not here, see
|
||||
//! `DevDocMD/implementation/partition_layer_cache.md`), [`dereplicate`]
|
||||
//! (two-phase split+merge deduplication), [`count`] (unique-kmer
|
||||
//! enumeration, MPHF, abundance counting).
|
||||
|
||||
mod count;
|
||||
mod dereplicate;
|
||||
@@ -13,7 +15,6 @@ mod router;
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
|
||||
pub use router::{partition_dir, KmerSpectrum, PartitionRouter};
|
||||
pub use router::{KmerSpectrum, PartitionRouter};
|
||||
|
||||
const SK_EXT: &str = "skmer.zst";
|
||||
pub const PARTITIONS_SUBDIR: &str = "partitions";
|
||||
|
||||
@@ -1,37 +1,32 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::fs;
|
||||
use std::io;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::path::PathBuf;
|
||||
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
|
||||
use std::sync::Arc;
|
||||
use std::time::Instant;
|
||||
|
||||
use obisys::progress_bar;
|
||||
|
||||
use obikindex::KmerIndex;
|
||||
use obikseq::RoutableSuperKmer;
|
||||
use obilayeredmap::Layer;
|
||||
use obiskio::SKResult;
|
||||
use obisys::{progress_bar, Progress};
|
||||
use rayon::prelude::*;
|
||||
use remove_dir_all::remove_dir_all;
|
||||
use sysinfo::System;
|
||||
use tracing::info;
|
||||
|
||||
use niffler::Level;
|
||||
use niffler::send::compression::Format;
|
||||
use obiskio::SKFileWriter;
|
||||
|
||||
use obipipeline::{throttle, ThrottleGuard, Throttled};
|
||||
use obiread::NucPage;
|
||||
|
||||
use crate::kmer_sort::chunk_size_from_ram;
|
||||
|
||||
use super::count::count_partition;
|
||||
use super::dereplicate::{dereplicate_partition, optimal_buckets};
|
||||
use super::{PARTITIONS_SUBDIR, SK_EXT};
|
||||
|
||||
/// Path of partition `i`'s directory under `root` — the single source of
|
||||
/// truth for the `part_{i:05}` on-disk naming convention. Shared by
|
||||
/// `PartitionRouter` (which writes here) and, one crate up, `KmerIndex`
|
||||
/// (which builds `index_dir`/`layer_dir` on top of this same root, once
|
||||
/// `PartitionRouter` has finished writing) — the only piece of
|
||||
/// partition-directory knowledge that genuinely needs to cross the
|
||||
/// crate boundary, since both sides must agree on where a partition lives.
|
||||
pub fn partition_dir(root: &Path, i: usize) -> PathBuf {
|
||||
root.join(PARTITIONS_SUBDIR).join(format!("part_{i:05}"))
|
||||
}
|
||||
use super::SK_EXT;
|
||||
|
||||
pub struct KmerSpectrum {
|
||||
pub f0: u64,
|
||||
@@ -39,75 +34,124 @@ pub struct KmerSpectrum {
|
||||
pub counts: BTreeMap<u32, u64>,
|
||||
}
|
||||
|
||||
/// Routes raw super-kmers into per-partition files, then dereplicates and
|
||||
/// counts them — this crate's entire job now that layer/query/merge/select/
|
||||
/// rebuild/dump/distance concerns have moved to `obikindex` (they operate on
|
||||
/// built layers, which don't exist yet at this stage — see
|
||||
/// `DevDocMD/implementation/partition_layer_cache.md`). Transient: it makes
|
||||
/// sense only while raw partition files are being written or processed, a
|
||||
/// phase that always precedes any `Layer`.
|
||||
pub struct PartitionRouter {
|
||||
root_path: PathBuf,
|
||||
// ── Pipeline plumbing, private to `run` ─────────────────────────────────────
|
||||
|
||||
/// Carrier enum for `obipipeline::make_pipe!`'s two-stage transform — local
|
||||
/// to this crate, not `obikmer`'s own `PipelineData` (which stays scoped to
|
||||
/// its other CLI commands): a library crate can't depend on the binary
|
||||
/// that consumes it, so this is a self-contained duplicate of the same
|
||||
/// shape, not a shared type.
|
||||
enum PipelineData {
|
||||
Path(Throttled<PathBuf>),
|
||||
NucPage(NucPage),
|
||||
Batch(Vec<RoutableSuperKmer>),
|
||||
}
|
||||
|
||||
unsafe impl Send for PipelineData {}
|
||||
unsafe impl Sync for PipelineData {}
|
||||
|
||||
/// Keeps a file's throttle-slot guard alive until the file's page iterator
|
||||
/// is exhausted, so the next queued file can only start once this one has
|
||||
/// actually finished producing pages, not merely been dequeued.
|
||||
struct GuardedIter {
|
||||
inner: Box<dyn Iterator<Item = NucPage> + Send>,
|
||||
_guard: ThrottleGuard,
|
||||
flat_active: Arc<AtomicU32>,
|
||||
}
|
||||
|
||||
impl Iterator for GuardedIter {
|
||||
type Item = NucPage;
|
||||
fn next(&mut self) -> Option<NucPage> {
|
||||
self.inner.next()
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for GuardedIter {
|
||||
fn drop(&mut self) {
|
||||
self.flat_active.fetch_sub(1, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
// ── PartitionRouter ──────────────────────────────────────────────────────────
|
||||
|
||||
/// Routes raw super-kmers into per-partition, layer-0 files, then
|
||||
/// dereplicates and counts them — this crate's entire job now that
|
||||
/// layer/query/merge/select/rebuild/dump/distance concerns have moved to
|
||||
/// `obikindex` (they operate on built layers, which don't exist yet at this
|
||||
/// stage — see `DevDocMD/implementation/partition_layer_cache.md`).
|
||||
///
|
||||
/// Holds `&mut KmerIndex` — this is an algorithm operating on an index, not
|
||||
/// a data structure of its own; it owns no path-naming knowledge (every
|
||||
/// path comes from `index.index_dir`/`obilayeredmap::layer_dir`/
|
||||
/// `Layer::create`), only the transient routing/dereplication/counting
|
||||
/// state a run needs.
|
||||
///
|
||||
/// Two-phase construction: `new` + optional setters (`level_max`/`theta`/
|
||||
/// `workers`/`max_open`) configure the run, `run` executes it. `run` takes
|
||||
/// an `Option` progress callback — the router reports raw `(position,
|
||||
/// total)` ticks via [`obisys::Progress`] and stays unaware of *how* (or
|
||||
/// whether) the caller displays them; rendering a spinner/progress bar is
|
||||
/// the caller's decision, not this crate's.
|
||||
pub struct PartitionRouter<'a> {
|
||||
index: &'a mut KmerIndex,
|
||||
n_partitions: usize,
|
||||
partitions_mask: u64,
|
||||
writers: Vec<Option<SKFileWriter>>,
|
||||
level: Level,
|
||||
closed: bool,
|
||||
level_max: usize,
|
||||
theta: f64,
|
||||
workers: usize,
|
||||
max_open: usize,
|
||||
}
|
||||
|
||||
impl PartitionRouter {
|
||||
/// Create a fresh partition layout at `root_path` for `n_partitions =
|
||||
/// 2^n_bits` partitions.
|
||||
pub fn create(root_path: &Path, n_bits: usize, force: bool) -> SKResult<Self> {
|
||||
// `root_path` itself may already exist as a bare directory: callers
|
||||
// typically hold an index-level lock file there before creating the
|
||||
// partition layout. What actually signals a pre-existing partition
|
||||
// set is the `PARTITIONS_SUBDIR` subdirectory, not the root itself.
|
||||
if root_path.join(PARTITIONS_SUBDIR).exists() {
|
||||
if force {
|
||||
remove_dir_all(root_path)?;
|
||||
} else {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::AlreadyExists,
|
||||
format!("{}: partition directory already exists", root_path.display()),
|
||||
)
|
||||
.into());
|
||||
}
|
||||
}
|
||||
fs::create_dir_all(root_path.join(PARTITIONS_SUBDIR))?;
|
||||
Self::new(root_path, n_bits)
|
||||
}
|
||||
|
||||
/// Reopen the partition layout at `root_path` for further routing (or
|
||||
/// for `dereplicate`/`count_kmer`, which don't need `writers` but reuse
|
||||
/// this same handle for consistency). Every caller of `open` wants to
|
||||
/// write or process, never just read paths (those live on `KmerIndex`
|
||||
/// directly), so this always starts open too.
|
||||
pub fn open(root_path: &Path, n_bits: usize) -> SKResult<Self> {
|
||||
if !root_path.exists() {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::NotFound,
|
||||
format!("{}: partition directory not found", root_path.display()),
|
||||
)
|
||||
.into());
|
||||
}
|
||||
Self::new(root_path, n_bits)
|
||||
}
|
||||
|
||||
fn new(root_path: &Path, n_bits: usize) -> SKResult<Self> {
|
||||
impl<'a> PartitionRouter<'a> {
|
||||
/// Configure a router for `index`'s partition layout. Doesn't touch
|
||||
/// disk by itself — partitions and their layer-0 shells are created
|
||||
/// lazily, on the first super-kmer routed to each of them.
|
||||
pub fn new(index: &'a mut KmerIndex) -> Self {
|
||||
let n_bits = index.n_bits();
|
||||
let n_partitions = 1usize << n_bits;
|
||||
let writers = (0..n_partitions).map(|_| None).collect();
|
||||
Ok(Self {
|
||||
root_path: root_path.to_owned(),
|
||||
let workers = obisys::effective_parallelism();
|
||||
Self {
|
||||
index,
|
||||
n_partitions,
|
||||
partitions_mask: (1u64 << n_bits) - 1,
|
||||
writers,
|
||||
writers: (0..n_partitions).map(|_| None).collect(),
|
||||
level: Level::One,
|
||||
closed: false,
|
||||
})
|
||||
level_max: 6,
|
||||
theta: 0.7,
|
||||
workers,
|
||||
max_open: (workers / 4).max(1),
|
||||
}
|
||||
}
|
||||
|
||||
/// Route and write one super-kmer to its partition file.
|
||||
/// Maximum sub-word size for entropy computation (superkmer building).
|
||||
pub fn level_max(mut self, v: usize) -> Self {
|
||||
self.level_max = v;
|
||||
self
|
||||
}
|
||||
|
||||
/// Entropy threshold (k-mers with score ≤ theta are rejected).
|
||||
pub fn theta(mut self, v: f64) -> Self {
|
||||
self.theta = v;
|
||||
self
|
||||
}
|
||||
|
||||
/// Number of worker threads for `run`'s file-reading/superkmer-building pipeline.
|
||||
pub fn workers(mut self, v: usize) -> Self {
|
||||
self.workers = v;
|
||||
self
|
||||
}
|
||||
|
||||
/// Maximum number of input files `run` opens simultaneously.
|
||||
pub fn max_open(mut self, v: usize) -> Self {
|
||||
self.max_open = v;
|
||||
self
|
||||
}
|
||||
|
||||
/// Route and write one super-kmer to its partition's raw super-kmer file.
|
||||
pub fn write(&mut self, rsk: RoutableSuperKmer) -> SKResult<()> {
|
||||
self.check_not_closed()?;
|
||||
let partition = (rsk.minimizer().seq_hash() & self.partitions_mask) as usize;
|
||||
@@ -149,6 +193,82 @@ impl PartitionRouter {
|
||||
!self.closed
|
||||
}
|
||||
|
||||
/// Run the full scatter pipeline: normalise every file in
|
||||
/// `path_source` -> build super-kmers -> route -> write, then close.
|
||||
/// `on_progress`, when set, is called periodically (rate-limited, not
|
||||
/// once per batch) with cumulative bases processed so far —
|
||||
/// `total: None`, since the total isn't known without pre-scanning
|
||||
/// every input file.
|
||||
pub fn run(
|
||||
&mut self,
|
||||
path_source: impl Iterator<Item = PathBuf> + Send + 'static,
|
||||
mut on_progress: Option<impl FnMut(Progress)>,
|
||||
) -> SKResult<()> {
|
||||
let k = self.index.kmer_size();
|
||||
let level_max = self.level_max;
|
||||
let theta = self.theta;
|
||||
let n_workers = self.workers;
|
||||
let max_open = self.max_open;
|
||||
|
||||
// Throttle in the source thread — never in a worker — to prevent deadlock.
|
||||
let throttled = throttle(path_source, max_open);
|
||||
|
||||
let file_count = Arc::new(AtomicU64::new(0));
|
||||
let flat_active = Arc::new(AtomicU32::new(0));
|
||||
let transform_active = Arc::new(AtomicU32::new(0));
|
||||
|
||||
let pipe = obipipeline::make_pipe! {
|
||||
PipelineData : Throttled<PathBuf> => Vec<RoutableSuperKmer>,
|
||||
||? {
|
||||
let file_count = Arc::clone(&file_count);
|
||||
let flat_active = Arc::clone(&flat_active);
|
||||
move |pw: Throttled<PathBuf>| {
|
||||
let path = pw.item;
|
||||
let guard = pw.guard;
|
||||
let n = file_count.fetch_add(1, Ordering::Relaxed) + 1;
|
||||
info!("indexing [{}]: {}", n, path.display());
|
||||
let path_str = path.to_str().unwrap_or("").to_owned();
|
||||
flat_active.fetch_add(1, Ordering::Relaxed);
|
||||
obiread::open_nuc_stream(&path_str, k)
|
||||
.map(|iter| GuardedIter { inner: iter, _guard: guard, flat_active: Arc::clone(&flat_active) })
|
||||
}
|
||||
} : Path => NucPage,
|
||||
| {
|
||||
let transform_active = Arc::clone(&transform_active);
|
||||
move |page| {
|
||||
transform_active.fetch_add(1, Ordering::Relaxed);
|
||||
let result = obiskbuilder::build_superkmers_page(page, k, level_max, theta);
|
||||
transform_active.fetch_sub(1, Ordering::Relaxed);
|
||||
result
|
||||
}
|
||||
} : NucPage => Batch,
|
||||
};
|
||||
|
||||
let mut total_bases: u64 = 0;
|
||||
let mut last_report = Instant::now();
|
||||
let kmer_overlap = (k - 1) as u64;
|
||||
const REPORT_INTERVAL: f64 = 0.1;
|
||||
|
||||
for batch in pipe.apply(throttled, n_workers, 1) {
|
||||
total_bases += batch
|
||||
.iter()
|
||||
.map(|sk| (sk.seql() as u64).saturating_sub(kmer_overlap))
|
||||
.sum::<u64>();
|
||||
if let Some(cb) = on_progress.as_mut() {
|
||||
let now = Instant::now();
|
||||
if now.duration_since(last_report).as_secs_f64() > REPORT_INTERVAL {
|
||||
last_report = now;
|
||||
cb(Progress { position: total_bases, total: None });
|
||||
}
|
||||
}
|
||||
self.write_batch(batch)?;
|
||||
}
|
||||
if let Some(cb) = on_progress.as_mut() {
|
||||
cb(Progress { position: total_bases, total: None });
|
||||
}
|
||||
self.close()
|
||||
}
|
||||
|
||||
/// Deduplicate all `raw.{ext}` files in parallel, replacing each with a
|
||||
/// `dereplicated.{ext}` file where identical canonical sequences are merged
|
||||
/// and their counts summed.
|
||||
@@ -181,7 +301,7 @@ impl PartitionRouter {
|
||||
let results: Vec<SKResult<()>> = (0..self.n_partitions)
|
||||
.into_par_iter()
|
||||
.map(|i| {
|
||||
let dir = partition_dir(&self.root_path, i);
|
||||
let dir = self.layer0_dir(i);
|
||||
if !dir.exists() {
|
||||
pb.inc(1);
|
||||
return Ok(());
|
||||
@@ -226,7 +346,7 @@ impl PartitionRouter {
|
||||
let results: Vec<SKResult<()>> = (0..self.n_partitions)
|
||||
.into_par_iter()
|
||||
.map(|i| {
|
||||
let dir = partition_dir(&self.root_path, i);
|
||||
let dir = self.layer0_dir(i);
|
||||
let dedup_path = dir.join(format!("dereplicated.{SK_EXT}"));
|
||||
if !dedup_path.exists() {
|
||||
pb.inc(1);
|
||||
@@ -251,7 +371,7 @@ impl PartitionRouter {
|
||||
let mut f1: u64 = 0;
|
||||
|
||||
for i in 0..self.n_partitions {
|
||||
let path = partition_dir(&self.root_path, i).join("kmer_spectrum_raw.json");
|
||||
let path = self.layer0_dir(i).join("kmer_spectrum_raw.json");
|
||||
if !path.exists() {
|
||||
continue;
|
||||
}
|
||||
@@ -276,6 +396,14 @@ impl PartitionRouter {
|
||||
|
||||
// ── private ───────────────────────────────────────────────────────────────
|
||||
|
||||
/// Directory of partition `i`'s layer 0 — every raw/dereplicated
|
||||
/// superkmer file and provisional `mphf1.bin`/`counts1.bin` this router
|
||||
/// produces lives here, alongside where `build_index_layer`
|
||||
/// (`obikindex`) will later turn it into the real layer 0.
|
||||
fn layer0_dir(&self, i: usize) -> PathBuf {
|
||||
obilayeredmap::layer_dir(&self.index.index_dir(i), 0)
|
||||
}
|
||||
|
||||
fn check_not_closed(&self) -> SKResult<()> {
|
||||
if self.closed {
|
||||
Err(io::Error::new(io::ErrorKind::BrokenPipe, "write to closed PartitionRouter").into())
|
||||
@@ -286,8 +414,8 @@ impl PartitionRouter {
|
||||
|
||||
fn ensure_writer(&mut self, partition: usize) -> SKResult<&mut SKFileWriter> {
|
||||
if self.writers[partition].is_none() {
|
||||
let dir = partition_dir(&self.root_path, partition);
|
||||
fs::create_dir_all(&dir)?;
|
||||
let dir = self.layer0_dir(partition);
|
||||
Layer::create(&dir).map_err(|e| io::Error::other(e.to_string()))?;
|
||||
let file_path = dir.join(format!("raw.{SK_EXT}"));
|
||||
let writer = SKFileWriter::create_with(file_path, Format::Zstd, self.level)?;
|
||||
self.writers[partition] = Some(writer);
|
||||
@@ -296,7 +424,7 @@ impl PartitionRouter {
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for PartitionRouter {
|
||||
impl Drop for PartitionRouter<'_> {
|
||||
fn drop(&mut self) {
|
||||
let _ = self.close();
|
||||
}
|
||||
|
||||
@@ -1,16 +1,30 @@
|
||||
use std::collections::HashMap;
|
||||
use std::fs;
|
||||
|
||||
use obikindex::{IndexConfig, KmerIndex};
|
||||
use obikrope::Rope;
|
||||
use obikseq::SuperKmer;
|
||||
use obilayeredmap::IndexMode;
|
||||
use obiskbuilder::build_superkmers;
|
||||
|
||||
use super::count::count_partition;
|
||||
use super::{PartitionRouter, PARTITIONS_SUBDIR};
|
||||
use super::PartitionRouter;
|
||||
|
||||
const K: usize = 11;
|
||||
const M: usize = 5;
|
||||
|
||||
fn test_index(dir: &std::path::Path) -> KmerIndex {
|
||||
let config = IndexConfig {
|
||||
kmer_size: K,
|
||||
minimizer_size: M,
|
||||
n_bits: 0, // 1 partition — matches this suite's ground-truth setup
|
||||
with_counts: false,
|
||||
evidence: IndexMode::Exact,
|
||||
block_bits: 0,
|
||||
};
|
||||
KmerIndex::create(dir, config, None).unwrap()
|
||||
}
|
||||
|
||||
fn setup() {
|
||||
obikseq::params::set_k(K);
|
||||
obikseq::params::set_m(M);
|
||||
@@ -46,12 +60,14 @@ fn pipeline_counts(seqs: &[&[u8]]) -> (u64, u64) {
|
||||
let superkmers: Vec<_> = build_superkmers(rope, K, 1, 0.0);
|
||||
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let mut kp = PartitionRouter::create(dir.path(), 0, true).unwrap();
|
||||
let mut index = test_index(&dir.path().join("idx"));
|
||||
let mut kp = PartitionRouter::new(&mut index);
|
||||
kp.write_batch(superkmers).unwrap();
|
||||
kp.close().unwrap();
|
||||
kp.dereplicate().unwrap();
|
||||
drop(kp); // ends the borrow of `index` early — `PartitionRouter`'s `Drop` impl would otherwise extend it to the end of scope
|
||||
|
||||
let part_dir = dir.path().join(PARTITIONS_SUBDIR).join("part_00000");
|
||||
let part_dir = index.layer_dir(0, 0);
|
||||
let dedup_path = part_dir.join("dereplicated.skmer.zst");
|
||||
if !dedup_path.exists() {
|
||||
return (0, 0);
|
||||
|
||||
@@ -7,6 +7,7 @@ use obisys::Reporter;
|
||||
use tempfile::tempdir;
|
||||
|
||||
use obikindex::{GenomeInfo, IndexConfig, KmerIndex, MergeMode};
|
||||
use obikpartitionner::PartitionRouter;
|
||||
|
||||
use super::alignment::SnpAlignmentExt;
|
||||
use super::build::SiblingAnnexBuildExt;
|
||||
@@ -64,19 +65,24 @@ fn build_single_genome_index(dir: &Path, label: &str, seq: &[u8]) -> KmerIndex {
|
||||
evidence: obilayeredmap::IndexMode::Exact,
|
||||
block_bits: 0,
|
||||
};
|
||||
let mut idx = KmerIndex::create(&index_path, config, Some(GenomeInfo::new(label)), false)
|
||||
let mut idx = KmerIndex::create(&index_path, config, Some(GenomeInfo::new(label)))
|
||||
.expect("create");
|
||||
|
||||
let mut rep = Reporter::new();
|
||||
let stream = obiread::open_nuc_stream(fasta_path.to_str().unwrap(), K).expect("open fasta");
|
||||
let mut router = idx.partition_router().expect("partition_router");
|
||||
let mut router = PartitionRouter::new(&mut idx);
|
||||
for page in stream {
|
||||
let batch = obiskbuilder::build_superkmers_page(page, K, /* level_max */ 1, /* theta */ 0.0);
|
||||
router.write_batch(batch).expect("write_batch");
|
||||
}
|
||||
router.close().expect("close partition writers");
|
||||
router.dereplicate().expect("dereplicate");
|
||||
let spectrum = router.count_kmer(false).expect("count_kmer");
|
||||
drop(router); // ends the borrow of `idx` early — `PartitionRouter`'s `Drop` impl would otherwise extend it to the end of scope
|
||||
|
||||
idx.mark_scattered().expect("mark_scattered");
|
||||
idx.dereplicate_and_count(false, &mut rep).expect("dereplicate_and_count");
|
||||
idx.write_spectrum(spectrum.f0, spectrum.f1, &spectrum.counts).expect("write_spectrum");
|
||||
idx.mark_counted().expect("mark_counted");
|
||||
idx.build_layers(1, None, false, &mut rep).expect("build_layers");
|
||||
idx
|
||||
}
|
||||
|
||||
@@ -9,6 +9,6 @@ mod stage;
|
||||
|
||||
pub use budget::MemoryBudget;
|
||||
pub use lock::DirLock;
|
||||
pub use progress::{TracedBar, progress_bar, spinner};
|
||||
pub use progress::{Progress, TracedBar, progress_bar, spinner};
|
||||
pub use resources::{CpuSample, IoSample, available_memory_bytes, effective_parallelism, peak_rss_bytes};
|
||||
pub use stage::{Reporter, Stage, StageStats};
|
||||
|
||||
@@ -5,6 +5,27 @@ use std::time::{Duration, Instant};
|
||||
use indicatif::{ProgressBar, ProgressStyle};
|
||||
use tracing::{debug, info};
|
||||
|
||||
// ── Generic progress reporting ──────────────────────────────────────────────
|
||||
|
||||
/// One tick of progress from a long-running algorithm, passed to a
|
||||
/// caller-supplied callback — the standard shape every algo crate's `run()`
|
||||
/// should use, so a CLI command only has to learn to read this once
|
||||
/// regardless of which algorithm is reporting.
|
||||
///
|
||||
/// `total: None` means the algorithm doesn't know its total work up front
|
||||
/// (e.g. bases scattered from input files of unknown size until read) — the
|
||||
/// caller should render a spinner. `total: Some(n)` means it does — the
|
||||
/// caller can render an actual progress bar. The algorithm itself never
|
||||
/// decides which; it only reports what it knows, `position` and `total`,
|
||||
/// and stays unaware of how — or whether — the caller displays it.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct Progress {
|
||||
pub position: u64,
|
||||
pub total: Option<u64>,
|
||||
}
|
||||
|
||||
// ── Terminal progress bar/spinner ────────────────────────────────────────────
|
||||
|
||||
const BRAILLE: &[&str] = &["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"];
|
||||
const ETA_REFRESH_MS: u64 = 500;
|
||||
const ETA_MIN_ELAPSED_MS: u64 = 1000;
|
||||
|
||||
Reference in New Issue
Block a user