OBIKmers/obikmer
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 34 Wiki Activity

refactor: rename obikindex to obikindexer

Browse Source
This commit is contained in:
Eric Coissac
2026-08-21 05:29:34 +02:00
parent 5c1584967f
commit 96b6517541
22 changed files with 132 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/obikindex/Cargo.toml
+1 -4
View File
@@ -12,8 +12,6 @@ obicompactvec = { path = "../obicompactvec" }
obidebruinj = { path = "../obidebruinj" }
obipipeline = { path = "../obipipeline" }
obikentropy = { path = "../obikentropy" }
obiread = { path = "../obiread" }
obiskbuilder = { path = "../obiskbuilder" }
cacheline-ef = "1.1"
epserde = "0.8"
ptr_hash = "1.1"
@@ -27,12 +25,11 @@ serde_json = "1"
indicatif = "0.18"
tracing = "0.1.44"
bitvec = "1"
sysinfo = "0.39"
hwlocality = { version = "1.0.0-alpha.11", features = ["vendored"], optional = true }
[dev-dependencies]
obiread = { path = "../obiread" }
obikseq = { path = "../obikseq", features = ["test-utils"] }
obikrope = { path = "../obikrope" }
tempfile = "3"
tracing-subscriber = { version = "0.3", features = ["fmt", "env-filter"] }
anyhow = "1"
src/obikindex/src/algorithms/dereplicator/dereplicate.rs
-155
View File
@@ -1,155 +0,0 @@
//! Per-partition dereplication mechanics — private to this module.
//! [`crate::algorithms::dereplicator::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;
use std::path::{Path, PathBuf};
use tracing::debug;
use niffler::send::compression::Format;
use niffler::Level;
use obikseq::superkmer::SuperKmer;
use obikseq::Sequence;
use crate::layer::{dereplicated_superkmers_path, raw_superkmers_path};
use obiskio::{SKFileMeta, SKFileReader, SKFileWriter, SKResult};
/// Scratch-file extension for this algorithm's own intermediate split
/// buckets — never read by anything outside [`dereplicate_partition`],
/// unlike `raw`/`dereplicated` (see `crate::layer::{raw_superkmers_path,
/// dereplicated_superkmers_path}`, the actual cross-module 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.
///
/// Memory per HashMap entry:
/// key Box (1 + avg_seq_bytes) + SuperKmer header (4 B) + avg seq bytes + u64 count (8 B),
/// multiplied by 1.5 for hashbrown load-factor overhead.
///
/// Returns 1 if the partition fits comfortably in memory (no split needed).
/// Always returns a power of two.
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;
let meta = match SKFileMeta::read(raw_path) {
Ok(Some(m)) if m.instances > 0 => m,
_ => return 1,
};
let avg_seq_bytes = ((meta.length_sum + meta.instances - 1) / meta.instances + 3) / 4;
// SuperKmer: header (4 B) + Box<[u8]> ptr+len (16 B) + heap seq bytes; value: u64 (8 B); ×1.5 for hashbrown overhead.
let bytes_per_entry = ((4 + 16 + avg_seq_bytes + 8) as f64 * 1.5) as u64;
let estimated = meta.instances * bytes_per_entry;
if estimated <= budget {
debug!("Dereplication: estimated={estimated} budget={budget} n_temp=1");
return 1;
}
// Round up to the next power of two.
let n = (estimated + budget - 1) / budget;
debug!("Dereplication: estimated={estimated} budget={budget} n_temp={n}");
n.next_power_of_two() as usize
}
/// Remove a SuperKmer file and its sidecar (if present).
fn remove_skmer_file(path: &Path) -> SKResult<()> {
fs::remove_file(path)?;
let sidecar = SKFileMeta::sidecar_path(path);
match fs::remove_file(&sidecar) {
Ok(()) => {}
Err(e) if e.kind() == io::ErrorKind::NotFound => {}
Err(e) => return Err(e.into()),
}
Ok(())
}
/// 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'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 = dereplicated_superkmers_path(dir);
let mut writer = SKFileWriter::create_with(&out_path, Format::Zstd, level)?;
if n_temp == 1 {
// ── Direct path: partition fits in memory, no split needed ────────────
let map = load_bucket(&raw_path)?;
remove_skmer_file(&raw_path)?;
flush_map(map, &mut writer)?;
} else {
// ── 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}.{TEMP_EXT}")))
.collect();
{
let mut writers: Vec<SKFileWriter> = temp_paths
.iter()
.map(|p| SKFileWriter::create_with(p, Format::Zstd, level))
.collect::<SKResult<_>>()?;
let mut reader = SKFileReader::open(&raw_path)?;
while let Some(sk) = reader.read()? {
let bucket = (sk.seq_hash() & temp_mask) as usize;
writers[bucket].write(&sk)?;
}
for w in &mut writers {
w.close()?;
}
}
remove_skmer_file(&raw_path)?;
// ── Phase 2: merge each temp bucket into the output ───────────────────
for temp_path in &temp_paths {
let map = load_bucket(temp_path)?;
remove_skmer_file(temp_path)?;
flush_map(map, &mut writer)?;
}
}
writer.close()?;
Ok(())
}
/// Read a SuperKmer file into a deduplication map (already canonical).
fn load_bucket(path: &Path) -> SKResult<HashMap<SuperKmer, u64>> {
let capacity = SKFileMeta::read(path)
.ok()
.flatten()
.map(|m| m.instances as usize)
.unwrap_or(0);
let mut map: HashMap<SuperKmer, u64> = HashMap::with_capacity(capacity);
let mut reader = SKFileReader::open(path)?;
while let Some(sk) = reader.read()? {
let count = sk.count() as u64;
*map.entry(sk).or_insert(0) += count;
}
Ok(map)
}
/// Write all entries of a deduplication map to `writer`, splitting oversized counts.
fn flush_map(map: HashMap<SuperKmer, u64>, writer: &mut SKFileWriter) -> SKResult<()> {
for (mut sk, mut total) in map {
while total > MAX_SK_COUNT {
sk.set_count(MAX_SK_COUNT as u32);
writer.write(&sk)?;
total -= MAX_SK_COUNT;
}
sk.set_count(total as u32);
writer.write(&sk)?;
}
Ok(())
}
src/obikindex/src/algorithms/dereplicator/mod.rs
-106
View File
@@ -1,106 +0,0 @@
//! Superkmer dereplication — the sibling algorithm to
//! [`crate::algorithms::partitionner`]: runs after
//! `partitionner::PartitionRouter::run` (scatter) has written each
//! partition's raw superkmer file, before
//! `partitionner::PartitionRouter::count_kmer` (counting) reads the
//! result — see `DevDocMD/implementation/partition_layer_cache.md`.
mod dereplicate;
use std::sync::atomic::{AtomicU64, Ordering};
use niffler::Level;
use obiskio::SKResult;
use obisys::Progress;
use rayon::prelude::*;
use sysinfo::System;
use crate::index::KmerIndex;
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 algorithm
/// never renders anything itself.
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 = crate::layer::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(())
}
}
src/obikindex/src/algorithms/mod.rs
-9
View File
@@ -1,9 +0,0 @@
//! Indexing-pipeline algorithms: code that operates on a
//! [`crate::index::KmerIndex`] to build or transform its content, as
//! opposed to the `index`/`partition`/`layer` modules, which are the data
//! model itself. Each algorithm is its own submodule: [`partitionner`]
//! (routing raw super-kmers into partitions, then counting), [`dereplicator`]
//! (deduplicating a partition's raw super-kmers before counting).
pub mod dereplicator;
pub mod partitionner;
src/obikindex/src/algorithms/partitionner/count.rs
-80
View File
@@ -1,80 +0,0 @@
use std::collections::BTreeMap;
use std::fs;
use std::io;
use std::path::Path;
use cacheline_ef::{CachelineEf, CachelineEfVec};
use epserde::ser::Serialize as EpSerialize;
use memmap2::Mmap;
use obicompactvec::PersistentCompactIntVecBuilder;
use obiskio::{SKFileReader, SKResult};
use ptr_hash::{PtrHash, PtrHashParams, bucket_fn::CubicEps, hash::Xx64};
use tracing::debug;
use super::kmer_sort::sort_unique_kmers;
pub(super) type Mphf = PtrHash<u64, CubicEps, CachelineEfVec<Vec<CachelineEf>>, Xx64, Vec<u8>>;
fn build_mphf(unique_path: &Path, f0: usize) -> io::Result<Mphf> {
let file = fs::File::open(unique_path)?;
let mmap = unsafe { Mmap::map(&file)? };
let kmers: &[u64] = unsafe {
std::slice::from_raw_parts(mmap.as_ptr() as *const u64, f0)
};
// Sequential constructor: the outer par_iter over partitions already saturates
// the Rayon pool. new_from_par_iter would get no additional threads and adds
// coordination overhead. try_new accesses the same mmap'd pages at zero extra cost.
Mphf::try_new(kmers, PtrHashParams::<CubicEps>::default())
.ok_or_else(|| io::Error::other("ptr_hash construction failed"))
}
pub(super) fn count_partition(dir: &Path, dedup_path: &Path, chunk_kmers: usize) -> SKResult<()> {
let unique_path = dir.join("sorted_unique.bin");
let f0 = sort_unique_kmers(dedup_path, dir, &unique_path, chunk_kmers)?;
if f0 == 0 {
return Ok(());
}
debug!("{}: f0={f0} unique kmers sorted", dir.display());
let mphf = build_mphf(&unique_path, f0)?;
fs::remove_file(&unique_path)?;
let counts_path = dir.join("counts1.bin");
let mut builder = PersistentCompactIntVecBuilder::new(f0, &counts_path)?;
{
let mut reader = SKFileReader::open(dedup_path)?;
while let Some(sk) = reader.read()? {
let sk_count = sk.count();
for kmer in sk.iter_canonical_kmers() {
let slot = mphf.index(&kmer.raw());
builder.set(slot, builder.get(slot).saturating_add(sk_count));
}
}
}
let mut spectrum: BTreeMap<u32, u64> = BTreeMap::new();
for slot in 0..f0 {
let c = builder.get(slot);
if c > 0 {
*spectrum.entry(c).or_insert(0) += 1;
}
}
let f1: u64 = spectrum.iter().map(|(&c, &f)| c as u64 * f).sum();
builder.close()?;
let spectrum_map: BTreeMap<String, u64> = spectrum
.iter()
.map(|(&c, &f)| (format!("{c:010}"), f))
.collect();
serde_json::to_writer_pretty(
fs::File::create(dir.join("kmer_spectrum_raw.json"))?,
&serde_json::json!({ "f0": f0 as u64, "f1": f1, "spectrum": &spectrum_map }),
)
.map_err(io::Error::other)?;
EpSerialize::store(&mphf, &dir.join("mphf1.bin"))
.map_err(|e| io::Error::other(e.to_string()))?;
Ok(())
}
src/obikindex/src/algorithms/partitionner/kmer_sort.rs
-126
View File
@@ -1,126 +0,0 @@
use std::cmp::Reverse;
use std::collections::BinaryHeap;
use std::fs;
use std::io::{self, BufWriter, Write};
use std::path::{Path, PathBuf};
use memmap2::Mmap;
use obiskio::{SKFileReader, SKResult};
/// Extract all canonical kmers from a dereplicated superkmer file,
/// sort them with an external merge sort, deduplicate, and write
/// unique u64 values to `unique_path`. Returns f0 (distinct kmers).
pub fn sort_unique_kmers(
dedup_path: &Path,
work_dir: &Path,
unique_path: &Path,
chunk_kmers: usize,
) -> SKResult<usize> {
let chunk_paths = extract_and_sort_chunks(dedup_path, work_dir, chunk_kmers)?;
if chunk_paths.is_empty() {
fs::write(unique_path, [])?;
return Ok(0);
}
let f0 = merge_and_dedup(&chunk_paths, unique_path)?;
for p in &chunk_paths {
fs::remove_file(p)?;
}
Ok(f0)
}
/// Number of kmers per sort chunk from available RAM (per thread).
pub fn chunk_size_from_ram(available_bytes: u64) -> usize {
// 40% of available RAM; each kmer is 8 bytes (u64)
let n = ((available_bytes as f64 * 0.40) / 8.0) as usize;
n.max(1 << 20) // minimum 1M kmers (~8 MB)
}
// ── private ───────────────────────────────────────────────────────────────────
fn extract_and_sort_chunks(
dedup_path: &Path,
work_dir: &Path,
chunk_kmers: usize,
) -> SKResult<Vec<PathBuf>> {
let mut reader = SKFileReader::open(dedup_path)?;
let mut buf: Vec<u64> = Vec::with_capacity(chunk_kmers);
let mut paths: Vec<PathBuf> = Vec::new();
while let Some(sk) = reader.read()? {
for kmer in sk.iter_canonical_kmers() {
buf.push(kmer.raw());
if buf.len() >= chunk_kmers {
paths.push(flush_sorted_chunk(&mut buf, work_dir, paths.len())?);
}
}
}
if !buf.is_empty() {
paths.push(flush_sorted_chunk(&mut buf, work_dir, paths.len())?);
}
Ok(paths)
}
fn flush_sorted_chunk(buf: &mut Vec<u64>, work_dir: &Path, idx: usize) -> io::Result<PathBuf> {
buf.sort_unstable();
let path = work_dir.join(format!("kmer_sort_{idx:05}.bin"));
let mut w = BufWriter::new(fs::File::create(&path)?);
for &v in buf.iter() {
w.write_all(&v.to_le_bytes())?;
}
buf.clear();
Ok(path)
}
/// K-way merge of sorted chunk files + in-line dedup → `unique_path`.
fn merge_and_dedup(chunk_paths: &[PathBuf], unique_path: &Path) -> io::Result<usize> {
struct ChunkReader {
mmap: Mmap,
pos: usize,
}
impl ChunkReader {
fn open(path: &Path) -> io::Result<Self> {
let f = fs::File::open(path)?;
Ok(Self { mmap: unsafe { Mmap::map(&f)? }, pos: 0 })
}
fn peek(&self) -> Option<u64> {
let off = self.pos * 8;
if off + 8 <= self.mmap.len() {
Some(u64::from_le_bytes(self.mmap[off..off + 8].try_into().unwrap()))
} else {
None
}
}
fn advance(&mut self) { self.pos += 1; }
}
let mut readers: Vec<ChunkReader> = chunk_paths.iter()
.map(|p| ChunkReader::open(p))
.collect::<io::Result<_>>()?;
let mut heap: BinaryHeap<Reverse<(u64, usize)>> = BinaryHeap::new();
for (i, r) in readers.iter().enumerate() {
if let Some(v) = r.peek() {
heap.push(Reverse((v, i)));
}
}
let mut w = BufWriter::new(fs::File::create(unique_path)?);
let mut f0 = 0usize;
let mut prev: Option<u64> = None;
while let Some(Reverse((val, idx))) = heap.pop() {
if prev != Some(val) {
w.write_all(&val.to_le_bytes())?;
prev = Some(val);
f0 += 1;
}
readers[idx].advance();
if let Some(next_val) = readers[idx].peek() {
heap.push(Reverse((next_val, idx)));
}
}
Ok(f0)
}
src/obikindex/src/algorithms/partitionner/mod.rs
-20
View File
@@ -1,20 +0,0 @@
//! K-mer partitioning: routing super-kmers into per-partition, layer-0
//! files, and counting unique canonical k-mers. Dereplication of the raw
//! super-kmers in between is [`crate::algorithms::dereplicator`], a
//! sibling algorithm, not a step of this one — see
//! `DevDocMD/implementation/partition_layer_cache.md`.
//!
//! Submodules: [`router`] (`PartitionRouter`, `KmerSpectrum`, the
//! routing/counting lifecycle API — partition/layer path naming itself
//! lives on `crate::index::KmerIndex`, not here), [`count`] (unique-kmer
//! enumeration, MPHF, abundance counting), `kmer_sort` (external sort
//! support for `count`).
mod count;
mod kmer_sort;
mod router;
#[cfg(test)]
mod tests;
pub use router::{KmerSpectrum, PartitionRouter};
src/obikindex/src/algorithms/partitionner/router.rs
-375
View File
@@ -1,375 +0,0 @@
use std::collections::BTreeMap;
use std::fs;
use std::io;
use std::path::PathBuf;
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Instant;
use crate::index::KmerIndex;
use obikseq::RoutableSuperKmer;
use crate::layer::Layer;
use obiskio::SKResult;
use obisys::{progress_bar, Progress};
use rayon::prelude::*;
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 super::kmer_sort::chunk_size_from_ram;
use super::count::count_partition;
pub struct KmerSpectrum {
pub f0: u64,
pub f1: u64,
pub counts: BTreeMap<u32, u64>,
}
// ── 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 (via the sibling [`crate::algorithms::dereplicator`]
/// algorithm) and counts them — the entry point of the indexing pipeline,
/// operating on layer-0 content that doesn'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`/`crate::layer::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<'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 workers = obisys::effective_parallelism();
Self {
index,
n_partitions,
partitions_mask: (1u64 << n_bits) - 1,
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),
}
}
/// 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;
let sk = rsk.into_superkmer();
self.ensure_writer(partition)?.write(&sk)
}
/// Route and write a batch of super-kmers.
pub fn write_batch(&mut self, rsks: Vec<RoutableSuperKmer>) -> SKResult<()> {
self.check_not_closed()?;
for rsk in rsks {
let partition = (rsk.minimizer().seq_hash() & self.partitions_mask) as usize;
let sk = rsk.into_superkmer();
self.ensure_writer(partition)?.write(&sk)?;
}
Ok(())
}
pub fn flush(&mut self) -> SKResult<()> {
self.check_not_closed()?;
for writer in self.writers.iter_mut().flatten() {
writer.flush()?;
}
Ok(())
}
pub fn close(&mut self) -> SKResult<()> {
if self.closed {
return Ok(());
}
self.closed = true;
for writer in self.writers.iter_mut().flatten() {
writer.close()?;
}
Ok(())
}
pub fn is_open(&self) -> bool {
!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()
}
/// 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.
/// 3. Writes a flat binary count file (`counts1.bin`, one `u32` per slot,
/// memory-mapped) accumulating kmer abundances from the superkmer counts.
/// 4. Persists the MPHF to `mphf1.bin` for downstream use.
///
/// Returns the aggregated `KmerSpectrum`. Per-partition spectrum files are
/// deleted after aggregation unless `keep_partial` is true.
pub fn count_kmer(&self, keep_partial: bool) -> SKResult<KmerSpectrum> {
let sys = System::new_all();
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 chunk_kmers = chunk_size_from_ram(available / n_threads);
let pb = progress_bar("counting", 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);
let dedup_path = crate::layer::dereplicated_superkmers_path(&dir);
if !dedup_path.exists() {
pb.inc(1);
return Ok(());
}
let t = Instant::now();
let result = count_partition(&dir, &dedup_path, chunk_kmers);
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?;
}
// Aggregate per-partition spectra.
let mut counts: BTreeMap<u32, u64> = BTreeMap::new();
let mut f0: u64 = 0;
let mut f1: u64 = 0;
for i in 0..self.n_partitions {
let path = self.layer0_dir(i).join("kmer_spectrum_raw.json");
if !path.exists() {
continue;
}
let v: serde_json::Value =
serde_json::from_str(&fs::read_to_string(&path)?).map_err(io::Error::other)?;
f0 += v["f0"].as_u64().unwrap_or(0);
f1 += v["f1"].as_u64().unwrap_or(0);
if let Some(obj) = v["spectrum"].as_object() {
for (c_str, freq) in obj {
if let (Ok(c), Some(f)) = (c_str.parse::<u32>(), freq.as_u64()) {
*counts.entry(c).or_insert(0) += f;
}
}
}
if !keep_partial {
let _ = fs::remove_file(&path);
}
}
Ok(KmerSpectrum { f0, f1, counts })
}
// ── 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`
/// (`crate::index`) will later turn it into the real layer 0.
fn layer0_dir(&self, i: usize) -> PathBuf {
crate::layer::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())
} else {
Ok(())
}
}
fn ensure_writer(&mut self, partition: usize) -> SKResult<&mut SKFileWriter> {
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 = crate::layer::raw_superkmers_path(&dir);
let writer = SKFileWriter::create_with(file_path, Format::Zstd, self.level)?;
self.writers[partition] = Some(writer);
}
Ok(self.writers[partition].as_mut().unwrap())
}
}
impl Drop for PartitionRouter<'_> {
fn drop(&mut self) {
let _ = self.close();
}
}
src/obikindex/src/algorithms/partitionner/tests.rs
-131
View File
@@ -1,131 +0,0 @@
use std::collections::HashMap;
use std::fs;
use crate::algorithms::dereplicator::Dereplicator;
use crate::{IndexConfig, KmerIndex};
use crate::layer::IndexMode;
use obikrope::Rope;
use obikseq::SuperKmer;
use obiskbuilder::build_superkmers;
use super::count::count_partition;
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);
}
/// Direct canonical k-mer counts from ASCII sequences — ground truth.
fn direct_counts(seqs: &[&[u8]]) -> (u64, u64) {
let mut counts: HashMap<Vec<u8>, u64> = HashMap::new();
for seq in seqs {
for i in 0..seq.len().saturating_sub(K - 1) {
let km = SuperKmer::from_ascii(&seq[i..i + K]).to_ascii();
*counts.entry(km).or_insert(0) += 1;
}
}
let f0 = counts.len() as u64;
let f1: u64 = counts.values().sum();
(f0, f1)
}
/// Run the full pipeline on a list of sequences and return (f0, f1) from
/// the `kmer_spectrum_raw.json` produced by `count_partition`.
fn pipeline_counts(seqs: &[&[u8]]) -> (u64, u64) {
setup();
let mut rope_data: Vec<u8> = Vec::new();
for seq in seqs {
rope_data.extend_from_slice(seq);
rope_data.push(0x00);
}
let mut rope = Rope::new(None);
rope.push(rope_data);
let superkmers: Vec<_> = build_superkmers(rope, K, 1, 0.0);
let dir = tempfile::tempdir().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();
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 = crate::layer::dereplicated_superkmers_path(&part_dir);
if !dedup_path.exists() {
return (0, 0);
}
count_partition(&part_dir, &dedup_path, 1 << 20).unwrap();
let spec: serde_json::Value =
serde_json::from_reader(fs::File::open(part_dir.join("kmer_spectrum_raw.json")).unwrap())
.unwrap();
let f0 = spec["f0"].as_u64().unwrap_or(0);
let f1 = spec["f1"].as_u64().unwrap_or(0);
(f0, f1)
}
#[test]
fn single_sequence_f0_f1_match() {
let seqs: &[&[u8]] = &[b"ACGTACGTACGTACGTACGT"];
let (ef0, ef1) = direct_counts(seqs);
let (gf0, gf1) = pipeline_counts(seqs);
assert_eq!(gf0, ef0, "f0 wrong: expected {ef0}, got {gf0}");
assert_eq!(gf1, ef1, "f1 wrong: expected {ef1}, got {gf1}");
}
#[test]
fn two_sequences_f0_f1_match() {
let seqs: &[&[u8]] = &[b"ACGTACGTACGTACGTACGT", b"TGCATGCATGCATGCATGCA"];
let (ef0, ef1) = direct_counts(seqs);
let (gf0, gf1) = pipeline_counts(seqs);
assert_eq!(gf0, ef0, "f0 wrong: expected {ef0}, got {gf0}");
assert_eq!(gf1, ef1, "f1 wrong: expected {ef1}, got {gf1}");
}
#[test]
fn repeated_sequence_f1_doubles() {
let seq = b"ACGTACGTACGTACGTACGT";
let seqs: &[&[u8]] = &[seq, seq];
let (ef0, ef1) = direct_counts(seqs);
let (gf0, gf1) = pipeline_counts(seqs);
assert_eq!(gf0, ef0, "f0 wrong: expected {ef0}, got {gf0}");
assert_eq!(gf1, ef1, "f1 wrong: expected {ef1}, got {gf1}");
}
#[test]
fn many_sequences_f0_f1_match() {
// 20 distinct sequences of length 40 — forces multiple super-kmers and
// multiple minimizer boundaries per sequence.
let bases = b"ACGT";
let seqs: Vec<Vec<u8>> = (0..20u32)
.map(|i| {
(0..40)
.map(|j| bases[((i * 7 + j * 3) % 4) as usize])
.collect()
})
.collect();
let seq_refs: Vec<&[u8]> = seqs.iter().map(|v| v.as_slice()).collect();
let (ef0, ef1) = direct_counts(&seq_refs);
let (gf0, gf1) = pipeline_counts(&seq_refs);
assert_eq!(gf0, ef0, "f0 wrong: expected {ef0}, got {gf0}");
assert_eq!(gf1, ef1, "f1 wrong: expected {ef1}, got {gf1}");
}
src/obikindex/src/index/kmer_index.rs
+5 -5
View File
@@ -159,10 +159,10 @@ impl KmerIndex {
/// Path of partition `i`'s raw directory (`partitions/part_{i:05}`) —
/// delegates to `crate::partition`, the Partition tier's own naming
/// primitive (mirrors `layer_dir` delegating to `crate::layer`).
/// `crate::algorithms::partitionner::PartitionRouter` reaches this same
/// directory only indirectly, through this method — it depends on
/// `KmerIndex`, not the other way around, even though both now live in
/// this crate — see `DevDocMD/implementation/partition_layer_cache.md`.
/// `obikindexer::algorithms::partitionner::PartitionRouter` reaches this
/// same directory only indirectly, through this method — `obikindexer`
/// depends on `KmerIndex`, not the other way around — see
/// `DevDocMD/implementation/partition_layer_cache.md`.
pub fn partition_dir(&self, i: usize) -> PathBuf {
crate::partition::partition_dir(&self.root_path, i)
}
@@ -232,7 +232,7 @@ impl KmerIndex {
/// Write `spectrums/{label}.json` from an already-computed kmer
/// spectrum (`f0`/`f1`/abundance histogram). Takes plain values rather
/// than `crate::algorithms::partitionner::KmerSpectrum` — `KmerIndex`
/// than `obikindexer::algorithms::partitionner::KmerSpectrum` — `KmerIndex`
/// is the data model, `PartitionRouter` the algorithm that depends on
/// it (see `DevDocMD/implementation/partition_layer_cache.md`), not
/// the other way around, and this is the only field of that type it
src/obikindex/src/layer/typed_layer.rs
+4 -4
View File
@@ -46,9 +46,9 @@ 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:
/// `crate::algorithms::partitionner::PartitionRouter`), read by whichever
/// `obikindexer::algorithms::partitionner::PartitionRouter`), read by whichever
/// algorithm dereplicates it (today:
/// `crate::algorithms::dereplicator::Dereplicator`). Naming this once
/// `obikindexer::algorithms::dereplicator::Dereplicator`). Naming this once
/// here, rather than in either algorithm submodule, is what lets the two
/// agree on the filename without depending on each other directly — see
/// `DevDocMD/implementation/partition_layer_cache.md`.
@@ -57,9 +57,9 @@ pub fn raw_superkmers_path(layer_dir: &Path) -> PathBuf {
}
/// Path of a layer's dereplicated superkmer file — written by
/// `crate::algorithms::dereplicator::Dereplicator`, read by whichever
/// `obikindexer::algorithms::dereplicator::Dereplicator`, read by whichever
/// algorithm counts kmer abundances from it (today:
/// `crate::algorithms::partitionner::PartitionRouter::count_kmer`) and,
/// `obikindexer::algorithms::partitionner::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}"))
src/obikindex/src/lib.rs
+4 -4
View File
@@ -3,11 +3,11 @@
//! section for the authoritative naming/scope discussion. Each tier is a
//! submodule: [`layer`] (was the `obilayeredmap` crate), [`partition`]
//! (was the `obikpartition` crate), [`index`] (this crate's original
//! content). [`algorithms`] (was the `obikpartitionner` and `obikderep`
//! crates) sits alongside the model: code that builds or transforms an
//! index's content rather than defining its shape.
//! content). Indexing-pipeline algorithms that build or transform an
//! index's content (the former `obikpartitionner`/`obikderep` crates) live
//! in the separate `obikindexer` crate, which depends on this one, not the
//! other way around.
pub mod algorithms;
pub mod index;
pub mod layer;
pub mod partition;