feat: enable index resumption and enforce directory creation

The command now supports reopening existing indexes instead of failing when the output file exists. Control flow branches between opening an existing index and constructing a new one, moving configuration setup exclusively to the creation path. Directory existence is enforced upfront with proper I/O error propagation. The --force flag retains its original semantics by removing the target directory before proceeding with a fresh build.
This commit is contained in:
Eric Coissac
2026-08-21 05:06:38 +02:00
parent abc51c2add
commit 5c1584967f
161 changed files with 5274 additions and 845 deletions
+6 -3
View File
@@ -5,15 +5,15 @@ edition = "2024"
[dependencies]
obikseq = { path = "../obikseq" }
obikpartition = { path = "../obikpartition" }
obitaxonomy = { path = "../obitaxonomy" }
obiskio = { path = "../obiskio" }
obisys = { path = "../obisys" }
obicompactvec = { path = "../obicompactvec" }
obilayeredmap = { path = "../obilayeredmap" }
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"
@@ -26,10 +26,13 @@ serde = { version = "1", features = ["derive"] }
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"
@@ -0,0 +1,155 @@
//! 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(())
}
@@ -0,0 +1,106 @@
//! 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(())
}
}
+9
View File
@@ -0,0 +1,9 @@
//! 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;
@@ -0,0 +1,80 @@
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(())
}
@@ -0,0 +1,126 @@
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)
}
@@ -0,0 +1,20 @@
//! 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};
@@ -0,0 +1,375 @@
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();
}
}
@@ -0,0 +1,131 @@
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}");
}
@@ -1,8 +1,8 @@
use std::path::Path;
use obicompactvec::{PersistentBitVecBuilder, PersistentCompactIntVecBuilder};
use obilayeredmap::meta::PartitionMeta;
use obilayeredmap::{layer_dir, IndexMode, OLMError};
use crate::layer::meta::PartitionMeta;
use crate::layer::{layer_dir, IndexMode, OLMError};
use obiskio::{SKError, SKResult};
// ── olm_to_sk ────────────────────────────────────────────────────────────────
@@ -1,10 +1,10 @@
use ndarray::Array2;
use obicompactvec::traits::{BitPartials, CountPartials};
use obilayeredmap::LayeredStore;
use crate::layer::LayeredStore;
use rayon::prelude::*;
use crate::error::{OKIError, OKIResult};
use crate::index::KmerIndex;
use crate::index::error::{OKIError, OKIResult};
use crate::index::kmer_index::KmerIndex;
// ── Public API ────────────────────────────────────────────────────────────────
@@ -3,9 +3,9 @@ use std::sync::atomic::{AtomicUsize, Ordering};
use rayon::prelude::*;
use crate::error::{OKIError, OKIResult};
use crate::index::KmerIndex;
use crate::KmerFilter;
use crate::index::error::{OKIError, OKIResult};
use crate::index::kmer_index::KmerIndex;
use crate::index::KmerFilter;
impl KmerIndex {
/// Write a CSV table of all indexed kmers to `out`.
@@ -1,10 +1,10 @@
use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix};
use obikseq::CanonicalKmer;
use obilayeredmap::{IndexMode, MphfLayer, OLMError};
use crate::layer::{IndexMode, MphfLayer, OLMError};
use obiskio::{SKError, SKResult, UnitigFileReader};
use crate::filter::{KmerFilter, passes_all};
use crate::index::KmerIndex;
use crate::index::filter::{KmerFilter, passes_all};
use crate::index::kmer_index::KmerIndex;
fn olm_to_sk(e: OLMError) -> SKError {
match e {
@@ -10,10 +10,10 @@ use obipipeline::{
use obidebruinj::GraphDeBruijn;
use obikseq::CanonicalKmer;
use obilayeredmap::{IndexMode, TypedLayer};
use crate::layer::{IndexMode, TypedLayer};
use obiskio::{SKError, SKResult};
use crate::common::olm_to_sk;
use crate::index::common::olm_to_sk;
// ── KmerGraphData ─────────────────────────────────────────────────────────────
@@ -5,14 +5,14 @@ use cacheline_ef::{CachelineEf, CachelineEfVec};
use epserde::prelude::*;
use obicompactvec::{PersistentCompactIntMatrix, PersistentCompactIntVec};
use obidebruinj::GraphDeBruijn;
use obilayeredmap::meta::PartitionMeta;
use obilayeredmap::{IndexMode, layer::TypedLayer};
use crate::layer::meta::PartitionMeta;
use crate::layer::{IndexMode, TypedLayer};
use obiskio::{SKError, SKFileMeta, SKFileReader};
use ptr_hash::{PtrHash, bucket_fn::CubicEps, hash::Xx64};
use crate::common::olm_to_sk;
use crate::graph_pipeline::{materialize_layer, write_graph_as_unitigs};
use crate::index::KmerIndex;
use crate::index::common::olm_to_sk;
use crate::index::graph_pipeline::{materialize_layer, write_graph_as_unitigs};
use crate::index::kmer_index::KmerIndex;
type Mphf = PtrHash<u64, CubicEps, CachelineEfVec<Vec<CachelineEf>>, Xx64, Vec<u8>>;
@@ -43,7 +43,7 @@ impl KmerIndex {
block_bits: u8,
) -> Result<usize, SKError> {
let layer0_dir = self.layer_dir(i, 0);
let dedup_path = obilayeredmap::dereplicated_superkmers_path(&layer0_dir);
let dedup_path = crate::layer::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 = obilayeredmap::dereplicated_superkmers_path(&layer0_dir);
let dedup = crate::layer::dereplicated_superkmers_path(&layer0_dir);
remove_if_exists(&SKFileMeta::sidecar_path(&dedup));
remove_if_exists(&dedup);
remove_if_exists(&layer0_dir.join("mphf1.bin"));
@@ -2,17 +2,17 @@ use std::collections::BTreeMap;
use std::fs;
use std::path::{Path, PathBuf};
use obilayeredmap::meta::PartitionMeta;
use crate::layer::meta::PartitionMeta;
use obisys::{Reporter, Stage, progress_bar};
use rayon::prelude::*;
use tracing::info;
use obikseq::{set_k, set_m};
use crate::common::load_meta;
use crate::error::{OKIError, OKIResult};
use crate::meta::{GenomeInfo, IndexConfig, IndexMeta};
use crate::state::{IndexState, SENTINEL_COUNTED, SENTINEL_INDEXED, SENTINEL_SCATTERED};
use crate::index::common::load_meta;
use crate::index::error::{OKIError, OKIResult};
use crate::index::meta::{GenomeInfo, IndexConfig, IndexMeta};
use crate::index::state::{IndexState, SENTINEL_COUNTED, SENTINEL_INDEXED, SENTINEL_SCATTERED};
pub struct KmerIndex {
pub(crate) root_path: PathBuf,
@@ -30,6 +30,7 @@ impl KmerIndex {
genome_info: Option<GenomeInfo>,
) -> OKIResult<Self> {
let root_path = path.as_ref().to_owned();
fs::create_dir_all(&root_path).map_err(OKIError::Io)?;
set_k(config.kmer_size);
set_m(config.minimizer_size);
let mut meta = IndexMeta::new(config);
@@ -130,7 +131,7 @@ impl KmerIndex {
pub fn with_counts(&self) -> bool {
self.meta.config.with_counts
}
pub fn evidence_mode(&self) -> &obilayeredmap::IndexMode {
pub fn evidence_mode(&self) -> &crate::layer::IndexMode {
&self.meta.config.evidence
}
/// Coarse evidence mode (`Exact`/`Approx`/`Hybrid`, no `b`/`z`
@@ -138,11 +139,11 @@ impl KmerIndex {
/// full `IndexMode`. Same discriminant as `Layer::evidence_kind`, so a
/// layer's own evidence kind can be compared directly against the
/// index's.
pub fn evidence_kind(&self) -> obilayeredmap::EvidenceKind {
pub fn evidence_kind(&self) -> crate::layer::EvidenceKind {
match self.meta.config.evidence {
obilayeredmap::IndexMode::Exact => obilayeredmap::EvidenceKind::Exact,
obilayeredmap::IndexMode::Approx { .. } => obilayeredmap::EvidenceKind::Approx,
obilayeredmap::IndexMode::Hybrid { .. } => obilayeredmap::EvidenceKind::Hybrid,
crate::layer::IndexMode::Exact => crate::layer::EvidenceKind::Exact,
crate::layer::IndexMode::Approx { .. } => crate::layer::EvidenceKind::Approx,
crate::layer::IndexMode::Hybrid { .. } => crate::layer::EvidenceKind::Hybrid,
}
}
pub fn block_bits(&self) -> u8 {
@@ -156,31 +157,32 @@ impl KmerIndex {
}
/// Path of partition `i`'s raw directory (`partitions/part_{i:05}`) —
/// 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`).
/// 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`.
pub fn partition_dir(&self, i: usize) -> PathBuf {
obikpartition::partition_dir(&self.root_path, i)
crate::partition::partition_dir(&self.root_path, i)
}
/// Path of partition `i`'s layered-index directory (`<partition>/index`).
pub fn index_dir(&self, i: usize) -> PathBuf {
obikpartition::index_dir(&self.root_path, i)
crate::partition::index_dir(&self.root_path, i)
}
/// Path of layer `l` within partition `i`'s layered index.
pub fn layer_dir(&self, i: usize, l: usize) -> PathBuf {
obilayeredmap::layer_dir(&self.index_dir(i), l)
crate::layer::layer_dir(&self.index_dir(i), l)
}
/// Partition `i`'s metadata (layer count, evidence mode). Returns
/// `obiskio::SKResult`, not `OKIResult` — matches the error convention
/// of the partition/layer-construction code below (moved here from
/// `obikpartitionner`, which predates `OKIError`); `?` still converts
/// it to `OKIResult` at any call site that needs one (`OKIError: From<SKError>`).
/// the former `obikpartitionner` crate, which predates `OKIError`);
/// `?` still converts it to `OKIResult` at any call site that needs
/// one (`OKIError: From<SKError>`).
pub fn partition_meta(&self, i: usize) -> obiskio::SKResult<PartitionMeta> {
load_meta(&self.index_dir(i), "partition_meta")
}
@@ -195,7 +197,7 @@ impl KmerIndex {
/// [`evidence_mode`](Self::evidence_mode): that one is the index-level
/// config, this one is per-partition ground truth (the two agree in
/// practice, but this is what `Layer::open` needs).
pub fn partition_mode(&self, i: usize) -> obiskio::SKResult<obilayeredmap::IndexMode> {
pub fn partition_mode(&self, i: usize) -> obiskio::SKResult<crate::layer::IndexMode> {
Ok(self.partition_meta(i)?.mode)
}
@@ -230,10 +232,11 @@ impl KmerIndex {
/// 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.
/// than `crate::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
/// actually uses.
pub fn write_spectrum(&self, f0: u64, f1: u64, counts: &BTreeMap<u32, u64>) -> OKIResult<()> {
let label = self
.meta
@@ -276,7 +279,7 @@ impl KmerIndex {
let pb = progress_bar("index", n as u64, "partitions");
let order: Vec<usize> = (0..n).collect();
let runner = crate::numa::PartitionRunner::new();
let runner = crate::index::numa::PartitionRunner::new();
runner
.run(
&order,
@@ -335,7 +338,7 @@ impl KmerIndex {
let n = self.n_partitions();
let order: Vec<usize> = (0..n).collect();
let pb = progress_bar("pack", n as u64, "partitions");
crate::numa::PartitionRunner::new().run(
crate::index::numa::PartitionRunner::new().run(
&order,
|i| -> OKIResult<()> {
let index_dir = self.index_dir(i);
@@ -1,9 +1,9 @@
use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix};
use obilayeredmap::{LayeredStore, open_data};
use crate::layer::{LayeredStore, open_data};
use obiskio::SKResult;
use crate::common::{load_meta, olm_to_sk};
use crate::index::KmerIndex;
use crate::index::common::{load_meta, olm_to_sk};
use crate::index::kmer_index::KmerIndex;
impl KmerIndex {
/// Open all count matrices for partition `part`, one per layer.
@@ -6,14 +6,14 @@ use std::path::Path;
use obisys::{Reporter, Stage, progress_bar, spinner};
use tracing::{debug, info};
use obilayeredmap::IndexMode;
use crate::layer::IndexMode;
use crate::error::{OKIError, OKIResult};
use crate::index::KmerIndex;
use crate::meta::{GenomeInfo, IndexMeta};
use crate::state::{IndexState, SENTINEL_INDEXED};
use crate::index::error::{OKIError, OKIResult};
use crate::index::kmer_index::KmerIndex;
use crate::index::meta::{GenomeInfo, IndexMeta};
use crate::index::state::{IndexState, SENTINEL_INDEXED};
pub use crate::merge_layer::MergeMode;
pub use crate::index::merge_layer::MergeMode;
// ── per-partition diagnostic record ──────────────────────────────────────────
@@ -217,7 +217,7 @@ impl KmerIndex {
let srcs = &srcs;
let evidence = &evidence;
let runner = crate::numa::PartitionRunner::new();
let runner = crate::index::numa::PartitionRunner::new();
let mut part_stats: Vec<PartStat> = Vec::with_capacity(n_partitions);
runner
@@ -20,12 +20,12 @@ use tracing::debug;
use obicompactvec::{PersistentBitMatrixBuilder, PersistentCompactIntMatrixBuilder};
use obikseq::CanonicalKmer;
use obilayeredmap::{IndexMode, TypedLayer, LayeredMap, MphfOnly, layer_dir};
use crate::layer::{IndexMode, TypedLayer, LayeredMap, MphfOnly, layer_dir};
use obiskio::{SKError, SKResult, UnitigFileReader};
use crate::common::{ColBuilder, load_meta, olm_to_sk};
use crate::graph_pipeline::{build_graph, materialize_layer};
use crate::index::KmerIndex;
use crate::index::common::{ColBuilder, load_meta, olm_to_sk};
use crate::index::graph_pipeline::{build_graph, materialize_layer};
use crate::index::kmer_index::KmerIndex;
mod src_layer;
@@ -2,10 +2,10 @@ use std::path::Path;
use obicompactvec::{MatrixGroupOps, PersistentBitMatrix, PersistentCompactIntMatrix};
use obikseq::CanonicalKmer;
use obilayeredmap::MphfOnly;
use crate::layer::MphfOnly;
use obiskio::{SKError, SKResult};
use crate::common::olm_to_sk;
use crate::index::common::olm_to_sk;
use super::MergeMode;
@@ -3,7 +3,7 @@ use std::fs;
use std::io;
use std::path::Path;
use obilayeredmap::IndexMode;
use crate::layer::IndexMode;
use serde::{Deserialize, Serialize};
pub const META_FILENAME: &str = "index.meta";
+36
View File
@@ -0,0 +1,36 @@
pub mod error;
pub mod meta;
pub mod predicate;
pub mod state;
mod common;
mod distance;
mod dump;
mod dump_layer;
pub mod filter;
mod graph_pipeline;
mod kmer_index;
mod index_layer;
mod matrix_store;
mod merge;
mod merge_layer;
mod numa;
mod query_layer;
mod rebuild;
mod rebuild_layer;
mod reindex;
mod select;
mod select_layer;
mod stats;
pub use error::{OKIError, OKIResult};
pub use distance::{DistanceMetric, DistanceOutput};
pub use filter::{GroupQuorumFilter, KmerFilter, passes_all};
pub use kmer_index::KmerIndex;
pub use merge_layer::MergeMode;
pub use meta::{validate_label, GenomeInfo, IndexConfig, IndexMeta, META_FILENAME};
pub use predicate::{GroupFilterParams, MetaPred};
pub use query_layer::{KmerDesc, QueryHit, QueryStats};
pub use select_layer::{AggOp, OutputCol};
pub use state::{IndexState, SENTINEL_COUNTED, SENTINEL_INDEXED, SENTINEL_SCATTERED};
pub use stats::IndexBitsPerKmer;
pub use numa::PartitionRunner;
@@ -1,9 +1,9 @@
use std::collections::HashMap;
use crate::GroupQuorumFilter;
use crate::index::GroupQuorumFilter;
use obitaxonomy::{TaxPath, TaxPattern};
use crate::meta::{GenomeInfo, IndexMeta};
use crate::index::meta::{GenomeInfo, IndexMeta};
// ── Operator ──────────────────────────────────────────────────────────────────
@@ -3,10 +3,10 @@ use std::path::Path;
use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix};
use obikseq::CanonicalKmer;
use obilayeredmap::{IndexMode, MphfLayer, OLMError};
use crate::layer::{IndexMode, MphfLayer, OLMError};
use obiskio::{SKError, SKResult};
use crate::index::KmerIndex;
use crate::index::kmer_index::KmerIndex;
fn olm_to_sk(e: OLMError) -> SKError {
match e {
@@ -1,13 +1,13 @@
use std::path::Path;
use crate::{KmerFilter, MergeMode};
use crate::index::{KmerFilter, MergeMode};
use obisys::{Reporter, Stage, progress_bar};
use tracing::info;
use crate::error::{OKIError, OKIResult};
use crate::index::KmerIndex;
use crate::meta::IndexMeta;
use crate::state::IndexState;
use crate::index::error::{OKIError, OKIResult};
use crate::index::kmer_index::KmerIndex;
use crate::index::meta::IndexMeta;
use crate::index::state::IndexState;
impl KmerIndex {
/// Rebuild `src` into a new compact single-layer index at `output`.
@@ -62,7 +62,7 @@ impl KmerIndex {
let block_bits = meta.config.block_bits;
let order: Vec<usize> = (0..n_partitions).collect();
let runner = crate::numa::PartitionRunner::new();
let runner = crate::index::numa::PartitionRunner::new();
runner.run(
&order,
|i| dst_partition.rebuild_partition(src, i, filters, mode, n_genomes, block_bits),
@@ -6,15 +6,15 @@ use obicompactvec::{
};
use obidebruinj::GraphDeBruijn;
use obikseq::CanonicalKmer;
use obilayeredmap::meta::PartitionMeta;
use obilayeredmap::{IndexMode, MphfLayer, layer_dir};
use crate::layer::meta::PartitionMeta;
use crate::layer::{IndexMode, MphfLayer, layer_dir};
use obiskio::{SKError, SKResult, UnitigFileReader};
use crate::common::{load_meta, olm_to_sk};
use crate::filter::KmerFilter;
use crate::graph_pipeline::materialize_layer;
use crate::merge_layer::{MergeMode, SrcLayerData};
use crate::index::KmerIndex;
use crate::index::common::{load_meta, olm_to_sk};
use crate::index::filter::KmerFilter;
use crate::index::graph_pipeline::materialize_layer;
use crate::index::merge_layer::{MergeMode, SrcLayerData};
use crate::index::kmer_index::KmerIndex;
// ── Builders — pair matrix builder + column builders for one mode ─────────────
@@ -1,18 +1,18 @@
use obilayeredmap::{IndexMode, layer::TypedLayer};
use crate::layer::{IndexMode, TypedLayer};
use obisys::{Reporter, Stage, progress_bar};
use std::fs;
use std::path::Path;
use tracing::info;
use crate::error::{OKIError, OKIResult};
use crate::index::KmerIndex;
use crate::state::IndexState;
use crate::index::error::{OKIError, OKIResult};
use crate::index::kmer_index::KmerIndex;
use crate::index::state::IndexState;
const EVIDENCE_FILE: &str = "evidence.bin";
const FINGERPRINT_FILE: &str = "fingerprint.bin";
const UNITIG_IDX_FILE: &str = "unitigs.bin.idx";
fn olm_to_oki(e: obilayeredmap::OLMError) -> OKIError {
fn olm_to_oki(e: crate::layer::OLMError) -> OKIError {
OKIError::InvalidInput(e.to_string())
}
@@ -44,7 +44,7 @@ impl KmerIndex {
let pb = progress_bar("reindex", n as u64, "partitions");
let order: Vec<usize> = (0..n).collect();
let runner = crate::numa::PartitionRunner::new();
let runner = crate::index::numa::PartitionRunner::new();
runner.run(
&order,
|i| {
@@ -1,13 +1,13 @@
use std::path::Path;
use crate::OutputCol;
use crate::index::OutputCol;
use obisys::{Reporter, Stage, progress_bar};
use tracing::info;
use crate::error::{OKIError, OKIResult};
use crate::index::KmerIndex;
use crate::meta::{GenomeInfo, IndexMeta};
use crate::state::IndexState;
use crate::index::error::{OKIError, OKIResult};
use crate::index::kmer_index::KmerIndex;
use crate::index::meta::{GenomeInfo, IndexMeta};
use crate::index::state::IndexState;
impl KmerIndex {
/// Create a new index at `output` by projecting/aggregating the genome columns
@@ -56,7 +56,7 @@ impl KmerIndex {
let pb = progress_bar("select", n_partitions as u64, "partitions");
let order: Vec<usize> = (0..n_partitions).collect();
let runner = crate::numa::PartitionRunner::new();
let runner = crate::index::numa::PartitionRunner::new();
runner
.run(
&order,
@@ -111,7 +111,7 @@ impl KmerIndex {
let pb = progress_bar("select", n_partitions as u64, "partitions");
let order: Vec<usize> = (0..n_partitions).collect();
let runner = crate::numa::PartitionRunner::new();
let runner = crate::index::numa::PartitionRunner::new();
runner
.run(
&order,
@@ -6,10 +6,10 @@ use obicompactvec::{
ColGroup, MatrixGroupOps, PersistentBitMatrix, PersistentBitMatrixBuilder,
PersistentCompactIntMatrix, PersistentCompactIntMatrixBuilder,
};
use obilayeredmap::OLMError;
use crate::layer::OLMError;
use obiskio::{SKError, SKResult};
use crate::index::KmerIndex;
use crate::index::kmer_index::KmerIndex;
// ── AggOp ─────────────────────────────────────────────────────────────────────
@@ -1,6 +1,6 @@
use std::path::Path;
use crate::meta::META_FILENAME;
use crate::index::meta::META_FILENAME;
pub const SENTINEL_SCATTERED: &str = "scatter.done";
pub const SENTINEL_COUNTED: &str = "count.done";
@@ -5,8 +5,8 @@ use obicompactvec::{LayerMeta, PersistentBitMatrix, PersistentCompactIntMatrix};
use obicompactvec::traits::ColumnWeights;
use rayon::prelude::*;
use crate::error::OKIResult;
use crate::index::KmerIndex;
use crate::index::error::OKIResult;
use crate::index::kmer_index::KmerIndex;
/// Bits per kmer broken down by index component.
pub struct IndexBitsPerKmer {
@@ -153,12 +153,12 @@ impl KmerIndex {
if this_layer_dir.join("counts").exists()
&& !this_layer_dir.join("presence").exists()
{
match obilayeredmap::open_data::<PersistentCompactIntMatrix>(&index_dir, l) {
match crate::layer::open_data::<PersistentCompactIntMatrix>(&index_dir, l) {
Ok(m) => Box::new(m),
Err(_) => continue,
}
} else {
match obilayeredmap::open_data::<PersistentBitMatrix>(&index_dir, l) {
match crate::layer::open_data::<PersistentBitMatrix>(&index_dir, l) {
Ok(m) => Box::new(m),
Err(_) => continue,
}
@@ -1,5 +1,5 @@
use super::*;
use crate::meta::IndexConfig;
use crate::index::meta::IndexConfig;
// ── QueryStats::AddAssign ───────────────────────────────────────────────────
@@ -50,7 +50,7 @@ fn query_partition_with_missing_index_dir_returns_default_stats() {
minimizer_size: 9,
n_bits: 2,
with_counts: false,
evidence: obilayeredmap::IndexMode::Exact,
evidence: crate::layer::IndexMode::Exact,
block_bits: 0,
};
let index = KmerIndex::create(tmp.path().join("idx"), config, None).expect("create index");
@@ -78,7 +78,7 @@ fn query_partition_with_empty_kmers_is_a_noop() {
minimizer_size: 9,
n_bits: 2,
with_counts: false,
evidence: obilayeredmap::IndexMode::Exact,
evidence: crate::layer::IndexMode::Exact,
block_bits: 0,
};
let index = KmerIndex::create(tmp.path().join("idx"), config, None).expect("create index");
+174
View File
@@ -0,0 +1,174 @@
//! [`Layer`] — the format-erased layer handle every caller outside this
//! crate should reach for. `TypedLayer<D>` (`layer.rs`) is monomorphic: a
//! `Vec<TypedLayer<D>>` needs one `D` fixed at compile time for every
//! element, which real partitions don't respect (`with_counts` can differ
//! layer to layer across merges; see
//! `DevDocMD/implementation/partition_layer_cache.md`). `Layer` is the
//! sum type that lets a partition's layers be held together regardless —
//! `Count`/`Presence` each wrap the one concrete `TypedLayer<D>` that
//! content implies, chosen by [`Layer::open`]'s own disk probe.
//!
//! This is exactly the shape `obikphylo::siblings::cache::Mat` used to
//! reimplement locally, minus its one sibling-specific method
//! (`iter_minorants_batch`, which stays an extension trait over there —
//! `obikindex::layer` has no business knowing about sibling annexes).
//!
//! `Layer` is meant to eventually represent a layer's *whole* life —
//! empty shell, under construction, ready to read — not just the
//! ready-to-read state, so the same type follows a layer from creation
//! through querying instead of construction living as disconnected free
//! functions elsewhere that happen to write into the same directory (see
//! `DevDocMD/implementation/partition_layer_cache.md`). [`Layer::Empty`]
//! is the first step: a directory and nothing else, able to hand out the
//! paths a builder needs, not yet able to build anything itself.
use std::path::{Path, PathBuf};
use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix};
use obikseq::CanonicalKmer;
use crate::layer::error::OLMResult;
use crate::layer::typed_layer::{LayerContent, TypedLayer, COUNTS_DIR, PRESENCE_DIR};
use crate::layer::meta::IndexMode;
use crate::layer::mphf_layer::{EvidenceKind, EVIDENCE_FILE, FINGERPRINT_FILE, MPHF_FILE, UNITIGS_FILE};
/// One layer, at any point in its life — see the module docs. Only
/// [`Empty`](Layer::Empty) and the two ready-to-read states
/// (`Count`/`Presence`) exist so far; states in between (unitigs written,
/// MPHF built, evidence built, matrix not yet built) are not represented
/// yet.
pub enum Layer {
/// Just a directory — nothing built yet. Every method below other than
/// the path accessors panics on this variant: calling them means the
/// caller assumed a layer was ready when it wasn't, an implementation
/// error to surface loudly, not paper over with a default value.
Empty { dir: PathBuf },
Count(TypedLayer<PersistentCompactIntMatrix>),
Presence(TypedLayer<PersistentBitMatrix>),
}
impl Layer {
/// An empty shell at `dir` — creates the directory (if it doesn't
/// already exist) but nothing inside it. The starting state for
/// building a new layer.
pub fn create(dir: &Path) -> std::io::Result<Self> {
std::fs::create_dir_all(dir)?;
Ok(Layer::Empty { dir: dir.to_owned() })
}
/// Open one layer, auto-detecting count vs. presence from what is
/// actually on disk (`counts/` present and wanted, else presence) — the
/// single source of truth every caller that opens a layer's own matrix
/// should share, so two callers can never disagree about which format a
/// given layer is. Dense vs. sparse presence storage is
/// `PersistentBitMatrix::open`'s own concern (it's a 4-way
/// `Columnar`/`Packed`/`Sparse`/`Implicit` enum internally), not decided
/// here.
pub fn open(layer_dir: &Path, mode: &IndexMode, with_counts: bool) -> OLMResult<Self> {
if with_counts && layer_dir.join(COUNTS_DIR).exists() {
return TypedLayer::<PersistentCompactIntMatrix>::open(layer_dir, mode).map(Layer::Count);
}
TypedLayer::<PersistentBitMatrix>::open(layer_dir, mode).map(Layer::Presence)
}
// ── Paths — only meaningful before anything is built ───────────────
//
// Once a layer is `Count`/`Presence`, its caller already knows the
// directory (it had to pass it to `open`) — these exist for builder
// code holding an `Empty` layer, so the `mphf.bin`/`unitigs.bin`/…
// naming stays defined once, here, rather than re-declared as
// string literals at every call site that writes into a layer
// directory (the same duplication `layer_dir`/`index_dir` fixed one
// level up — see `DevDocMD/implementation/partition_layer_cache.md`).
/// This layer's own directory. Panics on `Count`/`Presence` — by that
/// point the caller already has the directory it opened with; asking
/// again here would mean it lost track of its own state.
pub fn dir(&self) -> &Path {
match self {
Layer::Empty { dir } => dir,
_ => panic!("Layer::dir() only available on Empty — caller already has this path"),
}
}
pub fn mphf_path(&self) -> PathBuf { self.dir().join(MPHF_FILE) }
pub fn unitigs_path(&self) -> PathBuf { self.dir().join(UNITIGS_FILE) }
pub fn evidence_path(&self) -> PathBuf { self.dir().join(EVIDENCE_FILE) }
pub fn fingerprint_path(&self) -> PathBuf { self.dir().join(FINGERPRINT_FILE) }
pub fn counts_dir(&self) -> PathBuf { self.dir().join(COUNTS_DIR) }
pub fn presence_dir(&self) -> PathBuf { self.dir().join(PRESENCE_DIR) }
// ── Ready-only surface ───────────────────────────────────────────────
pub fn content(&self) -> LayerContent {
match self {
Layer::Count(_) => LayerContent::Count,
Layer::Presence(_) => LayerContent::Presence,
Layer::Empty { .. } => panic!("Layer::content() called on an Empty layer"),
}
}
pub fn evidence_kind(&self) -> EvidenceKind {
match self {
Layer::Count(l) => l.evidence_kind(),
Layer::Presence(l) => l.evidence_kind(),
Layer::Empty { .. } => panic!("Layer::evidence_kind() called on an Empty layer"),
}
}
pub fn n(&self) -> usize {
match self {
Layer::Count(l) => l.n(),
Layer::Presence(l) => l.n(),
Layer::Empty { .. } => panic!("Layer::n() called on an Empty layer"),
}
}
pub fn find_slot(&self, kmer: CanonicalKmer) -> Option<usize> {
match self {
Layer::Count(l) => l.find_slot(kmer),
Layer::Presence(l) => l.find_slot(kmer),
Layer::Empty { .. } => panic!("Layer::find_slot() called on an Empty layer"),
}
}
/// Raw MPHF batch lookup: kmer → slot, no membership check — for
/// callers that already know every kmer is a member of *this* layer, so
/// the evidence check `find_slot`/`find` would perform is redundant.
pub fn index_batch(&self, kmers: &[CanonicalKmer]) -> Vec<usize> {
match self {
Layer::Count(l) => l.index_batch(kmers),
Layer::Presence(l) => l.index_batch(kmers),
Layer::Empty { .. } => panic!("Layer::index_batch() called on an Empty layer"),
}
}
pub fn n_cols(&self) -> usize {
match self {
Layer::Count(l) => l.n_cols(),
Layer::Presence(l) => l.n_cols(),
Layer::Empty { .. } => panic!("Layer::n_cols() called on an Empty layer"),
}
}
/// Batch, genome-major "carries" for a set of `slots` — `out[g][i]` =
/// whether genome `g` (0..`out.len()`) carries `slots[i]`. `out` must
/// have one entry per genome column, each resized to `slots.len()`.
/// `Count` needs one intermediate `Vec<Vec<u32>>` fetch (the underlying
/// store only has an int sub-matrix, not a bool one), converted to
/// presence (`!= 0`) in place.
pub fn fill_sub_matrix_carries(&self, slots: &[usize], out: &mut [Vec<bool>]) {
match self {
Layer::Presence(l) => l.fill_sub_matrix(slots, out),
Layer::Count(l) => {
let mut counts: Vec<Vec<u32>> = out.iter().map(|_| Vec::new()).collect();
l.fill_sub_matrix(slots, &mut counts);
for (o, c) in out.iter_mut().zip(counts.iter()) {
o.clear();
o.extend(c.iter().map(|&v| v != 0));
}
}
Layer::Empty { .. } => panic!("Layer::fill_sub_matrix_carries() called on an Empty layer"),
}
}
}
+50
View File
@@ -0,0 +1,50 @@
use std::fmt;
use std::io;
#[derive(Debug)]
pub enum OLMError {
Io(io::Error),
Json(serde_json::Error),
Mphf(String),
InvalidLayer(String),
}
pub type OLMResult<T> = Result<T, OLMError>;
impl fmt::Display for OLMError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
OLMError::Io(e) => write!(f, "I/O error: {e}"),
OLMError::Json(e) => write!(f, "JSON error: {e}"),
OLMError::Mphf(s) => write!(f, "MPHF error: {s}"),
OLMError::InvalidLayer(s) => write!(f, "invalid layer: {s}"),
}
}
}
impl std::error::Error for OLMError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
OLMError::Io(e) => Some(e),
OLMError::Json(e) => Some(e),
_ => None,
}
}
}
impl From<io::Error> for OLMError {
fn from(e: io::Error) -> Self { OLMError::Io(e) }
}
impl From<serde_json::Error> for OLMError {
fn from(e: serde_json::Error) -> Self { OLMError::Json(e) }
}
impl From<obiskio::SKError> for OLMError {
fn from(e: obiskio::SKError) -> Self {
match e {
obiskio::SKError::Io(io_err) => OLMError::Io(io_err),
other => OLMError::InvalidLayer(other.to_string()),
}
}
}
+60
View File
@@ -0,0 +1,60 @@
// u32 per MPHF slot: bits [31:7] = chunk_id (25 bits), bits [6:0] = rank (7 bits).
use std::fs::File;
use std::io::{BufWriter, Write};
use std::path::Path;
use memmap2::Mmap;
use crate::layer::error::{OLMError, OLMResult};
pub struct Evidence {
mmap: Mmap,
}
impl Evidence {
pub fn open(path: &Path) -> OLMResult<Self> {
let f = File::open(path)?;
let mmap = unsafe { Mmap::map(&f)? };
Ok(Self { mmap })
}
#[inline]
pub fn decode(&self, slot: usize) -> (u32, u8) {
let off = slot * 4;
let raw = u32::from_le_bytes(self.mmap[off..off + 4].try_into().unwrap());
(raw >> 7, (raw & 0x7F) as u8)
}
pub fn len(&self) -> usize {
self.mmap.len() / 4
}
}
#[inline]
pub fn encode(chunk_id: u32, rank: u8) -> u32 {
(chunk_id << 7) | (rank as u32 & 0x7F)
}
pub struct EvidenceWriter {
buf: Vec<u32>,
}
impl EvidenceWriter {
pub fn new(n_slots: usize) -> Self {
Self { buf: vec![0u32; n_slots] }
}
#[inline]
pub fn set(&mut self, slot: usize, chunk_id: u32, rank: u8) {
self.buf[slot] = encode(chunk_id, rank);
}
pub fn write(self, path: &Path) -> OLMResult<()> {
let mut f = BufWriter::new(File::create(path)?);
for v in self.buf {
f.write_all(&v.to_le_bytes()).map_err(OLMError::Io)?;
}
Ok(())
}
}
+151
View File
@@ -0,0 +1,151 @@
// Packed B-bit fingerprint vector, one entry per MPHF slot.
//
// File format (fingerprint.bin):
// magic: b"FPVF" (4 bytes)
// b: u8 (bits per fingerprint, 1..=64)
// padding: [0u8; 3]
// n: u64 LE (number of slots)
// data: packed bits, ceil(n*b/8) bytes, Lsb0 order
use std::fs::File;
use std::io::{BufWriter, Write};
use std::path::Path;
use bitvec::prelude::*;
use memmap2::Mmap;
use crate::layer::error::{OLMError, OLMResult};
const MAGIC: &[u8; 4] = b"FPVF";
const HEADER: usize = 16;
// ── Reader ────────────────────────────────────────────────────────────────────
pub struct FingerprintVec {
_mmap: Mmap,
bits: &'static BitSlice<u8, Lsb0>,
n: usize,
b: u8,
mask: u64,
}
impl FingerprintVec {
pub fn open(path: &Path) -> OLMResult<Self> {
let f = File::open(path)?;
let mmap = unsafe { Mmap::map(&f)? };
if mmap.len() < HEADER || &mmap[..4] != MAGIC {
return Err(OLMError::InvalidLayer("bad fingerprint magic".into()));
}
let b = mmap[4];
if b == 0 || b > 64 {
return Err(OLMError::InvalidLayer("invalid fingerprint width".into()));
}
let n = u64::from_le_bytes(mmap[8..16].try_into().unwrap()) as usize;
let mask: u64 = if b == 64 { u64::MAX } else { (1u64 << b) - 1 };
// SAFETY: the mmap lives as long as Self (kept via _mmap); data is read-only.
let data: &'static [u8] = unsafe {
std::slice::from_raw_parts(mmap[HEADER..].as_ptr(), (n * b as usize + 7) / 8)
};
let bits = BitSlice::<u8, Lsb0>::from_slice(data);
Ok(Self { _mmap: mmap, bits, n, b, mask })
}
#[inline]
pub fn get(&self, slot: usize) -> u64 {
debug_assert!(slot < self.n);
let lo = slot * self.b as usize;
self.bits[lo .. lo + self.b as usize].load_le::<u64>()
}
#[inline]
pub fn matches(&self, slot: usize, fingerprint: u64) -> bool {
self.get(slot) == (fingerprint & self.mask)
}
pub fn n(&self) -> usize { self.n }
pub fn b(&self) -> u8 { self.b }
}
// ── Writer ────────────────────────────────────────────────────────────────────
pub struct FingerprintVecWriter {
buf: Vec<u8>,
n: usize,
b: u8,
}
impl FingerprintVecWriter {
pub fn new(n: usize, b: u8) -> Self {
assert!(b > 0 && b <= 64, "fingerprint width must be 1..=64");
let data_bytes = (n * b as usize + 7) / 8;
Self { buf: vec![0u8; data_bytes], n, b }
}
#[inline]
pub fn set(&mut self, slot: usize, fingerprint: u64) {
debug_assert!(slot < self.n);
let lo = slot * self.b as usize;
let bits = BitSlice::<u8, Lsb0>::from_slice_mut(&mut self.buf);
bits[lo .. lo + self.b as usize].store_le(fingerprint);
}
pub fn write(self, path: &Path) -> OLMResult<()> {
let mut f = BufWriter::new(File::create(path).map_err(OLMError::Io)?);
f.write_all(MAGIC).map_err(OLMError::Io)?;
f.write_all(&[self.b, 0, 0, 0]).map_err(OLMError::Io)?;
f.write_all(&(self.n as u64).to_le_bytes()).map_err(OLMError::Io)?;
f.write_all(&self.buf).map_err(OLMError::Io)?;
Ok(())
}
}
// ── Tests ─────────────────────────────────────────────────────────────────────
#[cfg(test)]
mod tests {
use super::*;
use tempfile::tempdir;
fn roundtrip(n: usize, b: u8, values: &[u64]) {
let dir = tempdir().unwrap();
let path = dir.path().join("fp.bin");
let mask: u64 = if b == 64 { u64::MAX } else { (1u64 << b) - 1 };
let mut w = FingerprintVecWriter::new(n, b);
for (i, &v) in values.iter().enumerate() {
w.set(i, v);
}
w.write(&path).unwrap();
let r = FingerprintVec::open(&path).unwrap();
assert_eq!(r.n(), n);
assert_eq!(r.b(), b);
for (i, &v) in values.iter().enumerate() {
assert_eq!(r.get(i), v & mask, "slot {i} b={b}");
}
}
#[test]
fn roundtrip_b1() { roundtrip(8, 1, &[1, 0, 1, 1, 0, 0, 1, 0]); }
#[test]
fn roundtrip_b8() { roundtrip(4, 8, &[0, 127, 200, 255]); }
#[test]
fn roundtrip_b16() { roundtrip(3, 16, &[0, 0xABCD, 0xFFFF]); }
#[test]
fn roundtrip_b7() { roundtrip(5, 7, &[0, 63, 127, 1, 42]); }
#[test]
fn roundtrip_b13_unaligned() {
let vals: Vec<u64> = (0..20).map(|i| (i * 317) % (1 << 13)).collect();
roundtrip(20, 13, &vals);
}
#[test]
fn matches_returns_true_for_exact() {
let dir = tempdir().unwrap();
let path = dir.path().join("fp.bin");
let mut w = FingerprintVecWriter::new(4, 8);
w.set(0, 42); w.set(1, 0); w.set(2, 255); w.set(3, 17);
w.write(&path).unwrap();
let r = FingerprintVec::open(&path).unwrap();
assert!( r.matches(0, 42));
assert!(!r.matches(0, 43));
assert!( r.matches(2, 255));
}
}
+100
View File
@@ -0,0 +1,100 @@
use ndarray::{Array1, Array2};
use rayon::prelude::*;
use obicompactvec::traits::{BitPartials, ColumnWeights, CountPartials};
/// A store that aggregates a `Vec<S>` — one entry per layer (within a partition)
/// or one entry per partition.
///
/// Blanket impls of `ColumnWeights`, `CountPartials`, and `BitPartials` propagate
/// automatically: `LayeredStore<LayeredStore<S>>` implements the same traits as
/// `LayeredStore<S>`, giving the partitioned level for free.
pub struct LayeredStore<S>(pub Vec<S>);
impl<S> LayeredStore<S> {
pub fn new(layers: Vec<S>) -> Self { Self(layers) }
pub fn layers(&self) -> &[S] { &self.0 }
pub fn n_layers(&self) -> usize { self.0.len() }
pub fn is_empty(&self) -> bool { self.0.is_empty() }
}
// ── ColumnWeights ─────────────────────────────────────────────────────────────
impl<S: ColumnWeights> ColumnWeights for LayeredStore<S> {
fn col_weights(&self) -> Array1<u64> {
self.0.par_iter()
.map(|s| s.col_weights())
.reduce_with(|a, b| a + b)
.unwrap_or_else(|| Array1::zeros(0))
}
}
// ── CountPartials ─────────────────────────────────────────────────────────────
impl<S: CountPartials> CountPartials for LayeredStore<S> {
fn partial_bray(&self) -> Array2<u64> {
self.0.par_iter()
.map(|s| s.partial_bray())
.reduce_with(|a, b| a + b)
.unwrap()
}
fn partial_euclidean(&self) -> Array2<f64> {
self.0.par_iter()
.map(|s| s.partial_euclidean())
.reduce_with(|a, b| a + b)
.unwrap()
}
fn partial_threshold_jaccard(&self, threshold: u32) -> (Array2<u64>, Array2<u64>) {
self.0.par_iter()
.map(|s| s.partial_threshold_jaccard(threshold))
.reduce_with(|(ai, au), (bi, bu)| (ai + bi, au + bu))
.unwrap()
}
fn partial_relfreq_bray(&self, global: &Array1<u64>) -> Array2<f64> {
self.0.par_iter()
.map(|s| s.partial_relfreq_bray(global))
.reduce_with(|a, b| a + b)
.unwrap()
}
fn partial_relfreq_euclidean(&self, global: &Array1<u64>) -> Array2<f64> {
self.0.par_iter()
.map(|s| s.partial_relfreq_euclidean(global))
.reduce_with(|a, b| a + b)
.unwrap()
}
fn partial_hellinger(&self, global: &Array1<u64>) -> Array2<f64> {
self.0.par_iter()
.map(|s| s.partial_hellinger(global))
.reduce_with(|a, b| a + b)
.unwrap()
}
}
// ── BitPartials ───────────────────────────────────────────────────────────────
impl<S: BitPartials> BitPartials for LayeredStore<S> {
fn partial_jaccard(&self) -> (Array2<u64>, Array2<u64>) {
self.0.par_iter()
.map(|s| s.partial_jaccard())
.reduce_with(|(ai, au), (bi, bu)| (ai + bi, au + bu))
.unwrap()
}
fn partial_hamming(&self) -> Array2<u64> {
self.0.par_iter()
.map(|s| s.partial_hamming())
.reduce_with(|a, b| a + b)
.unwrap()
}
}
// ── Tests ─────────────────────────────────────────────────────────────────────
#[cfg(test)]
#[path = "tests/layered_store.rs"]
mod tests;
+128
View File
@@ -0,0 +1,128 @@
use std::collections::HashMap;
use std::fs;
use std::path::{Path, PathBuf};
use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix, StorageKind};
use obikseq::CanonicalKmer;
use obiskio::{UnitigFileWriter, DEFAULT_BLOCK_BITS};
use crate::layer::error::{OLMError, OLMResult};
use crate::layer::typed_layer::{layer_dir, Hit, TypedLayer, LayerContent, LayerData};
use crate::layer::meta::{IndexMode, PartitionMeta};
use crate::layer::mphf_layer::EvidenceKind;
/// Layered kmer index for a single partition.
///
/// Each layer covers a disjoint kmer set. Queries probe layers in order;
/// the first match wins. Adding a dataset appends a new layer without
/// rebuilding existing ones.
pub struct LayeredMap<D: LayerData = ()> {
root: PathBuf,
meta: PartitionMeta,
layers: Vec<TypedLayer<D>>,
}
// ── Common methods ────────────────────────────────────────────────────────────
impl<D: LayerData> LayeredMap<D> {
/// Open an existing layered index at `root`.
/// The mode is read once from `PartitionMeta` and applied to all layers.
pub fn open(root: &Path) -> OLMResult<Self> {
let meta = PartitionMeta::load(root)?;
let layers = (0..meta.n_layers)
.map(|i| TypedLayer::<D>::open(&layer_dir(root, i), &meta.mode))
.collect::<OLMResult<Vec<_>>>()?;
Ok(Self { root: root.to_owned(), meta, layers })
}
/// Create a new, empty layered index at `root` with the given mode.
pub fn create(root: &Path, mode: IndexMode) -> OLMResult<Self> {
fs::create_dir_all(root)?;
let meta = PartitionMeta::new(mode);
meta.save(root)?;
Ok(Self { root: root.to_owned(), meta, layers: Vec::new() })
}
pub fn n_layers(&self) -> usize { self.layers.len() }
pub fn layer(&self, i: usize) -> &TypedLayer<D> { &self.layers[i] }
pub fn mode(&self) -> &IndexMode { &self.meta.mode }
/// TypedLayer `i`'s content (`Count`/`Presence`) — a lightweight disk probe,
/// no MPHF/matrix opened. For picking a `D` before committing to
/// [`TypedLayer::open`], or for reporting/diagnostics over an already-open
/// `LayeredMap` without re-opening every layer under a different `D`.
pub fn detect_layer_content(&self, i: usize) -> LayerContent {
LayerContent::detect(&layer_dir(&self.root, i))
}
/// TypedLayer `i`'s storage format — a lightweight disk probe, no
/// MPHF/matrix opened. See [`detect_layer_content`](Self::detect_layer_content).
pub fn detect_layer_storage(&self, i: usize) -> OLMResult<StorageKind> {
let dir = layer_dir(&self.root, i);
match LayerContent::detect(&dir) {
LayerContent::Count => PersistentCompactIntMatrix::detect_storage(&dir).map_err(OLMError::Io),
LayerContent::Presence => PersistentBitMatrix::detect_storage(&dir).map_err(OLMError::Io),
}
}
/// TypedLayer `i`'s evidence mode — a lightweight disk probe, no
/// MPHF/matrix opened. See [`detect_layer_content`](Self::detect_layer_content).
pub fn detect_layer_evidence(&self, i: usize) -> OLMResult<EvidenceKind> {
EvidenceKind::detect(&layer_dir(&self.root, i))
}
/// Query `kmer` across all layers. Returns `(layer_index, Hit)` on match.
pub fn query(&self, kmer: CanonicalKmer) -> Option<(usize, Hit<D::Item>)> {
self.layers
.iter()
.enumerate()
.find_map(|(i, layer)| layer.query(kmer).map(|hit| (i, hit)))
}
pub fn next_layer_writer(&self) -> OLMResult<UnitigFileWriter> {
let dir = layer_dir(&self.root, self.layers.len());
TypedLayer::<D>::unitig_writer(&dir)
}
fn append_layer(&mut self) -> OLMResult<()> {
let i = self.layers.len();
let dir = layer_dir(&self.root, i);
self.layers.push(TypedLayer::<D>::open(&dir, &self.meta.mode)?);
self.meta.n_layers = self.layers.len();
self.meta.save(&self.root)?;
Ok(())
}
}
// ── Mode 1 — set membership ───────────────────────────────────────────────────
impl LayeredMap<()> {
pub fn push_layer(&mut self) -> OLMResult<usize> {
let i = self.layers.len();
let dir = layer_dir(&self.root, i);
TypedLayer::<()>::build(&dir, DEFAULT_BLOCK_BITS, &self.meta.mode)?;
self.append_layer()?;
Ok(i)
}
}
// ── Mode 2 — count matrix ─────────────────────────────────────────────────────
impl LayeredMap<PersistentCompactIntMatrix> {
pub fn push_layer(&mut self, count_of: impl Fn(CanonicalKmer) -> u32) -> OLMResult<usize> {
let i = self.layers.len();
let dir = layer_dir(&self.root, i);
TypedLayer::<PersistentCompactIntMatrix>::build(&dir, DEFAULT_BLOCK_BITS, &self.meta.mode, count_of)?;
self.append_layer()?;
Ok(i)
}
pub fn push_layer_from_map(&mut self, counts: &HashMap<CanonicalKmer, u32>) -> OLMResult<usize> {
self.push_layer(|kmer| counts.get(&kmer).copied().unwrap_or(0))
}
}
#[cfg(test)]
#[path = "tests/map.rs"]
mod tests;
+63
View File
@@ -0,0 +1,63 @@
use std::fs::File;
use std::path::Path;
use serde::{Deserialize, Serialize};
use crate::layer::error::OLMResult;
const META_FILE: &str = "meta.json";
// ── IndexMode ─────────────────────────────────────────────────────────────────
/// Evidence mode for an entire partitioned index — homogeneous across all layers.
///
/// Determined once at build time; stored in `PartitionMeta` (`meta.json`).
/// All layers within an index share the same mode.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum IndexMode {
/// Exact evidence: `evidence.bin` + `unitigs.bin.idx`. Zero false positives.
Exact,
/// Approximate evidence: `fingerprint.bin` only.
/// `b` — fingerprint bits per slot; false-positive rate ≈ 1/2^b per query.
/// `z` — Findere consecutive-kmer parameter (build-time only; not used at query time).
Approx { b: u8, z: u8 },
/// Hybrid: both `fingerprint.bin` and `evidence.bin` + `unitigs.bin.idx`.
/// `find()` uses the fingerprint (O(1), approx); `find_strict()` uses exact evidence.
Hybrid { b: u8, z: u8 },
}
impl Default for IndexMode {
fn default() -> Self { Self::Exact }
}
// ── PartitionMeta ─────────────────────────────────────────────────────────────
/// Index-level metadata stored in `meta.json` at the root of a partition index.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PartitionMeta {
pub n_layers: usize,
#[serde(default)]
pub mode: IndexMode,
}
impl PartitionMeta {
pub fn new(mode: IndexMode) -> Self {
Self { n_layers: 0, mode }
}
pub fn load(dir: &Path) -> OLMResult<Self> {
let f = File::open(dir.join(META_FILE))?;
Ok(serde_json::from_reader(f)?)
}
pub fn save(&self, dir: &Path) -> OLMResult<()> {
let f = File::create(dir.join(META_FILE))?;
serde_json::to_writer_pretty(f, self)?;
Ok(())
}
}
impl Default for PartitionMeta {
fn default() -> Self { Self::new(IndexMode::Exact) }
}
+20
View File
@@ -0,0 +1,20 @@
pub mod content_layer;
pub mod error;
pub mod evidence;
pub mod fingerprint;
pub mod typed_layer;
pub mod layered_store;
pub mod map;
pub mod meta;
pub(crate) mod mphf_layer;
pub use content_layer::Layer;
pub use error::{OLMError, OLMResult};
pub use typed_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};
pub use mphf_layer::{EvidenceKind, KmerBatchIter, KmerIter, MphfLayer, MphfOnly};
+539
View File
@@ -0,0 +1,539 @@
use std::fs;
use std::iter::Enumerate;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use cacheline_ef::{CachelineEf, CachelineEfVec};
use epserde::prelude::*;
use obicompactvec::LayerMeta;
use obikseq::CanonicalKmer;
use obiskio::{CanonicalKmerIter, UnitigFileReader, UnitigFileWriter, build_unitig_idx};
use ptr_hash::{PtrHash, PtrHashParams, bucket_fn::CubicEps, hash::Xx64};
use crate::layer::error::{OLMError, OLMResult};
use crate::layer::evidence::{Evidence, EvidenceWriter};
use crate::layer::fingerprint::{FingerprintVec, FingerprintVecWriter};
use crate::layer::meta::IndexMode;
pub(crate) const MPHF_FILE: &str = "mphf.bin";
pub(crate) const UNITIGS_FILE: &str = "unitigs.bin";
pub(crate) const EVIDENCE_FILE: &str = "evidence.bin";
pub(crate) const FINGERPRINT_FILE: &str = "fingerprint.bin";
/// Owned MPHF — used only at build time (construction + store).
pub(crate) type Mphf = PtrHash<u64, CubicEps, CachelineEfVec<Vec<CachelineEf>>, Xx64, Vec<u8>>;
/// Zero-copy MPHF for querying — ε-deserialized view into a memory-mapped file.
/// `MemCase` owns the mmap backing; `'static` is sound because MemCase pins the memory.
type MphfEps = PtrHash<u64, CubicEps, CachelineEfVec<&'static [CachelineEf]>, Xx64, &'static [u8]>;
// ── LayerEvidence ─────────────────────────────────────────────────────────────
enum LayerEvidence {
Exact { evidence: Evidence, unitigs: Arc<UnitigFileReader> },
Approx { fingerprint: FingerprintVec, unitigs: Arc<UnitigFileReader>, unitigs_path: PathBuf },
Hybrid { evidence: Evidence, unitigs: Arc<UnitigFileReader>, fingerprint: FingerprintVec },
}
// ── EvidenceKind ──────────────────────────────────────────────────────────────
/// Coarse evidence mode of a layer — the `IndexMode` a layer was opened
/// with, without the `Approx`/`Hybrid` `b`/`z` parameters.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum EvidenceKind {
Exact,
Approx,
Hybrid,
}
impl EvidenceKind {
/// Lightweight disk probe, no MPHF/evidence opened: presence of
/// `evidence.bin`/`fingerprint.bin` alone determines the mode — same
/// signal `MphfLayer::open` uses, just without opening either file.
pub fn detect(layer_dir: &Path) -> OLMResult<EvidenceKind> {
let has_evidence = layer_dir.join(EVIDENCE_FILE).exists();
let has_fingerprint = layer_dir.join(FINGERPRINT_FILE).exists();
match (has_evidence, has_fingerprint) {
(true, false) => Ok(EvidenceKind::Exact),
(false, true) => Ok(EvidenceKind::Approx),
(true, true) => Ok(EvidenceKind::Hybrid),
(false, false) => Err(OLMError::InvalidLayer(format!(
"no evidence.bin or fingerprint.bin in {}", layer_dir.display()
))),
}
}
}
// ── MphfLayer ─────────────────────────────────────────────────────────────────
/// Autonomous kmer → slot mapping for one layer.
///
/// Two query methods:
/// - [`find`](Self::find) — O(1), uses fingerprint (Approx/Hybrid) or exact evidence (Exact).
/// - [`find_strict`](Self::find_strict) — always exact; O(1) on Exact/Hybrid layers,
/// O(n) sequential scan on Approx layers.
pub struct MphfLayer {
mphf: MemCase<MphfEps>,
ev: LayerEvidence,
n: usize,
}
impl MphfLayer {
/// Open a layer using the index-level `mode` determined at `LayeredMap` open time.
/// No per-layer metadata file is read.
pub fn open(dir: &Path, mode: &IndexMode) -> OLMResult<Self> {
let mphf: MemCase<MphfEps> = Mphf::mmap(&dir.join(MPHF_FILE), Flags::empty())
.map_err(|e| OLMError::InvalidLayer(e.to_string()))?;
let (ev, n) = match mode {
IndexMode::Exact => {
let evidence = Evidence::open(&dir.join(EVIDENCE_FILE))?;
let n = evidence.len();
let unitigs = Arc::new(UnitigFileReader::open(&dir.join(UNITIGS_FILE))?);
(LayerEvidence::Exact { evidence, unitigs }, n)
}
IndexMode::Approx { .. } => {
let fingerprint = FingerprintVec::open(&dir.join(FINGERPRINT_FILE))?;
let n = fingerprint.n();
let unitigs = Arc::new(UnitigFileReader::open(&dir.join(UNITIGS_FILE))?);
let unitigs_path = dir.join(UNITIGS_FILE);
(LayerEvidence::Approx { fingerprint, unitigs, unitigs_path }, n)
}
IndexMode::Hybrid { .. } => {
let evidence = Evidence::open(&dir.join(EVIDENCE_FILE))?;
let fingerprint = FingerprintVec::open(&dir.join(FINGERPRINT_FILE))?;
let n = evidence.len();
let unitigs = Arc::new(UnitigFileReader::open(&dir.join(UNITIGS_FILE))?);
(LayerEvidence::Hybrid { evidence, unitigs, fingerprint }, n)
}
};
Ok(Self { mphf, ev, n })
}
// ── Query API ─────────────────────────────────────────────────────────────
/// O(1) lookup — dispatches automatically:
/// - Exact: evidence + `verify_canonical_kmer`, zero false positives.
/// - Approx: fingerprint check, false-positive rate ≈ 1/2^b.
/// - Hybrid: fingerprint check (fast path), zero false positives via `find_strict`.
#[inline]
pub fn find(&self, kmer: CanonicalKmer) -> Option<usize> {
let slot = self.mphf.index(&kmer.raw());
if slot >= self.n { return None; }
match &self.ev {
LayerEvidence::Exact { evidence, unitigs, .. } => {
let (chunk_id, rank) = evidence.decode(slot);
if unitigs.verify_canonical_kmer(chunk_id as usize, rank as usize, kmer) {
Some(slot)
} else {
None
}
}
LayerEvidence::Approx { fingerprint, .. } |
LayerEvidence::Hybrid { fingerprint, .. } => {
if fingerprint.matches(slot, kmer.seq_hash()) { Some(slot) } else { None }
}
}
}
/// Always-exact lookup — zero false positives regardless of mode.
///
/// - Exact/Hybrid: O(1) via evidence + `verify_canonical_kmer`.
/// - Approx: O(n) sequential scan of `unitigs.bin` to confirm the kmer
/// that owns the slot, then exact comparison.
pub fn find_strict(&self, kmer: CanonicalKmer) -> Option<usize> {
let slot = self.mphf.index(&kmer.raw());
if slot >= self.n { return None; }
match &self.ev {
LayerEvidence::Exact { evidence, unitigs, .. } |
LayerEvidence::Hybrid { evidence, unitigs, .. } => {
let (chunk_id, rank) = evidence.decode(slot);
if unitigs.verify_canonical_kmer(chunk_id as usize, rank as usize, kmer) {
Some(slot)
} else {
None
}
}
LayerEvidence::Approx { unitigs_path, .. } => {
let reader = UnitigFileReader::open_sequential(unitigs_path).ok()?;
for (stored, _, _) in reader.iter_indexed_canonical_kmers() {
if self.mphf.index(&stored.raw()) == slot {
return if stored == kmer { Some(slot) } else { None };
}
}
None
}
}
}
/// Reconstruct the canonical k-mer stored at `slot` — the inverse of
/// [`find`](Self::find)/[`find_strict`](Self::find_strict) (k-mer → slot).
/// O(1) on `Exact`/`Hybrid` layers: `evidence.decode(slot)` gives
/// `(chunk_id, rank)` directly (no MPHF hashing, no file scan), then
/// `unitigs.canonical_raw_kmer` is a direct-access read. `None` on
/// `Approx` layers — fingerprints alone can't recover a k-mer, and
/// falling back to a full sequential scan here would silently make one
/// call cost O(n); callers needing this on an `Approx` layer should
/// scan `unitigs.bin` themselves and decide how to handle that cost
/// explicitly.
pub fn kmer_at(&self, slot: usize) -> Option<CanonicalKmer> {
if slot >= self.n {
return None;
}
match &self.ev {
LayerEvidence::Exact { evidence, unitigs, .. } |
LayerEvidence::Hybrid { evidence, unitigs, .. } => {
let (chunk_id, rank) = evidence.decode(slot);
let raw = unitigs.canonical_raw_kmer(chunk_id as usize, rank as usize);
Some(CanonicalKmer::from_raw_unchecked(raw))
}
LayerEvidence::Approx { .. } => None,
}
}
pub fn n(&self) -> usize { self.n }
/// This already-open layer's evidence mode — reads the discriminant
/// already in memory, no disk access.
pub fn evidence_kind(&self) -> EvidenceKind {
match &self.ev {
LayerEvidence::Exact { .. } => EvidenceKind::Exact,
LayerEvidence::Approx { .. } => EvidenceKind::Approx,
LayerEvidence::Hybrid { .. } => EvidenceKind::Hybrid,
}
}
/// Raw MPHF lookup: kmer → slot.
///
/// Returns the slot assigned by the MPHF without performing any membership
/// check against the layer's evidence. The returned slot is only meaningful
/// when the kmer is actually present in the layer; callers that need an
/// existence test should use [`find`](Self::find) instead.
pub fn index(&self, kmer: CanonicalKmer) -> usize {
self.mphf.index(&kmer.raw())
}
/// Batch raw MPHF lookup: kmers → slots.
///
/// Returns a [`Vec`] of slots, one per input kmer, in the same order as the
/// input slice. Like [`index`](Self::index), no membership check is
/// performed.
pub fn index_batch(&self, kmers: &[CanonicalKmer]) -> Vec<usize> {
kmers.iter().map(|k| self.mphf.index(&k.raw())).collect()
}
/// Iterate over all canonical kmers in the layer, in deterministic order.
///
/// The iteration order follows the physical layout of `unitigs.bin` and is
/// **not** correlated with MPHF slot numbers. Owns a clone of the
/// underlying `Arc<UnitigFileReader>` rather than borrowing `self` — `Send
/// + 'static`, so it can be handed to a threaded consumer (e.g.
/// `obipipeline`) directly, streamed from disk, with no need to collect
/// the layer's kmers into memory first. Multiple instances can be held
/// concurrently for as long as the underlying file exists, because each
/// carries its own cursor.
pub fn iter_kmers(&self) -> KmerIter {
let reader = match &self.ev {
LayerEvidence::Exact { unitigs, .. } => unitigs,
LayerEvidence::Approx { unitigs, .. } => unitigs,
LayerEvidence::Hybrid { unitigs, .. } => unitigs,
};
KmerIter { inner: Box::new(reader.iter_indexed_canonical_kmers_owned()) }
}
/// Iterate over all canonical kmers, each paired with its zero-based
/// sequence index in `unitigs.bin`.
///
/// Yields `(index, kmer)` where `index` starts at `0` for the first kmer
/// stored in the layer. This is the standard Rust `enumerate` adapter
/// applied to [`iter_kmers`](Self::iter_kmers).
pub fn enumerate_kmers(&self) -> Enumerate<KmerIter> {
self.iter_kmers().enumerate()
}
/// Iterate over the layer's canonical kmers in batches of `n`.
///
/// Each call to [`next`](Iterator::next) returns a [`Vec`] of up to `n`
/// kmers. The final batch may be shorter when the layer is exhausted.
pub fn iter_kmers_batch(&self, n: usize) -> KmerBatchIter {
KmerBatchIter { inner: self.iter_kmers(), batch_size: n }
}
/// Iterate over batches, each paired with the zero-based index of the
/// first kmer in the batch.
///
/// Yields `(batch_start_index, Vec<CanonicalKmer>)` where
/// `batch_start_index` is the iteration-order offset of the first kmer
/// in that batch within the full layer sequence — i.e. a multiple of `n`
/// except for the final (possibly shorter) batch.
pub fn enumerate_kmers_batch(&self, n: usize) -> impl Iterator<Item = (usize, Vec<CanonicalKmer>)> + Send + 'static {
let mut offset = 0usize;
self.iter_kmers_batch(n).map(move |batch| {
let base = offset;
offset += batch.len();
(base, batch)
})
}
}
// ── Iterator types ────────────────────────────────────────────────────────────
/// Iterator over the canonical kmers stored in a layer.
///
/// Produced by [`MphfLayer::iter_kmers`]. Owns an `Arc<UnitigFileReader>`
/// clone internally (via `iter_indexed_canonical_kmers_owned`) instead of
/// borrowing the parent layer — `Send + 'static`, streamed from disk one
/// kmer at a time, never materialised as a whole. Multiple `KmerIter`
/// instances can coexist concurrently, because each holds its own cursor.
pub struct KmerIter {
inner: Box<dyn Iterator<Item = (CanonicalKmer, usize, usize)> + Send>,
}
impl Iterator for KmerIter {
type Item = CanonicalKmer;
/// Return the next canonical kmer in iteration order.
fn next(&mut self) -> Option<Self::Item> {
self.inner.next().map(|(kmer, _, _)| kmer)
}
}
/// Iterator over batches of canonical kmers stored in a layer.
///
/// Produced by [`MphfLayer::iter_kmers_batch`]. Each call to [`next`](Self::next)
/// returns a [`Vec`] of up to `batch_size` kmers. The last batch may be shorter
/// than `batch_size` when the layer is exhausted.
pub struct KmerBatchIter {
inner: KmerIter,
batch_size: usize,
}
impl Iterator for KmerBatchIter {
type Item = Vec<CanonicalKmer>;
/// Return the next batch of kmers, or `None` if the layer is exhausted.
fn next(&mut self) -> Option<Self::Item> {
let mut batch = Vec::with_capacity(self.batch_size);
for _ in 0..self.batch_size {
if let Some(kmer) = self.inner.next() {
batch.push(kmer);
} else {
break;
}
}
if batch.is_empty() {
None
} else {
Some(batch)
}
}
}
// ── MphfOnly ──────────────────────────────────────────────────────────────────
/// Lightweight wrapper that loads only the MPHF file, without evidence or unitigs.
///
/// Use this when the caller guarantees that all queried kmers are in the MPHF
/// domain (e.g. when iterating the source's own unitigs during merge).
pub struct MphfOnly(MemCase<MphfEps>);
impl MphfOnly {
pub fn open(dir: &Path) -> OLMResult<Self> {
let mphf: MemCase<MphfEps> = Mphf::mmap(&dir.join(MPHF_FILE), Flags::empty())
.map_err(|e| OLMError::InvalidLayer(e.to_string()))?;
Ok(Self(mphf))
}
/// Return the slot for `kmer`. Only valid when `kmer` is in the MPHF domain.
#[inline]
pub fn index(&self, kmer: CanonicalKmer) -> usize {
self.0.index(&kmer.raw())
}
}
impl MphfLayer {
// ── Build helpers ─────────────────────────────────────────────────────────
pub fn unitig_writer(dir: &Path) -> OLMResult<UnitigFileWriter> {
fs::create_dir_all(dir)?;
Ok(UnitigFileWriter::create(&dir.join(UNITIGS_FILE))?)
}
/// Build `evidence.bin` + `unitigs.bin.idx` from `unitigs.bin` + `mphf.bin`.
pub fn build_exact_evidence(dir: &Path, block_bits: u8) -> OLMResult<usize> {
let unitig_path = dir.join(UNITIGS_FILE);
let unitigs = UnitigFileReader::open_sequential(&unitig_path)?;
let n = unitigs.n_kmers();
if n == 0 {
fs::File::create(dir.join(EVIDENCE_FILE))?;
build_unitig_idx(&unitig_path, block_bits)?;
return Ok(0);
}
let mphf: Mphf = Mphf::load_full(&dir.join(MPHF_FILE))
.map_err(|e| OLMError::InvalidLayer(e.to_string()))?;
let mut ev = EvidenceWriter::new(n);
let mut seen = vec![0u8; (n + 7) / 8];
for (kmer, chunk_id, rank) in unitigs.iter_indexed_canonical_kmers() {
let slot = mphf.index(&kmer.raw());
if slot >= n {
return Err(OLMError::Mphf("slot out of bounds".into()));
}
let byte = slot / 8;
let bit = 1u8 << (slot % 8);
if seen[byte] & bit != 0 {
return Err(OLMError::Mphf("duplicate slot".into()));
}
seen[byte] |= bit;
ev.set(slot, chunk_id as u32, rank as u8);
}
ev.write(&dir.join(EVIDENCE_FILE))?;
build_unitig_idx(&unitig_path, block_bits)?;
Ok(n)
}
/// Build `fingerprint.bin` from `unitigs.bin` + `mphf.bin`.
pub fn build_approx_evidence(dir: &Path, b: u8, z: u8) -> OLMResult<usize> {
if b == 0 || b > 64 {
return Err(OLMError::InvalidLayer("fingerprint width must be 1..=64".into()));
}
if z == 0 {
return Err(OLMError::InvalidLayer("z must be ≥ 1".into()));
}
let unitig_path = dir.join(UNITIGS_FILE);
let unitigs = UnitigFileReader::open_sequential(&unitig_path)?;
let n = unitigs.n_kmers();
if n == 0 {
FingerprintVecWriter::new(0, b).write(&dir.join(FINGERPRINT_FILE))?;
return Ok(0);
}
let mphf: Mphf = Mphf::load_full(&dir.join(MPHF_FILE))
.map_err(|e| OLMError::InvalidLayer(e.to_string()))?;
let mut fw = FingerprintVecWriter::new(n, b);
for (kmer, _, _) in unitigs.iter_indexed_canonical_kmers() {
let slot = mphf.index(&kmer.raw());
if slot >= n {
return Err(OLMError::Mphf("slot out of bounds".into()));
}
fw.set(slot, kmer.seq_hash());
}
fw.write(&dir.join(FINGERPRINT_FILE))?;
Ok(n)
}
/// Build MPHF + evidence from `unitigs.bin` already present in `dir`.
///
/// `fill_slot(slot, kmer)` is called once per kmer in all modes.
/// No `layer_meta.json` is written — the mode is an index-level property
/// stored in `PartitionMeta`.
pub(crate) fn build(
dir: &Path,
block_bits: u8,
mode: &IndexMode,
fill_slot: &mut impl FnMut(usize, CanonicalKmer) -> OLMResult<()>,
) -> OLMResult<usize> {
use rayon::prelude::*;
let unitig_path = dir.join(UNITIGS_FILE);
let n = UnitigFileReader::open_sequential(&unitig_path)?.n_kmers();
let sk_to_olm = |e: obiskio::SKError| match e {
obiskio::SKError::Io(io) => OLMError::Io(io),
e => OLMError::InvalidLayer(e.to_string()),
};
// ── Empty layer ───────────────────────────────────────────────────────
if n == 0 {
let mphf: Mphf =
Mphf::try_new(&[] as &[u64], PtrHashParams::<CubicEps>::default())
.ok_or_else(|| OLMError::Mphf("construction failed".into()))?;
mphf.store(&dir.join(MPHF_FILE))
.map_err(|e| OLMError::InvalidLayer(e.to_string()))?;
match mode {
IndexMode::Exact | IndexMode::Hybrid { .. } => {
fs::File::create(dir.join(EVIDENCE_FILE))?;
build_unitig_idx(&unitig_path, block_bits)?;
}
IndexMode::Approx { b, .. } => {
FingerprintVecWriter::new(0, *b).write(&dir.join(FINGERPRINT_FILE))?;
}
}
if let IndexMode::Hybrid { b, .. } = mode {
FingerprintVecWriter::new(0, *b).write(&dir.join(FINGERPRINT_FILE))?;
}
return Ok(0);
}
// ── Pass 1: MPHF via clonable mmap iterator ───────────────────────────
let keys = CanonicalKmerIter::new(&unitig_path).map_err(sk_to_olm)?;
let mphf: Mphf =
Mphf::new_from_par_iter(n, keys.map(|k| k.raw()).par_bridge(),
PtrHashParams::<CubicEps>::default());
mphf.store(&dir.join(MPHF_FILE))
.map_err(|e| OLMError::InvalidLayer(e.to_string()))?;
// ── Pass 2: fill evidence files + callback ────────────────────────────
let unitigs2 = UnitigFileReader::open_sequential(&unitig_path)?;
let mut seen = vec![0u8; (n + 7) / 8];
match mode {
IndexMode::Exact => {
let mut ev = EvidenceWriter::new(n);
for (kmer, chunk_id, rank) in unitigs2.iter_indexed_canonical_kmers() {
let slot = mphf.index(&kmer.raw());
if slot >= n { return Err(OLMError::Mphf("slot out of bounds".into())); }
let byte = slot / 8; let bit = 1u8 << (slot % 8);
if seen[byte] & bit != 0 { return Err(OLMError::Mphf("duplicate slot".into())); }
seen[byte] |= bit;
ev.set(slot, chunk_id as u32, rank as u8);
fill_slot(slot, kmer)?;
}
ev.write(&dir.join(EVIDENCE_FILE))?;
build_unitig_idx(&unitig_path, block_bits)?;
}
IndexMode::Approx { b, .. } => {
let mut fw = FingerprintVecWriter::new(n, *b);
for (kmer, _, _) in unitigs2.iter_indexed_canonical_kmers() {
let slot = mphf.index(&kmer.raw());
if slot >= n { return Err(OLMError::Mphf("slot out of bounds".into())); }
let byte = slot / 8; let bit = 1u8 << (slot % 8);
if seen[byte] & bit != 0 { return Err(OLMError::Mphf("duplicate slot".into())); }
seen[byte] |= bit;
fw.set(slot, kmer.seq_hash());
fill_slot(slot, kmer)?;
}
fw.write(&dir.join(FINGERPRINT_FILE))?;
}
IndexMode::Hybrid { b, .. } => {
let mut ev = EvidenceWriter::new(n);
let mut fw = FingerprintVecWriter::new(n, *b);
for (kmer, chunk_id, rank) in unitigs2.iter_indexed_canonical_kmers() {
let slot = mphf.index(&kmer.raw());
if slot >= n { return Err(OLMError::Mphf("slot out of bounds".into())); }
let byte = slot / 8; let bit = 1u8 << (slot % 8);
if seen[byte] & bit != 0 { return Err(OLMError::Mphf("duplicate slot".into())); }
seen[byte] |= bit;
ev.set(slot, chunk_id as u32, rank as u8);
fw.set(slot, kmer.seq_hash());
fill_slot(slot, kmer)?;
}
ev.write(&dir.join(EVIDENCE_FILE))?;
fw.write(&dir.join(FINGERPRINT_FILE))?;
build_unitig_idx(&unitig_path, block_bits)?;
}
}
LayerMeta::save(dir, n)?;
Ok(n)
}
}
@@ -0,0 +1,381 @@
use super::*;
use obicompactvec::{
PersistentBitMatrix, PersistentBitMatrixBuilder,
PersistentCompactIntMatrix, PersistentCompactIntMatrixBuilder,
};
use tempfile::tempdir;
fn make_int_matrix(cols: &[&[u32]]) -> (tempfile::TempDir, PersistentCompactIntMatrix) {
let n = cols.first().map_or(0, |c| c.len());
let dir = tempdir().unwrap();
let mut b = PersistentCompactIntMatrixBuilder::new(n, &dir.path().join("counts")).unwrap();
for &col in cols {
let mut cb = b.add_col().unwrap();
for (slot, &v) in col.iter().enumerate() { cb.set(slot, v); }
cb.close().unwrap();
}
b.close().unwrap();
let m = PersistentCompactIntMatrix::open(dir.path()).unwrap();
(dir, m)
}
fn make_bit_matrix(cols: &[&[bool]]) -> (tempfile::TempDir, PersistentBitMatrix) {
let n = cols.first().map_or(0, |c| c.len());
let dir = tempdir().unwrap();
let mut b = PersistentBitMatrixBuilder::new(n, &dir.path().join("presence")).unwrap();
for &col in cols {
let mut cb = b.add_col().unwrap();
for (slot, &v) in col.iter().enumerate() { cb.set(slot, v); }
cb.close().unwrap();
}
b.close().unwrap();
let m = PersistentBitMatrix::open(dir.path()).unwrap();
(dir, m)
}
// ── ColumnWeights ─────────────────────────────────────────────────────────
#[test]
fn col_weights_sums_across_layers() {
// layer 0: col0=[1,2], col1=[3,4] → weights [3, 7]
// layer 1: col0=[10,0], col1=[0,10] → weights [10, 10]
// combined: [13, 17]
let (_d0, m0) = make_int_matrix(&[&[1, 2], &[3, 4]]);
let (_d1, m1) = make_int_matrix(&[&[10, 0], &[0, 10]]);
let store = LayeredStore::new(vec![m0, m1]);
let w = store.col_weights();
assert_eq!(w[0], 13);
assert_eq!(w[1], 17);
}
#[test]
fn col_weights_bit_sums_across_layers() {
// layer 0: col0=[T,F,T], col1=[F,T,T] → counts [2, 2]
// layer 1: col0=[F,F,T], col1=[T,T,F] → counts [1, 2]
// combined: [3, 4]
let (_d0, m0) = make_bit_matrix(&[&[true, false, true], &[false, true, true]]);
let (_d1, m1) = make_bit_matrix(&[&[false, false, true], &[true, true, false]]);
let store = LayeredStore::new(vec![m0, m1]);
let w = store.col_weights();
assert_eq!(w[0], 3);
assert_eq!(w[1], 4);
}
// ── CountPartials — layered (one partition) ───────────────────────────────
#[test]
fn layered_bray_matches_combined() {
// Split [1,2,3,4,5] across two layers; bray dist should equal direct computation
// on [1,2,3,4,5] for each column pair.
// col0=[1,2,3,4,5], col1=[5,4,3,2,1]
let (_d0, m0) = make_int_matrix(&[&[1, 2], &[5, 4]]); // slots 0-1
let (_d1, m1) = make_int_matrix(&[&[3, 4, 5], &[3, 2, 1]]); // slots 2-4
let store = LayeredStore::new(vec![m0, m1]);
// direct on full data
let (_df, mf) = make_int_matrix(&[&[1, 2, 3, 4, 5], &[5, 4, 3, 2, 1]]);
let expected = CountPartials::bray_dist_matrix(&mf);
let got = CountPartials::bray_dist_matrix(&store);
assert!((got[[0, 1]] - expected[[0, 1]]).abs() < 1e-12, "bray [0,1]");
assert!((got[[1, 0]] - expected[[1, 0]]).abs() < 1e-12, "bray [1,0]");
}
#[test]
fn layered_relfreq_bray_matches_combined() {
let (_d0, m0) = make_int_matrix(&[&[1, 2], &[5, 4]]);
let (_d1, m1) = make_int_matrix(&[&[3, 4, 5], &[3, 2, 1]]);
let store = LayeredStore::new(vec![m0, m1]);
let (_df, mf) = make_int_matrix(&[&[1, 2, 3, 4, 5], &[5, 4, 3, 2, 1]]);
let expected = CountPartials::relfreq_bray_dist_matrix(&mf);
let got = CountPartials::relfreq_bray_dist_matrix(&store);
assert!((got[[0, 1]] - expected[[0, 1]]).abs() < 1e-12, "relfreq_bray [0,1]");
}
#[test]
fn layered_euclidean_matches_combined() {
let (_d0, m0) = make_int_matrix(&[&[3, 0], &[0, 4]]);
let (_d1, m1) = make_int_matrix(&[&[1, 1], &[2, 2]]);
let store = LayeredStore::new(vec![m0, m1]);
let (_df, mf) = make_int_matrix(&[&[3, 0, 1, 1], &[0, 4, 2, 2]]);
let expected = CountPartials::euclidean_dist_matrix(&mf);
let got = CountPartials::euclidean_dist_matrix(&store);
assert!((got[[0, 1]] - expected[[0, 1]]).abs() < 1e-12, "euclidean [0,1]");
}
// ── CountPartials — partitioned (LayeredStore<LayeredStore<_>>) ───────────
#[test]
fn partitioned_bray_matches_combined() {
// partition 0: slots [1,2,3,4,5] col0 vs col1
// partition 1: slots [10,20] col0 vs col1
let (_d0, p0) = make_int_matrix(&[&[1, 2, 3, 4, 5], &[5, 4, 3, 2, 1]]);
let (_d1, p1) = make_int_matrix(&[&[10, 20], &[20, 10]]);
let partitioned = LayeredStore::new(vec![
LayeredStore::new(vec![p0]),
LayeredStore::new(vec![p1]),
]);
let (_df, mf) = make_int_matrix(&[&[1, 2, 3, 4, 5, 10, 20], &[5, 4, 3, 2, 1, 20, 10]]);
let expected = CountPartials::bray_dist_matrix(&mf);
let got = CountPartials::bray_dist_matrix(&partitioned);
assert!((got[[0, 1]] - expected[[0, 1]]).abs() < 1e-12, "partitioned bray [0,1]");
}
#[test]
fn partitioned_threshold_jaccard_off_diagonal_is_pairwise() {
// 3 genomes, 2 partitions, 1 layer each — mirrors distance.rs's
// LayeredStore<LayeredStore<PersistentCompactIntMatrix>> shape.
// partition 0: col0=[3,0], col1=[0,3], col2=[3,3]
// partition 1: col0=[1,1], col1=[1,0], col2=[0,1]
let (_d0, p0) = make_int_matrix(&[&[3, 0], &[0, 3], &[3, 3]]);
let (_d1, p1) = make_int_matrix(&[&[1, 1], &[1, 0], &[0, 1]]);
let partitioned = LayeredStore::new(vec![
LayeredStore::new(vec![p0]),
LayeredStore::new(vec![p1]),
]);
let (_df, mf) = make_int_matrix(&[&[3, 0, 1, 1], &[0, 3, 1, 0], &[3, 3, 0, 1]]);
let threshold = 1u32;
let (inter_p, union_p) = CountPartials::partial_threshold_jaccard(&partitioned, threshold);
let (inter_f, union_f) = CountPartials::partial_threshold_jaccard(&mf, threshold);
let n = 3;
for i in 0..n {
for j in 0..n {
assert_eq!(inter_p[[i, j]], inter_f[[i, j]], "inter[{i},{j}]");
assert_eq!(union_p[[i, j]], union_f[[i, j]], "union[{i},{j}]");
}
}
}
#[test]
fn partitioned_threshold_jaccard_packed_off_diagonal_is_pairwise() {
// Same as `partitioned_threshold_jaccard_off_diagonal_is_pairwise` but
// each partition matrix is packed into a single .pcmx file first —
// the on-disk format actually used in production after `pack_matrices`.
use obicompactvec::pack_compact_int_matrix;
let (d0, _p0) = make_int_matrix(&[&[3, 0], &[0, 3], &[3, 3]]);
pack_compact_int_matrix(&d0.path().join("counts")).unwrap();
let p0 = PersistentCompactIntMatrix::open(d0.path()).unwrap();
let (d1, _p1) = make_int_matrix(&[&[1, 1], &[1, 0], &[0, 1]]);
pack_compact_int_matrix(&d1.path().join("counts")).unwrap();
let p1 = PersistentCompactIntMatrix::open(d1.path()).unwrap();
let partitioned = LayeredStore::new(vec![
LayeredStore::new(vec![p0]),
LayeredStore::new(vec![p1]),
]);
let (_df, mf) = make_int_matrix(&[&[3, 0, 1, 1], &[0, 3, 1, 0], &[3, 3, 0, 1]]);
let threshold = 1u32;
let (inter_p, union_p) = CountPartials::partial_threshold_jaccard(&partitioned, threshold);
let (inter_f, union_f) = CountPartials::partial_threshold_jaccard(&mf, threshold);
let n = 3;
for i in 0..n {
for j in 0..n {
assert_eq!(inter_p[[i, j]], inter_f[[i, j]], "inter[{i},{j}]");
assert_eq!(union_p[[i, j]], union_f[[i, j]], "union[{i},{j}]");
}
}
}
#[test]
fn partitioned_multilayer_threshold_jaccard_off_diagonal_is_pairwise() {
// 2 partitions, 2 layers each — the shape production indexes actually
// have (MPHF collision layers within a partition).
// partition 0, layer 0: col0=[3,0], col1=[0,3], col2=[3,3]
// partition 0, layer 1: col0=[2,0], col1=[0,0], col2=[2,0]
// partition 1, layer 0: col0=[1,1], col1=[1,0], col2=[0,1]
// partition 1, layer 1: col0=[0,5], col1=[5,5], col2=[0,0]
let (_d0a, p0a) = make_int_matrix(&[&[3, 0], &[0, 3], &[3, 3]]);
let (_d0b, p0b) = make_int_matrix(&[&[2, 0], &[0, 0], &[2, 0]]);
let (_d1a, p1a) = make_int_matrix(&[&[1, 1], &[1, 0], &[0, 1]]);
let (_d1b, p1b) = make_int_matrix(&[&[0, 5], &[5, 5], &[0, 0]]);
let partitioned = LayeredStore::new(vec![
LayeredStore::new(vec![p0a, p0b]),
LayeredStore::new(vec![p1a, p1b]),
]);
// Flattened equivalent: concatenate every layer's slots into one matrix.
let (_df, mf) = make_int_matrix(&[
&[3, 0, 2, 0, 1, 1, 0, 5],
&[0, 3, 0, 0, 1, 0, 5, 5],
&[3, 3, 2, 0, 0, 1, 0, 0],
]);
let threshold = 1u32;
let (inter_p, union_p) = CountPartials::partial_threshold_jaccard(&partitioned, threshold);
let (inter_f, union_f) = CountPartials::partial_threshold_jaccard(&mf, threshold);
let n = 3;
for i in 0..n {
for j in 0..n {
assert_eq!(inter_p[[i, j]], inter_f[[i, j]], "inter[{i},{j}]");
assert_eq!(union_p[[i, j]], union_f[[i, j]], "union[{i},{j}]");
}
}
}
// ── BitPartials ───────────────────────────────────────────────────────────
#[test]
fn layered_jaccard_matches_combined() {
let (_d0, m0) = make_bit_matrix(&[&[true, false], &[false, true]]);
let (_d1, m1) = make_bit_matrix(&[&[true, true], &[true, false]]);
let store = LayeredStore::new(vec![m0, m1]);
let (_df, mf) = make_bit_matrix(&[
&[true, false, true, true],
&[false, true, true, false],
]);
let expected = BitPartials::jaccard_dist_matrix(&mf);
let got = BitPartials::jaccard_dist_matrix(&store);
assert!((got[[0, 1]] - expected[[0, 1]]).abs() < 1e-12, "jaccard [0,1]");
}
#[test]
fn layered_hamming_matches_combined() {
let (_d0, m0) = make_bit_matrix(&[&[true, false], &[false, true]]);
let (_d1, m1) = make_bit_matrix(&[&[true, true], &[false, false]]);
let store = LayeredStore::new(vec![m0, m1]);
let (_df, mf) = make_bit_matrix(&[
&[true, false, true, true],
&[false, true, false, false],
]);
let expected = BitPartials::hamming_dist_matrix(&mf);
let got = BitPartials::hamming_dist_matrix(&store);
assert_eq!(got[[0, 1]], expected[[0, 1]], "hamming [0,1]");
}
#[test]
fn partitioned_bit_jaccard_off_diagonal_is_pairwise() {
// Same shape as the count-based `partitioned_multilayer_threshold_jaccard_*`
// tests, but for the presence/bit path (`with_counts = false` — what
// `all_specifics` actually uses in production).
// 4 genomes, 3 partitions, 2 layers in the last one.
let (_d0, p0) = make_bit_matrix(&[
&[true, false, true],
&[false, true, true],
&[true, true, false],
&[false, false, true],
]);
let (_d1, p1) = make_bit_matrix(&[
&[true, true],
&[false, true],
&[true, false],
&[true, true],
]);
let (_d2a, p2a) = make_bit_matrix(&[
&[false, true],
&[true, true],
&[false, false],
&[true, false],
]);
let (_d2b, p2b) = make_bit_matrix(&[
&[true],
&[false],
&[true],
&[true],
]);
let partitioned = LayeredStore::new(vec![
LayeredStore::new(vec![p0]),
LayeredStore::new(vec![p1]),
LayeredStore::new(vec![p2a, p2b]),
]);
// Flattened equivalent: concatenate every partition/layer's slots.
let (_df, mf) = make_bit_matrix(&[
&[true, false, true, true, true, false, true, true],
&[false, true, true, false, true, true, true, false],
&[true, true, false, true, false, false, false, true],
&[false, false, true, true, true, true, false, true],
]);
let (inter_p, union_p) = BitPartials::partial_jaccard(&partitioned);
let (inter_f, union_f) = BitPartials::partial_jaccard(&mf);
let n = 4;
for i in 0..n {
for j in 0..n {
assert_eq!(inter_p[[i, j]], inter_f[[i, j]], "inter[{i},{j}]");
assert_eq!(union_p[[i, j]], union_f[[i, j]], "union[{i},{j}]");
}
}
}
#[test]
fn partitioned_bit_jaccard_packed_off_diagonal_is_pairwise() {
// Same as `partitioned_bit_jaccard_off_diagonal_is_pairwise` but every
// partition's presence matrix is packed into a single .pbmx file —
// the on-disk format actually used in production after `pack_matrices`.
use obicompactvec::pack_bit_matrix;
let (d0, _p0) = make_bit_matrix(&[
&[true, false, true],
&[false, true, true],
&[true, true, false],
&[false, false, true],
]);
pack_bit_matrix(&d0.path().join("presence")).unwrap();
let p0 = PersistentBitMatrix::open(d0.path()).unwrap();
let (d1, _p1) = make_bit_matrix(&[
&[true, true],
&[false, true],
&[true, false],
&[true, true],
]);
pack_bit_matrix(&d1.path().join("presence")).unwrap();
let p1 = PersistentBitMatrix::open(d1.path()).unwrap();
let (d2a, _p2a) = make_bit_matrix(&[
&[false, true],
&[true, true],
&[false, false],
&[true, false],
]);
pack_bit_matrix(&d2a.path().join("presence")).unwrap();
let p2a = PersistentBitMatrix::open(d2a.path()).unwrap();
let (d2b, _p2b) = make_bit_matrix(&[
&[true],
&[false],
&[true],
&[true],
]);
pack_bit_matrix(&d2b.path().join("presence")).unwrap();
let p2b = PersistentBitMatrix::open(d2b.path()).unwrap();
let partitioned = LayeredStore::new(vec![
LayeredStore::new(vec![p0]),
LayeredStore::new(vec![p1]),
LayeredStore::new(vec![p2a, p2b]),
]);
let (_df, mf) = make_bit_matrix(&[
&[true, false, true, true, true, false, true, true],
&[false, true, true, false, true, true, true, false],
&[true, true, false, true, false, false, false, true],
&[false, false, true, true, true, true, false, true],
]);
let (inter_p, union_p) = BitPartials::partial_jaccard(&partitioned);
let (inter_f, union_f) = BitPartials::partial_jaccard(&mf);
let n = 4;
for i in 0..n {
for j in 0..n {
assert_eq!(inter_p[[i, j]], inter_f[[i, j]], "inter[{i},{j}]");
assert_eq!(union_p[[i, j]], union_f[[i, j]], "union[{i},{j}]");
}
}
}
+219
View File
@@ -0,0 +1,219 @@
use super::*;
use obicompactvec::{pack_bit_matrix, pack_sparse_bit_matrix, PersistentBitMatrix, PersistentCompactIntMatrix, StorageKind};
use obikseq::{set_k, Sequence as _, Unitig};
use obiskio::DEFAULT_BLOCK_BITS;
use crate::layer::typed_layer::LayerContent;
use crate::layer::meta::IndexMode;
use crate::layer::mphf_layer::EvidenceKind;
use tempfile::tempdir;
fn push_unitigs_and_layer(
map: &mut LayeredMap<PersistentCompactIntMatrix>,
seqs: &[&[u8]],
count: u32,
) {
let mut w = map.next_layer_writer().unwrap();
for s in seqs {
w.write(&Unitig::from_ascii(s)).unwrap();
}
w.close().unwrap();
map.push_layer(|_| count).unwrap();
}
fn canonical(ascii: &[u8]) -> CanonicalKmer {
obikseq::Kmer::from_ascii(ascii).unwrap().canonical()
}
#[test]
fn create_empty_map() {
set_k(4);
let dir = tempdir().unwrap();
let map = LayeredMap::<()>::create(dir.path(), IndexMode::Exact).unwrap();
assert_eq!(map.n_layers(), 0);
}
#[test]
fn open_reloads_layer_count() {
set_k(4);
let dir = tempdir().unwrap();
{
let mut map = LayeredMap::<PersistentCompactIntMatrix>::create(dir.path(), IndexMode::Exact).unwrap();
push_unitigs_and_layer(&mut map, &[b"AAAACGT"], 1);
}
let map = LayeredMap::<PersistentCompactIntMatrix>::open(dir.path()).unwrap();
assert_eq!(map.n_layers(), 1);
}
#[test]
fn query_finds_kmer_in_layer_zero() {
set_k(4);
let dir = tempdir().unwrap();
let mut map = LayeredMap::<PersistentCompactIntMatrix>::create(dir.path(), IndexMode::Exact).unwrap();
push_unitigs_and_layer(&mut map, &[b"AAAACGT"], 3);
let kmer = canonical(b"AAAC");
let (layer_idx, hit) = map.query(kmer).expect("kmer must be found");
assert_eq!(layer_idx, 0);
assert_eq!(hit.data[0], 3);
}
#[test]
fn query_finds_kmer_in_correct_layer() {
set_k(4);
let dir = tempdir().unwrap();
let mut map = LayeredMap::<PersistentCompactIntMatrix>::create(dir.path(), IndexMode::Exact).unwrap();
push_unitigs_and_layer(&mut map, &[b"AAAACGT"], 1);
push_unitigs_and_layer(&mut map, &[b"GGGACGT"], 2);
assert_eq!(map.n_layers(), 2);
let (li, hit) = map.query(canonical(b"AAAA")).expect("AAAA must be found");
assert_eq!(li, 0);
assert_eq!(hit.data[0], 1);
let (li, hit) = map.query(canonical(b"GGGA")).expect("GGGA must be found");
assert_eq!(li, 1);
assert_eq!(hit.data[0], 2);
}
#[test]
fn query_absent_returns_none() {
set_k(4);
let dir = tempdir().unwrap();
let mut map = LayeredMap::<PersistentCompactIntMatrix>::create(dir.path(), IndexMode::Exact).unwrap();
push_unitigs_and_layer(&mut map, &[b"AAAACGT"], 1);
let absent = canonical(b"CCCC");
assert!(map.query(absent).is_none());
}
#[test]
fn push_layer_from_map_convenience() {
set_k(4);
let dir = tempdir().unwrap();
let mut map = LayeredMap::<PersistentCompactIntMatrix>::create(dir.path(), IndexMode::Exact).unwrap();
let mut w = map.next_layer_writer().unwrap();
w.write(&Unitig::from_ascii(b"AAAACGT")).unwrap();
w.close().unwrap();
let counts: HashMap<CanonicalKmer, u32> = vec![
(canonical(b"AAAA"), 10u32),
].into_iter().collect();
map.push_layer_from_map(&counts).unwrap();
let (_, hit) = map.query(canonical(b"AAAA")).unwrap();
assert_eq!(hit.data[0], 10);
}
// ── detect_layer_content / detect_layer_storage / detect_layer_evidence ────
//
// All built directly on disk (bypassing `LayeredMap<D>::push_layer`, which
// only exists for modes 1 and 2 — mode 3/presence has no `push_layer`),
// then reopened as `LayeredMap<()>` — `()`'s `LayerData::open` never reads
// the matrix, so it's the right handle for exercising the disk-probing
// `detect_layer_*` methods regardless of what content/storage the layer
// underneath actually is.
fn write_unitigs_at(dir: &Path, seqs: &[&[u8]]) {
fs::create_dir_all(dir).unwrap();
let mut w = UnitigFileWriter::create(&dir.join(crate::layer::typed_layer::UNITIGS_FILE)).unwrap();
for s in seqs {
w.write(&Unitig::from_ascii(s)).unwrap();
}
w.close().unwrap();
}
/// Build a one-layer partition root with a presence layer at `layer_0`,
/// save `meta.json`, and reopen it as `LayeredMap<()>`.
fn presence_partition(seqs: &[&[u8]], n_genomes: usize) -> (tempfile::TempDir, LayeredMap<()>) {
let dir = tempdir().unwrap();
let mode = IndexMode::Exact;
let layer0 = layer_dir(dir.path(), 0);
write_unitigs_at(&layer0, seqs);
TypedLayer::<PersistentBitMatrix>::build_presence(&layer0, DEFAULT_BLOCK_BITS, &mode, n_genomes, |kmer, g| {
(kmer.raw().wrapping_add(g as u64)) % 2 == 0
}).unwrap();
PartitionMeta { n_layers: 1, mode }.save(dir.path()).unwrap();
let map = LayeredMap::<()>::open(dir.path()).unwrap();
(dir, map)
}
#[test]
fn detect_layer_content_count() {
set_k(4);
let dir = tempdir().unwrap();
let mut map = LayeredMap::<PersistentCompactIntMatrix>::create(dir.path(), IndexMode::Exact).unwrap();
push_unitigs_and_layer(&mut map, &[b"AAAACGT"], 3);
assert_eq!(map.detect_layer_content(0), LayerContent::Count);
}
#[test]
fn detect_layer_content_presence() {
set_k(4);
let (_dir, map) = presence_partition(&[b"AAATCTA", b"CTTCGCC", b"TGATACG"], 3);
assert_eq!(map.detect_layer_content(0), LayerContent::Presence);
}
#[test]
fn detect_layer_storage_columnar_then_packed_then_sparse() {
set_k(4);
let (dir, map) = presence_partition(&[b"AAATCTA", b"CTTCGCC", b"TGATACG"], 3);
let presence_dir = layer_dir(dir.path(), 0).join("presence");
assert_eq!(map.detect_layer_storage(0).unwrap(), StorageKind::Columnar);
pack_bit_matrix(&presence_dir).unwrap();
assert_eq!(map.detect_layer_storage(0).unwrap(), StorageKind::Packed);
pack_sparse_bit_matrix(&presence_dir).unwrap();
assert_eq!(map.detect_layer_storage(0).unwrap(), StorageKind::Sparse);
}
#[test]
fn detect_layer_storage_implicit() {
set_k(4);
let dir = tempdir().unwrap();
let mode = IndexMode::Exact;
let layer0 = layer_dir(dir.path(), 0);
write_unitigs_at(&layer0, &[b"AAAACGT"]);
// Mode-1 build: MPHF + evidence + `layer_meta.json`, no matrix at all —
// must read back as `Presence`/`Implicit`, never a third "empty" state.
TypedLayer::<()>::build(&layer0, DEFAULT_BLOCK_BITS, &mode).unwrap();
PartitionMeta { n_layers: 1, mode }.save(dir.path()).unwrap();
let map = LayeredMap::<()>::open(dir.path()).unwrap();
assert_eq!(map.detect_layer_content(0), LayerContent::Presence);
assert_eq!(map.detect_layer_storage(0).unwrap(), StorageKind::Implicit);
}
#[test]
fn detect_layer_storage_count_columnar_then_packed() {
set_k(4);
let dir = tempdir().unwrap();
let mut map = LayeredMap::<PersistentCompactIntMatrix>::create(dir.path(), IndexMode::Exact).unwrap();
push_unitigs_and_layer(&mut map, &[b"AAAACGT"], 3);
drop(map);
let map = LayeredMap::<()>::open(dir.path()).unwrap();
assert_eq!(map.detect_layer_storage(0).unwrap(), StorageKind::Columnar);
obicompactvec::pack_compact_int_matrix(&layer_dir(dir.path(), 0).join("counts")).unwrap();
assert_eq!(map.detect_layer_storage(0).unwrap(), StorageKind::Packed);
}
#[test]
fn detect_layer_evidence_exact_vs_approx() {
set_k(4);
let dir = tempdir().unwrap();
let layer0 = layer_dir(dir.path(), 0);
write_unitigs_at(&layer0, &[b"AAAACGT"]);
TypedLayer::<()>::build(&layer0, DEFAULT_BLOCK_BITS, &IndexMode::Exact).unwrap();
PartitionMeta { n_layers: 1, mode: IndexMode::Exact }.save(dir.path()).unwrap();
let map = LayeredMap::<()>::open(dir.path()).unwrap();
assert_eq!(map.detect_layer_evidence(0).unwrap(), EvidenceKind::Exact);
let approx = IndexMode::Approx { b: 8, z: 1 };
let dir2 = tempdir().unwrap();
let layer0b = layer_dir(dir2.path(), 0);
write_unitigs_at(&layer0b, &[b"AAAACGT"]);
TypedLayer::<()>::build(&layer0b, DEFAULT_BLOCK_BITS, &approx).unwrap();
PartitionMeta { n_layers: 1, mode: approx }.save(dir2.path()).unwrap();
let map2 = LayeredMap::<()>::open(dir2.path()).unwrap();
assert_eq!(map2.detect_layer_evidence(0).unwrap(), EvidenceKind::Approx);
}
@@ -0,0 +1,142 @@
use super::*;
use obicompactvec::PersistentSparseBitMatrixBuilder;
use obikseq::{set_k, Unitig};
use obiskio::DEFAULT_BLOCK_BITS;
use tempfile::tempdir;
fn write_unitigs(dir: &Path, seqs: &[&[u8]]) {
let mut w = UnitigFileWriter::create(&dir.join(UNITIGS_FILE)).unwrap();
for s in seqs {
w.write(&Unitig::from_ascii(s)).unwrap();
}
w.close().unwrap();
}
fn all_canonical_kmers(dir: &Path) -> Vec<CanonicalKmer> {
UnitigFileReader::open_sequential(&dir.join(UNITIGS_FILE)).unwrap()
.iter_indexed_canonical_kmers()
.map(|(kmer, _, _)| kmer)
.collect()
}
// ── Iterator-order consistency tests ─────────────────────────────────────────
// These exercise the unitig-file iterators directly, without constructing a
// layer/MPHF. The sibling-annex builder relies on the invariant that
// `enumerate_kmers_batch()` yields the same sequence as `enumerate_kmers()`,
// in the same order as `UnitigFileReader::iter_indexed_canonical_kmers()`.
// A mismatch between any of these would silently write garbage into the
// annex's order-indexed mask array, producing exactly the "minorant-only"
// sentinel (0x10) that triggers the stats.rs panic.
#[test]
fn canonical_kmer_iter_matches_reader() {
set_k(4);
let dir = tempdir().unwrap();
write_unitigs(dir.path(), &[b"AAAACGT", b"TTTTGCA"]);
let from_iter: Vec<CanonicalKmer> = obiskio::CanonicalKmerIter::new(&dir.path().join(UNITIGS_FILE))
.unwrap()
.collect();
let from_reader: Vec<CanonicalKmer> = all_canonical_kmers(dir.path());
assert_eq!(from_iter.len(), from_reader.len(), "different kmer counts");
assert_eq!(from_iter, from_reader, "CanonicalKmerIter and UnitigFileReader disagree");
}
// ── Generic `TypedLayer<D>` over dense vs. sparse presence storage ───────────────
#[test]
fn presence_layer_generic_over_sparse_matches_dense() {
// k=4, matching every other test in this crate's binary: `K`/`M` are
// process-wide (not thread-local) in test builds — see
// `obikseq::params` — so a test using a different k here would race
// against `map.rs`'s k=4 tests running concurrently in the same
// binary. These three sequences were chosen (randomised search) to
// have no shared canonical 4-mer among them, required for `ptr_hash`
// MPHF construction (duplicate keys make it panic).
set_k(4);
let dir = tempdir().unwrap();
write_unitigs(dir.path(), &[b"AAATCTA", b"CTTCGCC", b"TGATACG"]);
let n_genomes = 3;
let mode = IndexMode::Exact;
// Deterministic, arbitrary presence function — genome `g` carries a
// kmer iff `(kmer's raw bits + g)` is even. Doesn't need to be
// biologically meaningful, just the same on both sides of the
// dense/sparse comparison below.
TypedLayer::<PersistentBitMatrix>::build_presence(dir.path(), DEFAULT_BLOCK_BITS, &mode, n_genomes, |kmer, g| {
(kmer.raw().wrapping_add(g as u64)) % 2 == 0
}).unwrap();
let dense_layer = TypedLayer::<PersistentBitMatrix>::open(dir.path(), &mode).unwrap();
assert!(dense_layer.n_cols() >= 1);
// Build the sparse form directly into the same `presence/` dir the
// dense form already lives in (matches how `pack --sparse` will work:
// same directory, distinct filenames — see
// `DevDocMD/architecture/siblings.md`'s sparse-matrix section).
let dense_matrix = obicompactvec::PersistentBitMatrix::open(dir.path()).unwrap();
PersistentSparseBitMatrixBuilder::build_from_dense(&dense_matrix, &dir.path().join(PRESENCE_DIR))
.unwrap()
.close()
.unwrap();
let sparse_layer = TypedLayer::<PersistentSparseBitMatrix>::open(dir.path(), &mode).unwrap();
// Same generic methods, same results, different concrete `D`.
assert_eq!(dense_layer.n_cols(), sparse_layer.n_cols());
let n = dense_layer.n();
assert_eq!(n, sparse_layer.n());
let slots: Vec<usize> = (0..n).collect();
assert_eq!(dense_layer.sub_matrix(&slots), sparse_layer.sub_matrix(&slots));
// The matrix-agnostic, MPHF-only surface (from the generic
// `impl<D: LayerData> TypedLayer<D>`) must also agree: same kmer set,
// looked up through either concrete `D`.
for kmer in all_canonical_kmers(dir.path()) {
assert_eq!(dense_layer.find_slot(kmer), sparse_layer.find_slot(kmer));
}
}
// ── content / storage_kind / evidence_kind (runtime introspection) ─────────
#[test]
fn count_layer_reports_count_content_and_columnar_storage() {
set_k(4);
let dir = tempdir().unwrap();
write_unitigs(dir.path(), &[b"AAAACGT"]);
TypedLayer::<PersistentCompactIntMatrix>::build(dir.path(), DEFAULT_BLOCK_BITS, &IndexMode::Exact, |_| 1)
.unwrap();
let layer = TypedLayer::<PersistentCompactIntMatrix>::open(dir.path(), &IndexMode::Exact).unwrap();
assert_eq!(layer.content(), LayerContent::Count);
assert_eq!(layer.storage_kind(), obicompactvec::StorageKind::Columnar);
assert_eq!(layer.evidence_kind(), crate::layer::mphf_layer::EvidenceKind::Exact);
}
#[test]
fn presence_layer_reports_presence_content_and_columnar_storage() {
set_k(4);
let dir = tempdir().unwrap();
write_unitigs(dir.path(), &[b"AAATCTA", b"CTTCGCC", b"TGATACG"]);
TypedLayer::<PersistentBitMatrix>::build_presence(dir.path(), DEFAULT_BLOCK_BITS, &IndexMode::Exact, 2, |kmer, g| {
(kmer.raw().wrapping_add(g as u64)) % 2 == 0
}).unwrap();
let layer = TypedLayer::<PersistentBitMatrix>::open(dir.path(), &IndexMode::Exact).unwrap();
assert_eq!(layer.content(), LayerContent::Presence);
assert_eq!(layer.storage_kind(), obicompactvec::StorageKind::Columnar);
}
#[test]
fn enumerate_kmers_is_stable_across_calls() {
set_k(4);
let dir = tempdir().unwrap();
write_unitigs(dir.path(), &[b"AAAACGT", b"TTTTGCA"]);
let reader = UnitigFileReader::open_sequential(&dir.path().join(UNITIGS_FILE)).unwrap();
let first: Vec<CanonicalKmer> = reader.iter_indexed_canonical_kmers().map(|(k, _, _)| k).collect();
let reader2 = UnitigFileReader::open_sequential(&dir.path().join(UNITIGS_FILE)).unwrap();
let second: Vec<CanonicalKmer> = reader2.iter_indexed_canonical_kmers().map(|(k, _, _)| k).collect();
assert_eq!(first, second, "iter_indexed_canonical_kmers must be deterministic across calls");
}
+461
View File
@@ -0,0 +1,461 @@
use std::collections::HashMap;
use std::fs;
use std::path::{Path, PathBuf};
use obicompactvec::{
BinaryMatrix,
PersistentBitMatrix, PersistentBitMatrixBuilder,
PersistentCompactIntMatrix, PersistentCompactIntMatrixBuilder,
PersistentSparseBitMatrix,
};
use obikseq::CanonicalKmer;
use obiskio::{UnitigFileReader, UnitigFileWriter};
use crate::layer::error::{OLMError, OLMResult};
use crate::layer::meta::IndexMode;
use crate::layer::mphf_layer::MphfLayer;
pub(crate) use crate::layer::mphf_layer::UNITIGS_FILE;
pub(crate) const COUNTS_DIR: &str = "counts";
pub(crate) const PRESENCE_DIR: &str = "presence";
// ── Trait ─────────────────────────────────────────────────────────────────────
pub trait LayerData: Sized {
type Item;
fn open(layer_dir: &Path) -> OLMResult<Self>;
fn read(&self, slot: usize) -> Self::Item;
}
/// TypedLayer directory path for layer `i` within a partition's index root — the
/// single source of truth for the on-disk `layer_N` naming convention
/// (mirrors what `LayeredMap::open`/`push_layer` already use internally).
///
/// `obikindex::layer` operates within a single partition's index root; it has
/// no notion of "partition" at all. Turning a partition number into that
/// 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:
/// `crate::algorithms::partitionner::PartitionRouter`), read by whichever
/// algorithm dereplicates it (today:
/// `crate::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`.
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
/// `crate::algorithms::dereplicator::Dereplicator`, read by whichever
/// algorithm counts kmer abundances from it (today:
/// `crate::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}"))
}
/// 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.
/// `TypedLayer<D>::open` always pays for the MPHF too (and doesn't expose `data`
/// once open), so it's the wrong tool for these; this is the other half of
/// the same `D::open(layer_dir)` call, without the MPHF alongside it.
///
/// Takes `(root, i)`, not a pre-built path: the caller names a layer
/// *number* within a partition it already knows the root of, the same
/// vocabulary as [`layer_dir`] and `LayeredMap` — never the `layer_N`
/// naming convention itself, which stays private to this crate.
pub fn open_data<D: LayerData>(root: &Path, i: usize) -> OLMResult<D> {
D::open(&layer_dir(root, i))
}
impl LayerData for () {
type Item = ();
fn open(_layer_dir: &Path) -> OLMResult<Self> { Ok(()) }
fn read(&self, _slot: usize) {}
}
impl LayerData for PersistentCompactIntMatrix {
type Item = Box<[u32]>;
fn open(layer_dir: &Path) -> OLMResult<Self> {
PersistentCompactIntMatrix::open(layer_dir).map_err(OLMError::Io)
}
fn read(&self, slot: usize) -> Box<[u32]> { self.row(slot) }
}
impl LayerData for PersistentBitMatrix {
type Item = Box<[bool]>;
fn open(layer_dir: &Path) -> OLMResult<Self> {
PersistentBitMatrix::open(layer_dir).map_err(OLMError::Io)
}
fn read(&self, slot: usize) -> Box<[bool]> { self.row(slot) }
}
impl LayerData for PersistentSparseBitMatrix {
type Item = Box<[bool]>;
fn open(layer_dir: &Path) -> OLMResult<Self> {
PersistentSparseBitMatrix::open(&layer_dir.join(PRESENCE_DIR)).map_err(OLMError::Io)
}
fn read(&self, slot: usize) -> Box<[bool]> { self.row(slot) }
}
// ── LayerContent ─────────────────────────────────────────────────────────────
/// What a layer's data matrix represents — orthogonal to *how* it's stored
/// on disk (`obicompactvec::StorageKind`'s concern). Only meaningful for
/// `D` that actually carry a matrix: `TypedLayer<()>` (mode 1, set membership,
/// no matrix at all) has neither a `LayerContent` nor a `content()` method
/// — it's a write-time-only state, never a queryable content. Once a layer
/// is closed, "no matrix file" reads back as `Presence` via
/// `PersistentBitMatrix::open`'s own `Implicit` fallback, not as some third
/// "empty" content — see `DevDocMD/implementation/partition_layer_cache.md`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum LayerContent {
Count,
Presence,
}
impl LayerContent {
/// Lightweight disk probe, no MPHF/matrix opened: `counts/` present →
/// `Count`, otherwise `Presence` — an absent/`Implicit` presence
/// matrix still reads as `Presence`, never a separate content.
pub fn detect(layer_dir: &Path) -> LayerContent {
if layer_dir.join(COUNTS_DIR).exists() {
LayerContent::Count
} else {
LayerContent::Presence
}
}
}
/// Implemented only by `D` that carry a real matrix — gates
/// [`TypedLayer::content`](TypedLayer::content) so `TypedLayer<()>` doesn't get it.
pub trait HasLayerContent {
const CONTENT: LayerContent;
}
impl HasLayerContent for PersistentCompactIntMatrix {
const CONTENT: LayerContent = LayerContent::Count;
}
impl HasLayerContent for PersistentBitMatrix {
const CONTENT: LayerContent = LayerContent::Presence;
}
impl HasLayerContent for PersistentSparseBitMatrix {
const CONTENT: LayerContent = LayerContent::Presence;
}
// ── StorageKind passthrough ────────────────────────────────────────────────────
/// Implemented only by `D` that expose their own `storage_kind()` — gates
/// [`TypedLayer::storage_kind`]. `PersistentSparseBitMatrix` has a single,
/// fixed on-disk format (no `Columnar`/`Packed`/`Implicit` variants of its
/// own), so it's deliberately not given an impl: there's nothing to report
/// beyond what `content()` already says.
pub trait HasStorageKind {
fn storage_kind(&self) -> obicompactvec::StorageKind;
}
impl HasStorageKind for PersistentCompactIntMatrix {
fn storage_kind(&self) -> obicompactvec::StorageKind {
PersistentCompactIntMatrix::storage_kind(self)
}
}
impl HasStorageKind for PersistentBitMatrix {
fn storage_kind(&self) -> obicompactvec::StorageKind {
PersistentBitMatrix::storage_kind(self)
}
}
// ── Structures ────────────────────────────────────────────────────────────────
pub struct TypedLayer<D: LayerData = ()> {
mphf: MphfLayer,
data: D,
}
pub struct Hit<T = ()> {
pub slot: usize,
pub data: T,
}
// ── Common read path ──────────────────────────────────────────────────────────
impl<D: LayerData> TypedLayer<D> {
pub fn open(path: &Path, mode: &IndexMode) -> OLMResult<Self> {
let mphf = MphfLayer::open(path, mode)?;
let data = D::open(path)?;
Ok(Self { mphf, data })
}
pub fn query(&self, kmer: CanonicalKmer) -> Option<Hit<D::Item>> {
self.mphf.find(kmer).map(|slot| Hit { slot, data: self.data.read(slot) })
}
/// MPHF + evidence membership check only — no data read. For callers
/// that batch many lookups before touching the matrix at all (e.g. to
/// group hits by column for a later `sub_matrix`/`fill_sub_matrix`
/// sweep), so a plain `query` reading — and discarding — a full row
/// per lookup would be wasted work.
pub fn find_slot(&self, kmer: CanonicalKmer) -> Option<usize> {
self.mphf.find(kmer)
}
pub fn n(&self) -> usize { self.mphf.n() }
/// Raw MPHF lookup: kmer → slot, no membership check.
pub fn index(&self, kmer: CanonicalKmer) -> usize {
self.mphf.index(kmer)
}
/// Batch raw MPHF lookup: kmers → slots, no membership check.
pub fn index_batch(&self, kmers: &[CanonicalKmer]) -> Vec<usize> {
self.mphf.index_batch(kmers)
}
/// Iterate over all canonical kmers in the layer, in deterministic order.
pub fn iter_kmers(&self) -> crate::layer::mphf_layer::KmerIter {
self.mphf.iter_kmers()
}
/// Iterate over all canonical kmers, each paired with its zero-based
/// sequence index in `unitigs.bin`.
pub fn enumerate_kmers(&self) -> std::iter::Enumerate<crate::layer::mphf_layer::KmerIter> {
self.mphf.enumerate_kmers()
}
/// Iterate over the layer's canonical kmers in batches of `n`.
pub fn iter_kmers_batch(&self, n: usize) -> crate::layer::mphf_layer::KmerBatchIter {
self.mphf.iter_kmers_batch(n)
}
/// Iterate over batches, each paired with the zero-based index of the
/// first kmer in the batch.
pub fn enumerate_kmers_batch(&self, n: usize) -> impl Iterator<Item = (usize, Vec<CanonicalKmer>)> + Send + 'static {
self.mphf.enumerate_kmers_batch(n)
}
pub fn unitig_writer(out_dir: &Path) -> OLMResult<UnitigFileWriter> {
MphfLayer::unitig_writer(out_dir)
}
/// Build `unitigs.bin.idx` and `evidence.bin` from `unitigs.bin` and
/// `mphf.bin` already present in `layer_dir`.
/// `block_bits` controls the `.idx` block size (2^block_bits chunks/block).
pub fn build_exact_evidence(layer_dir: &Path, block_bits: u8) -> OLMResult<usize> {
MphfLayer::build_exact_evidence(layer_dir, block_bits)
}
/// Build `fingerprint.bin` from `unitigs.bin` and `mphf.bin` already
/// present in `layer_dir`. `b` — fingerprint bits (1..=64); `z` — Findere
/// consecutive k-mer parameter (≥1).
pub fn build_approx_evidence(layer_dir: &Path, b: u8, z: u8) -> OLMResult<usize> {
MphfLayer::build_approx_evidence(layer_dir, b, z)
}
/// This already-open layer's evidence mode — reads the discriminant
/// already in memory (`MphfLayer`'s own `LayerEvidence`), no disk
/// access. Available regardless of `D`, unlike `content`/`storage_kind`.
pub fn evidence_kind(&self) -> crate::layer::mphf_layer::EvidenceKind {
self.mphf.evidence_kind()
}
}
impl<D: LayerData + HasLayerContent> TypedLayer<D> {
/// This already-open layer's content (`Count`/`Presence`) — a
/// compile-time fact about `D`, not a runtime read.
pub const fn content(&self) -> LayerContent {
D::CONTENT
}
}
impl<D: LayerData + HasStorageKind> TypedLayer<D> {
/// This already-open layer's storage format — delegates to `D`'s own
/// `storage_kind()`, reading a discriminant already in memory.
pub fn storage_kind(&self) -> obicompactvec::StorageKind {
self.data.storage_kind()
}
}
// ── Mode 1 — set membership ───────────────────────────────────────────────────
impl TypedLayer<()> {
pub fn build(out_dir: &Path, block_bits: u8, mode: &IndexMode) -> OLMResult<usize> {
MphfLayer::build(out_dir, block_bits, mode, &mut |_, _| Ok(()))
}
/// Create a presence matrix for a set-membership layer (first merge).
pub fn init_presence_matrix(layer_dir: &Path, n_kmers: usize) -> OLMResult<()> {
let presence_dir = layer_dir.join(PRESENCE_DIR);
fs::create_dir_all(&presence_dir).map_err(OLMError::Io)?;
let mut mb = PersistentBitMatrixBuilder::new(n_kmers, &presence_dir).map_err(OLMError::Io)?;
mb.add_col_ones().map_err(OLMError::Io)?.close().map_err(OLMError::Io)?;
mb.close().map_err(OLMError::Io)
}
}
// ── Mode 2 — count matrix ─────────────────────────────────────────────────────
impl TypedLayer<PersistentCompactIntMatrix> {
pub fn build(
out_dir: &Path,
block_bits: u8,
mode: &IndexMode,
count_of: impl Fn(CanonicalKmer) -> u32,
) -> OLMResult<usize> {
let n = UnitigFileReader::open_sequential(&out_dir.join(UNITIGS_FILE))?.n_kmers();
let counts_dir = out_dir.join(COUNTS_DIR);
let mut mb = PersistentCompactIntMatrixBuilder::new(n, &counts_dir)
.map_err(OLMError::Io)?;
let mut col = mb.add_col().map_err(OLMError::Io)?;
let n_built = MphfLayer::build(out_dir, block_bits, mode, &mut |slot, kmer| {
col.set(slot, count_of(kmer));
Ok(())
})?;
col.close().map_err(OLMError::Io)?;
mb.close().map_err(OLMError::Io)?;
Ok(n_built)
}
pub fn build_from_map(
out_dir: &Path,
block_bits: u8,
mode: &IndexMode,
counts: &HashMap<CanonicalKmer, u32>,
) -> OLMResult<usize> {
Self::build(out_dir, block_bits, mode, |kmer| counts.get(&kmer).copied().unwrap_or(0))
}
}
// ── Mode 2 — count matrix column append ──────────────────────────────────────
impl TypedLayer<PersistentCompactIntMatrix> {
pub fn append_genome_column(
layer_dir: &Path,
value_of: impl Fn(usize) -> u32,
) -> OLMResult<()> {
PersistentCompactIntMatrix::append_column(&layer_dir.join(COUNTS_DIR), value_of)
.map_err(OLMError::Io)
}
/// Number of genome columns in this layer's count matrix.
pub fn n_cols(&self) -> usize {
self.data.n_cols()
}
/// Extract a sub-matrix of counts for the rows at `slots`.
///
/// Returns a column-first `Vec<Vec<u32>>`: one inner `Vec` per genome
/// column, containing the counts for the requested slots in order.
/// Column access is sequential for cache efficiency.
pub fn sub_matrix(&self, slots: &[usize]) -> Vec<Vec<u32>> {
self.data.sub_matrix(slots)
}
/// Like [`sub_matrix`](Self::sub_matrix), but fills caller-provided column
/// buffers to avoid allocating the outer `Vec`.
///
/// `out` must have length equal to the number of genome columns. Each
/// `out[c]` is cleared, resized to `slots.len()`, and filled with the
/// counts for column `c` in the same order as `slots`.
pub fn fill_sub_matrix(&self, slots: &[usize], out: &mut [Vec<u32>]) {
self.data.fill_sub_matrix(slots, out)
}
}
// ── Mode 3 — presence/absence matrix ─────────────────────────────────────────
// ── Mode 3 — presence/absence matrix, generic over dense/sparse storage ──────
//
// `n_cols`/`sub_matrix`/`fill_sub_matrix` are identical in shape for any
// `D` satisfying `BinaryMatrix` — genuinely generic (not just two
// near-identical impl blocks) so `TypedLayer<PersistentSparseBitMatrix>` gets
// them for free, matching `TypedLayer<PersistentBitMatrix>` at the same call
// sites (see `obikphylo::siblings::cache::Mat`). Construction
// (`append_genome_column`/`build_presence`) stays `PersistentBitMatrix`-only
// below: sparse matrices aren't built column-by-column, they're built
// row-by-row from an already-built dense layer
// (`PersistentSparseBitMatrixBuilder::build_from_dense`).
impl<D: LayerData<Item = Box<[bool]>> + BinaryMatrix> TypedLayer<D> {
/// Number of genome columns in this layer's presence matrix — see
/// `PersistentBitMatrix::n_cols`'s docs for the `Implicit` mono-genome
/// special case (always reports `1`, regardless of the index's real
/// genome count).
pub fn n_cols(&self) -> usize {
self.data.n_cols()
}
/// Extract a sub-matrix of presence/absence for the rows at `slots`.
///
/// Returns a column-first `Vec<Vec<bool>>`: one inner `Vec` per genome
/// column, containing the presence bits for the requested slots in order.
/// Column access is sequential for cache efficiency on the dense
/// storage — on sparse storage it's a naive row-by-row decode, see
/// `DevDocMD/architecture/siblings.md`'s sparse-matrix section.
pub fn sub_matrix(&self, slots: &[usize]) -> Vec<Vec<bool>> {
self.data.sub_matrix(slots)
}
/// Like [`sub_matrix`](Self::sub_matrix), but fills caller-provided column
/// buffers to avoid allocating the outer `Vec`.
///
/// `out` must have length equal to the number of genome columns. Each
/// `out[c]` is cleared, resized to `slots.len()`, and filled with the
/// presence bits for column `c` in the same order as `slots`.
pub fn fill_sub_matrix(&self, slots: &[usize], out: &mut [Vec<bool>]) {
self.data.fill_sub_matrix(slots, out)
}
}
impl TypedLayer<PersistentBitMatrix> {
pub fn append_genome_column(
layer_dir: &Path,
value_of: impl Fn(usize) -> bool,
) -> OLMResult<()> {
PersistentBitMatrix::append_column(&layer_dir.join(PRESENCE_DIR), value_of)
.map_err(OLMError::Io)
}
pub fn build_presence(
out_dir: &Path,
block_bits: u8,
mode: &IndexMode,
n_genomes: usize,
present_in: impl Fn(CanonicalKmer, usize) -> bool,
) -> OLMResult<usize> {
let n = UnitigFileReader::open_sequential(&out_dir.join(UNITIGS_FILE))?.n_kmers();
let presence_dir = out_dir.join(PRESENCE_DIR);
let mut mb = PersistentBitMatrixBuilder::new(n, &presence_dir).map_err(OLMError::Io)?;
let mut cols: Vec<_> = (0..n_genomes)
.map(|_| mb.add_col().map_err(OLMError::Io))
.collect::<OLMResult<_>>()?;
let n_built = MphfLayer::build(out_dir, block_bits, mode, &mut |slot, kmer| {
for (g, col) in cols.iter_mut().enumerate() {
col.set(slot, present_in(kmer, g));
}
Ok(())
})?;
for col in cols {
col.close().map_err(OLMError::Io)?;
}
mb.close().map_err(OLMError::Io)?;
Ok(n_built)
}
}
#[cfg(test)]
#[path = "tests/typed_layer.rs"]
mod tests;
+21 -35
View File
@@ -1,36 +1,22 @@
pub mod error;
pub mod meta;
pub mod predicate;
pub mod state;
mod common;
mod distance;
mod dump;
mod dump_layer;
pub mod filter;
mod graph_pipeline;
mod index;
mod index_layer;
mod matrix_store;
mod merge;
mod merge_layer;
mod numa;
mod query_layer;
mod rebuild;
mod rebuild_layer;
mod reindex;
mod select;
mod select_layer;
mod stats;
//! The `Index { Partition { Layer } }` model — see
//! `DevDocMD/implementation/partition_layer_cache.md`'s "Definitions"
//! 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.
pub use error::{OKIError, OKIResult};
pub use distance::{DistanceMetric, DistanceOutput};
pub use filter::{GroupQuorumFilter, KmerFilter, passes_all};
pub use index::KmerIndex;
pub use merge_layer::MergeMode;
pub use meta::{validate_label, GenomeInfo, IndexConfig, IndexMeta, META_FILENAME};
pub use predicate::{GroupFilterParams, MetaPred};
pub use query_layer::{KmerDesc, QueryHit, QueryStats};
pub use select_layer::{AggOp, OutputCol};
pub use state::{IndexState, SENTINEL_COUNTED, SENTINEL_INDEXED, SENTINEL_SCATTERED};
pub use stats::IndexBitsPerKmer;
pub use numa::PartitionRunner;
pub mod algorithms;
pub mod index;
pub mod layer;
pub mod partition;
pub use index::{
validate_label, AggOp, DistanceMetric, DistanceOutput, GenomeInfo, GroupFilterParams,
GroupQuorumFilter, IndexBitsPerKmer, IndexConfig, IndexMeta, IndexState, KmerDesc, KmerFilter,
KmerIndex, MergeMode, MetaPred, OKIError, OKIResult, OutputCol, PartitionRunner, QueryHit,
QueryStats, META_FILENAME, SENTINEL_COUNTED, SENTINEL_INDEXED, SENTINEL_SCATTERED,
passes_all,
};
pub use index::{filter, meta};
+89
View File
@@ -0,0 +1,89 @@
//! The **Partition** tier of the `Index { Partition { Layer } }` model —
//! see `DevDocMD/implementation/partition_layer_cache.md`'s "Definitions"
//! section for the authoritative naming/scope discussion this module
//! implements. [`crate::layer`] holds the **Layer** tier the same way;
//! this module holds the tier above it.
//!
//! [`KmerPartition`] represents one partition's already-open layers — a
//! read cache, opened once and held for the run, not rebuilt per lookup.
//! It does no path computation of its own beyond the shared
//! `crate::layer::layer_dir` naming primitive: which partition, which
//! `index_dir`, how many layers, and `IndexMode` are `crate::index::
//! KmerIndex`'s job to resolve and hand in as plain arguments — this
//! module depends only on `crate::layer` and below, never on
//! `crate::index`, so it cannot reach back for them itself.
//!
//! Not yet wired into any caller: `obikphylo::siblings::cache::
//! PartitionCache` (`Vec<Vec<crate::layer::Layer>>`) and
//! `crate::index::query_layer::QueryLayer` (uncached, bypasses `Layer`
//! entirely) both still reinvent a fragment of this. Migrating them is a
//! separate, deferred step.
use std::path::{Path, PathBuf};
use obikseq::CanonicalKmer;
use crate::layer::{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 [`crate::layer::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>,
}
impl KmerPartition {
/// Open every layer under `index_dir` (`index_dir/layer_0`,
/// `index_dir/layer_1`, ... up to `n_layers`), eagerly — not lazily on
/// first access, so the caller pays the mmap cost once, up front,
/// rather than at an unpredictable point during later lookups.
pub fn open(index_dir: &Path, mode: &IndexMode, n_layers: usize, with_counts: bool) -> OLMResult<Self> {
let layers = (0..n_layers)
.map(|l| Layer::open(&layer_dir(index_dir, l), mode, with_counts))
.collect::<OLMResult<Vec<_>>>()?;
Ok(Self { layers })
}
pub fn n_layers(&self) -> usize {
self.layers.len()
}
pub fn layer(&self, i: usize) -> &Layer {
&self.layers[i]
}
pub fn layers(&self) -> &[Layer] {
&self.layers
}
/// Existence lookup of `kmer` across this partition's layers: tries
/// each in turn, stopping at the first hit and reporting which layer it
/// was — `find_slot`, not a data read, for a plain existence check.
pub fn find(&self, kmer: CanonicalKmer) -> Option<usize> {
self.layers
.iter()
.enumerate()
.find_map(|(li, layer)| layer.find_slot(kmer).map(|_| li))
}
}
#[cfg(test)]
mod tests;
+90
View File
@@ -0,0 +1,90 @@
use super::*;
use obikseq::{set_k, Unitig};
use obiskio::{UnitigFileWriter, DEFAULT_BLOCK_BITS};
use tempfile::tempdir;
fn write_layer(root: &Path, l: usize, seqs: &[&[u8]], n_genomes: usize, mode: &IndexMode) {
let dir = layer_dir(root, l);
std::fs::create_dir_all(&dir).unwrap();
let mut w = UnitigFileWriter::create(&dir.join("unitigs.bin")).unwrap();
for s in seqs {
w.write(&Unitig::from_ascii(s)).unwrap();
}
w.close().unwrap();
crate::layer::TypedLayer::<obicompactvec::PersistentBitMatrix>::build_presence(
&dir,
DEFAULT_BLOCK_BITS,
mode,
n_genomes,
|_kmer, _g| true,
)
.unwrap();
}
#[test]
fn open_reads_every_layer_in_order() {
// k=4, process-wide in test builds — see obikindex::layer's own tests for
// why a fixed, unshared k matters here.
set_k(4);
let dir = tempdir().unwrap();
let mode = IndexMode::Exact;
write_layer(dir.path(), 0, &[b"AAATCTA"], 1, &mode);
write_layer(dir.path(), 1, &[b"CTTCGCC"], 1, &mode);
let partition = KmerPartition::open(dir.path(), &mode, 2, false).unwrap();
assert_eq!(partition.n_layers(), 2);
assert_eq!(partition.layers().len(), 2);
}
#[test]
fn find_reports_the_first_layer_that_carries_the_kmer() {
set_k(4);
let dir = tempdir().unwrap();
let mode = IndexMode::Exact;
write_layer(dir.path(), 0, &[b"AAATCTA"], 1, &mode);
write_layer(dir.path(), 1, &[b"CTTCGCC"], 1, &mode);
let partition = KmerPartition::open(dir.path(), &mode, 2, false).unwrap();
let in_layer_0 = obiskio::CanonicalKmerIter::new(&layer_dir(dir.path(), 0).join("unitigs.bin"))
.unwrap()
.next()
.unwrap();
let in_layer_1 = obiskio::CanonicalKmerIter::new(&layer_dir(dir.path(), 1).join("unitigs.bin"))
.unwrap()
.next()
.unwrap();
assert_eq!(partition.find(in_layer_0), Some(0));
assert_eq!(partition.find(in_layer_1), Some(1));
}
#[test]
fn find_returns_none_for_an_absent_kmer() {
set_k(4);
let dir = tempdir().unwrap();
let mode = IndexMode::Exact;
write_layer(dir.path(), 0, &[b"AAATCTA"], 1, &mode);
let partition = KmerPartition::open(dir.path(), &mode, 1, false).unwrap();
let absent = obiskio::CanonicalKmerIter::new(&layer_dir(dir.path(), 0).join("unitigs.bin"))
.unwrap()
.next()
.unwrap();
// Not actually absent from the MPHF domain in a general sense, but a
// kmer that was never written to this partition's only layer is what
// "absent" means here — build a second, disjoint layer purely to
// source a genuinely foreign kmer instead of relying on ptr_hash's
// undefined behaviour for out-of-domain queries.
let other_dir = tempdir().unwrap();
write_layer(other_dir.path(), 0, &[b"GGGGTAC"], 1, &mode);
let foreign = obiskio::CanonicalKmerIter::new(&layer_dir(other_dir.path(), 0).join("unitigs.bin"))
.unwrap()
.next()
.unwrap();
assert_eq!(partition.find(absent), Some(0));
assert_eq!(partition.find(foreign), None);
}