refactor: merge KmerPartitions into KmerIndex and rename obikpartition

Consolidates partition logic, metadata storage, and layer management directly into KmerIndex. Renames obikpartitionner to obikpartition, retaining only PartitionRouter for superkmer routing. Removes intermediate .partition() accessors in favor of direct methods on the index and updates PartitionCache::build to accept &KmerIndex directly. Derives n_partitions from config.n_bits and consolidates k-mer/minimizer sizes into IndexMeta.config. Fixes a regression where PartitionRouter::open incorrectly defaulted to closed.
This commit is contained in:
Eric Coissac
2026-08-20 20:23:58 +02:00
parent b5ec0122d0
commit 6c860f120f
49 changed files with 357 additions and 427 deletions
@@ -69,6 +69,79 @@ let a per-partition type hold layers of mixed `D` (today `LayeredMap<D>`
can't, being monomorphic — see "The gap in `obilayeredmap`'s existing
cache" below).
## Major restructuring (2026-08-20): `KmerPartitions` merged into `KmerIndex`
Prompted by a direct question: why keep `KmerIndex`/`KmerPartitions` split
when, one level down, `KmerPartitions` is going to directly hold
`Vec<KmerPartition>` rather than being split again into
"collection-holder" + "collection"? Investigating the actual justification
("`KmerPartitions` has an independent lifecycle, used before an index
exists") turned out to be **false** — `KmerPartitions::create` was called
in exactly one place, inside `KmerIndex::create`, and every
`open_with_config` reopen outside `KmerIndex`'s own constructors was a
redundant re-derivation of a `KmerPartitions` already reachable via
`index.partition()` (the exact kind of duplication this whole doc has been
tracking). Once that was gone, so was the reason to keep them separate.
Second correction, from the same conversation: `obikpartitionner` had
accumulated query/merge/select/rebuild/dump/distance logic that has
nothing to do with partitioning super-kmers — it operates on *layers*,
which don't exist yet at the phase `obikpartitionner` is actually
responsible for (scatter → dereplicate → count, all pre-layer). That
logic moved to `obikindex`, which already depends on `obilayeredmap` and
never needed `obikpartitionner` for it. No crate-dependency inversion was
needed — `obikindex → obikpartitionner` stays the same direction as before.
**Result:**
- `obikpartitionner` (renamed back from `obikpartition`) now contains only
`PartitionRouter` (superkmer routing: `write`/`write_batch`/`flush`/
`close`, `dereplicate`, `count_kmer`, `KmerSpectrum`) and the
`partition_dir(root, i)` naming primitive both `PartitionRouter` and
`KmerIndex` build on. `KmerPartitions` no longer exists as a type.
- `KmerIndex` (`obikindex`) absorbed `KmerPartitions`'s read-side entirely:
`partition_dir`/`index_dir`/`layer_dir`/`partition_meta`/`n_layers`/
`partition_mode`/`n_partitions` (the last now derived from
`2^config.n_bits`, no longer a stored, independently-set duplicate field
— `kmer_size`/`minimizer_size` used to be double-stored, in both
`KmerPartitions` and `IndexMeta.config`, a latent-drift risk flagged
earlier in this doc; now single-sourced from `IndexMeta.config`). Seven
whole files moved from `obikpartitionner` into `obikindex` verbatim as
`impl KmerIndex` blocks, kept as separate files (not merged into
existing same-topic files): `index_layer.rs`, `query_layer.rs`,
`merge_layer/`, `select_layer.rs`, `rebuild_layer.rs`, `dump_layer.rs`,
plus `distance.rs`'s `count_store`/`presence_store` (renamed
`matrix_store.rs` to avoid colliding with `obikindex`'s own pre-existing
`distance.rs`), and their shared support (`common.rs`'s `load_meta`/
`olm_to_sk`, `filter.rs`, `graph_pipeline.rs`).
- `obikphylo::siblings::cache::PartitionCache::build` now takes `&KmerIndex`
directly instead of a separately-opened `&KmerPartitions` — this deleted
the redundant-reopen pattern at all 8 call sites
(`alignment`/`build`/`cardinality`/`distance`/`entropy`×2/
`sankoff_bundle`/`stats`), the same bug flagged earlier in this
conversation as a side effect of investigating the false "independent
lifecycle" claim.
- `KmerIndex::partition()`/`partition_mut()` are gone; `scatter()`
(`obikmer`) and any write-side code get a transient `PartitionRouter` via
`KmerIndex::partition_router()`.
- A real bug caught by the test suite during this move:
`PartitionRouter::open` initially defaulted to `closed: true` (inherited
from `KmerPartitions::open_with_config`'s old read-only-reopen
semantics), which broke every write through a router obtained via
`partition_router()`. Fixed — `PartitionRouter` is exclusively a
write/processing tool now, so `open` always starts open.
Full workspace test suite green (0 failed) after, including all 27
`obikphylo::siblings` tests.
Still open: (1)/(2) themselves — `AnyLayer` (or whatever it ends up named;
`AnyLayer` was rejected as a placeholder, no replacement chosen yet) and
`KmerPartition` (singular, one partition's open layers) are not built yet.
`obikphylo::siblings::cache::{Mat, PartitionCache}` and
`obikpartitionner::query_layer` (now `obikindex::query_layer`)'s
`QueryLayer` still each independently bundle MPHF+matrix — unchanged by
this restructuring, which was purely about *where* code lives, not about
building the heterogeneous-layer cache itself.
## The problem
Reading a layer's data (MPHF + matrix) is not free: `MphfLayer::open` mmaps
+8 -5
View File
@@ -1595,18 +1595,26 @@ name = "obikindex"
version = "0.1.0"
dependencies = [
"anyhow",
"cacheline-ef",
"crossbeam-channel",
"epserde",
"hwlocality",
"indicatif",
"memmap2",
"ndarray",
"niffler",
"obicompactvec",
"obidebruinj",
"obikentropy",
"obikpartitionner",
"obikseq",
"obilayeredmap",
"obipipeline",
"obiread",
"obiskio",
"obisys",
"obitaxonomy",
"ptr_hash",
"rayon",
"serde",
"serde_json",
@@ -1653,16 +1661,11 @@ version = "0.1.0"
dependencies = [
"cacheline-ef",
"epserde",
"indicatif",
"memmap2",
"niffler",
"obicompactvec",
"obidebruinj",
"obikentropy",
"obikrope",
"obikseq",
"obilayeredmap",
"obipipeline",
"obiread",
"obiskbuilder",
"obiskio",
+8
View File
@@ -11,6 +11,14 @@ obiskio = { path = "../obiskio" }
obisys = { path = "../obisys" }
obicompactvec = { path = "../obicompactvec" }
obilayeredmap = { path = "../obilayeredmap" }
obidebruinj = { path = "../obidebruinj" }
obipipeline = { path = "../obipipeline" }
obikentropy = { path = "../obikentropy" }
cacheline-ef = "1.1"
epserde = "0.8"
ptr_hash = "1.1"
niffler = "3.0.0"
memmap2 = "0.9.10"
ndarray = "0.17"
rayon = "1"
crossbeam-channel = "0.5"
+4 -4
View File
@@ -25,16 +25,16 @@ fn main() -> anyhow::Result<()> {
let mut first_mismatch = None;
for part in 0..n_parts {
let index_dir_sparse = sparse.partition().index_dir(part);
let index_dir_dense = dense.partition().index_dir(part);
let index_dir_sparse = sparse.index_dir(part);
let index_dir_dense = dense.index_dir(part);
if !index_dir_sparse.exists() || !index_dir_dense.exists() {
continue;
}
for layer in 0..n_layers {
let layer_dir_sparse = sparse.partition().layer_dir(part, layer);
let layer_dir_dense = dense.partition().layer_dir(part, layer);
let layer_dir_sparse = sparse.layer_dir(part, layer);
let layer_dir_dense = dense.layer_dir(part, layer);
if !layer_dir_sparse.exists() || !layer_dir_dense.exists() {
continue;
+2 -2
View File
@@ -74,7 +74,7 @@ impl KmerIndex {
if use_counts {
let stores: Vec<_> = (0..n_parts)
.into_par_iter()
.map(|i| self.partition.count_store(i).map_err(OKIError::Partition))
.map(|i| self.count_store(i).map_err(OKIError::Partition))
.collect::<OKIResult<_>>()?;
let global = LayeredStore::new(stores);
@@ -105,7 +105,7 @@ impl KmerIndex {
} else {
let stores: Vec<_> = (0..n_parts)
.into_par_iter()
.map(|i| self.partition.presence_store(i).map_err(OKIError::Partition))
.map(|i| self.presence_store(i).map_err(OKIError::Partition))
.collect::<OKIResult<_>>()?;
let global = LayeredStore::new(stores);
+5 -5
View File
@@ -5,7 +5,7 @@ use rayon::prelude::*;
use crate::error::{OKIError, OKIResult};
use crate::index::KmerIndex;
use obikpartitionner::KmerFilter;
use crate::KmerFilter;
impl KmerIndex {
/// Write a CSV table of all indexed kmers to `out`.
@@ -69,14 +69,14 @@ impl KmerIndex {
}
};
if debug {
self.partition
self
.iter_partition_kmers_located(i, use_counts, n_genomes, filters, |part, layer, kmer, row| {
let seq = String::from_utf8(kmer.to_ascii()).unwrap_or_else(|_| "?".repeat(kmer_size));
try_write(&mut buf, &row, &format!("{part},{layer},{seq}"))
})
.map_err(OKIError::Partition)?;
} else {
self.partition
self
.iter_partition_kmers(i, use_counts, n_genomes, filters, |kmer, row| {
let seq = String::from_utf8(kmer.to_ascii()).unwrap_or_else(|_| "?".repeat(kmer_size));
try_write(&mut buf, &row, &seq)
@@ -91,7 +91,7 @@ impl KmerIndex {
(0..n).into_par_iter().map(|i| {
let mut buf = Vec::<u8>::new();
if debug {
self.partition
self
.iter_partition_kmers_located(i, use_counts, n_genomes, filters, |part, layer, kmer, row| {
let seq = String::from_utf8(kmer.to_ascii()).unwrap_or_else(|_| "?".repeat(kmer_size));
write_row(&mut buf, &row, &format!("{part},{layer},{seq}"));
@@ -99,7 +99,7 @@ impl KmerIndex {
})
.map_err(OKIError::Partition)?;
} else {
self.partition
self
.iter_partition_kmers(i, use_counts, n_genomes, filters, |kmer, row| {
let seq = String::from_utf8(kmer.to_ascii()).unwrap_or_else(|_| "?".repeat(kmer_size));
write_row(&mut buf, &row, &seq);
@@ -4,7 +4,7 @@ use obilayeredmap::{IndexMode, MphfLayer, OLMError};
use obiskio::{SKError, SKResult, UnitigFileReader};
use crate::filter::{KmerFilter, passes_all};
use crate::partition::KmerPartitions;
use crate::index::KmerIndex;
fn olm_to_sk(e: OLMError) -> SKError {
match e {
@@ -16,7 +16,7 @@ fn olm_to_sk(e: OLMError) -> SKError {
}
}
impl KmerPartitions {
impl KmerIndex {
/// Iterate all indexed kmers in partition `part`, calling `cb(kmer, row)` for each
/// kmer that passes every filter in `filters`.
///
@@ -41,7 +41,7 @@ impl KmerPartitions {
return Ok(true);
}
let index_mode = self.index_mode(part).unwrap_or(IndexMode::Exact);
let index_mode = self.partition_mode(part).unwrap_or(IndexMode::Exact);
let mut l = 0;
loop {
@@ -131,7 +131,7 @@ impl KmerPartitions {
return Ok(true);
}
let index_mode = self.index_mode(part).unwrap_or(IndexMode::Exact);
let index_mode = self.partition_mode(part).unwrap_or(IndexMode::Exact);
let mut layer = 0;
loop {
+76 -65
View File
@@ -2,13 +2,15 @@ use std::collections::BTreeMap;
use std::fs;
use std::path::{Path, PathBuf};
use obikpartitionner::{KmerPartitions, KmerSpectrum, PARTITIONS_SUBDIR};
use obikpartitionner::{KmerSpectrum, PartitionRouter};
use obilayeredmap::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};
@@ -16,7 +18,6 @@ use crate::state::{IndexState, SENTINEL_COUNTED, SENTINEL_INDEXED, SENTINEL_SCAT
pub struct KmerIndex {
pub(crate) root_path: PathBuf,
pub(crate) meta: IndexMeta,
pub(crate) partition: KmerPartitions,
}
impl KmerIndex {
@@ -31,13 +32,7 @@ impl KmerIndex {
force: bool,
) -> OKIResult<Self> {
let root_path = path.as_ref().to_owned();
let partition = KmerPartitions::create(
&root_path,
config.n_bits,
config.kmer_size,
config.minimizer_size,
force,
)?;
PartitionRouter::create(&root_path, config.n_bits, force)?;
set_k(config.kmer_size);
set_m(config.minimizer_size);
let mut meta = IndexMeta::new(config);
@@ -45,11 +40,7 @@ impl KmerIndex {
meta.genomes.push(info);
}
meta.write(&root_path)?;
Ok(Self {
root_path,
meta,
partition,
})
Ok(Self { root_path, meta })
}
pub fn open<P: AsRef<Path>>(path: P) -> OKIResult<Self> {
@@ -57,17 +48,7 @@ impl KmerIndex {
let meta = IndexMeta::read(&root_path).map_err(OKIError::Io)?;
set_k(meta.config.kmer_size);
set_m(meta.config.minimizer_size);
let partition = KmerPartitions::open_with_config(
&root_path,
meta.config.kmer_size,
meta.config.minimizer_size,
meta.config.n_bits,
)?;
Ok(Self {
root_path,
meta,
partition,
})
Ok(Self { root_path, meta })
}
/// Return `true` if `path` contains an `index.meta` file.
@@ -96,25 +77,17 @@ impl KmerIndex {
}
/// Lay out a fresh index skeleton at `output`: the root directory,
/// `index.meta` (from `meta`), and an opened, empty partition set.
/// `index.meta` (from `meta`), and an empty partition layout.
///
/// For construction paths that build partitions from scratch (`select`,
/// `rebuild`). `merge` bootstraps by copying a source index instead, so
/// it does not use this.
pub(crate) fn create_skeleton<P: AsRef<Path>>(
output: P,
meta: &IndexMeta,
) -> OKIResult<KmerPartitions> {
pub(crate) fn create_skeleton<P: AsRef<Path>>(output: P, meta: &IndexMeta) -> OKIResult<KmerIndex> {
let output = output.as_ref();
fs::create_dir_all(output).map_err(OKIError::Io)?;
meta.write(output).map_err(OKIError::Io)?;
fs::create_dir_all(output.join(PARTITIONS_SUBDIR)).map_err(OKIError::Io)?;
Ok(KmerPartitions::open_with_config(
output,
meta.config.kmer_size,
meta.config.minimizer_size,
meta.config.n_bits,
)?)
PartitionRouter::create(output, meta.config.n_bits, false)?;
Ok(KmerIndex { root_path: output.to_owned(), meta: meta.clone() })
}
/// Mark `output` as fully indexed, pack its column matrices, and reopen it.
@@ -139,9 +112,7 @@ impl KmerIndex {
IndexState::detect(&self.root_path).unwrap_or(IndexState::Empty)
}
/// The index's root directory — needed by out-of-crate extension code
/// (e.g. `obikphylo`) that opens its own `KmerPartition` handle onto
/// the same on-disk index.
/// The index's root directory.
pub fn root_path(&self) -> &Path {
&self.root_path
}
@@ -185,7 +156,48 @@ impl KmerIndex {
&self.meta.genomes
}
pub fn n_partitions(&self) -> usize {
self.partition.n_partitions()
1usize << self.meta.config.n_bits
}
/// Path of partition `i`'s raw directory (`partitions/part_{i:05}`) —
/// the on-disk naming convention `obikpartitionner::PartitionRouter`
/// also writes to (raw/dereplicated superkmer files, `mphf1.bin`,
/// `counts1.bin`), the single point of agreement between the two.
pub fn partition_dir(&self, i: usize) -> PathBuf {
obikpartitionner::partition_dir(&self.root_path, i)
}
/// Path of partition `i`'s layered-index directory (`<partition>/index`).
pub fn index_dir(&self, i: usize) -> PathBuf {
self.partition_dir(i).join("index")
}
/// 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)
}
/// 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>`).
pub fn partition_meta(&self, i: usize) -> obiskio::SKResult<PartitionMeta> {
load_meta(&self.index_dir(i), "partition_meta")
}
/// Number of layers in partition `i` — see [`partition_meta`](Self::partition_meta).
pub fn n_layers(&self, i: usize) -> obiskio::SKResult<usize> {
Ok(self.partition_meta(i)?.n_layers)
}
/// Evidence mode partition `i` was actually built with — see
/// [`partition_meta`](Self::partition_meta). Distinct from
/// [`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> {
Ok(self.partition_meta(i)?.mode)
}
/// Number of layers per partition.
@@ -194,13 +206,7 @@ impl KmerIndex {
/// homogeneous across all partitions — reading it off partition 0
/// is enough, no need to scan every partition.
pub fn n_layers_per_partition(&self) -> OKIResult<usize> {
Ok(self.partition.n_layers(0)?)
}
/// Expose the inner partition so the caller can run scatter into it.
/// Call `mark_scattered` once scatter is complete.
pub fn partition_mut(&mut self) -> &mut KmerPartitions {
&mut self.partition
Ok(self.n_layers(0)?)
}
/// Mark scatter as complete and write `scatter.done`.
@@ -217,6 +223,14 @@ impl KmerIndex {
Ok(())
}
/// Open a fresh [`PartitionRouter`] onto this index's partition layout
/// — the write-side handle for `scatter`, or for `dereplicate_and_count`
/// below. Transient: no state is kept in `KmerIndex` itself between
/// calls, only on disk.
pub fn partition_router(&self) -> OKIResult<PartitionRouter> {
Ok(PartitionRouter::open(&self.root_path, self.meta.config.n_bits)?)
}
/// Dereplicate all partitions then compute kmer counts.
///
/// Writes `spectrums/{label}.json` and touches `count.done` upon completion.
@@ -226,12 +240,14 @@ impl KmerIndex {
keep_intermediate: bool,
rep: &mut Reporter,
) -> OKIResult<()> {
let router = self.partition_router()?;
let t = Stage::start("dereplicate");
self.partition.dereplicate()?;
router.dereplicate()?;
rep.push(t.stop());
let t = Stage::start("count_kmer");
let spectrum = self.partition.count_kmer(keep_intermediate)?;
let spectrum = router.count_kmer(keep_intermediate)?;
rep.push(t.stop());
self.write_spectrum(&spectrum)?;
@@ -273,7 +289,7 @@ impl KmerIndex {
keep_intermediate: bool,
rep: &mut Reporter,
) -> OKIResult<()> {
let n = self.partition.n_partitions();
let n = self.n_partitions();
let t = Stage::start("index");
let with_counts = self.meta.config.with_counts;
let evidence = self.meta.config.evidence.clone();
@@ -287,7 +303,7 @@ impl KmerIndex {
.run(
&order,
|i| {
self.partition.build_index_layer(
self.build_index_layer(
i,
min_ab,
max_ab,
@@ -311,7 +327,7 @@ impl KmerIndex {
if !keep_intermediate {
for i in 0..n {
self.partition.remove_build_artifacts(i);
self.remove_build_artifacts(i);
}
}
@@ -320,14 +336,9 @@ impl KmerIndex {
Ok(())
}
/// Borrow the inner partition for direct superkmer-level queries.
pub fn partition(&self) -> &KmerPartitions {
&self.partition
}
/// Path to the unitigs file for partition `part`, layer `layer`.
pub fn layer_unitigs_path(&self, part: usize, layer: usize) -> PathBuf {
self.partition.layer_dir(part, layer).join("unitigs.bin")
self.layer_dir(part, layer).join("unitigs.bin")
}
/// Pack all partition matrices into single-file format (presence → .pbmx, counts → .pcmx).
@@ -349,13 +360,13 @@ impl KmerIndex {
crate::numa::PartitionRunner::new().run(
&order,
|i| -> OKIResult<()> {
let index_dir = self.partition.index_dir(i);
let index_dir = self.index_dir(i);
if !index_dir.exists() {
return Ok(());
}
let n_layers = self.partition.n_layers(i)?;
let n_layers = self.n_layers(i)?;
for l in 0..n_layers {
let layer_dir = self.partition.layer_dir(i, l);
let layer_dir = self.layer_dir(i, l);
let presence_dir = layer_dir.join("presence");
let counts_dir = layer_dir.join("counts");
if presence_dir.exists() {
@@ -391,11 +402,11 @@ impl KmerIndex {
let errors: Vec<_> = (0..n)
.into_par_iter()
.filter_map(|i| {
let index_dir = self.partition.index_dir(i);
let index_dir = self.index_dir(i);
if !index_dir.exists() {
return None;
}
let n_layers = match self.partition.n_layers(i) {
let n_layers = match self.n_layers(i) {
Ok(n) => n,
Err(e) => {
return Some(OKIError::Io(std::io::Error::new(
@@ -405,7 +416,7 @@ impl KmerIndex {
}
};
for l in 0..n_layers {
let layer_dir = self.partition.layer_dir(i, l);
let layer_dir = self.layer_dir(i, l);
let meta_path = layer_dir.join(LayerMeta::FILENAME);
if meta_path.exists() {
continue;
@@ -12,7 +12,7 @@ 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::partition::KmerPartitions;
use crate::index::KmerIndex;
type Mphf = PtrHash<u64, CubicEps, CachelineEfVec<Vec<CachelineEf>>, Xx64, Vec<u8>>;
@@ -24,7 +24,7 @@ fn remove_if_exists(path: &std::path::Path) {
}
}
impl KmerPartitions {
impl KmerIndex {
/// Build the layered MPHF index for partition `i`.
///
/// Returns the number of canonical k-mers indexed, or 0 if the partition
+14 -1
View File
@@ -2,22 +2,35 @@ 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;
pub use error::{OKIError, OKIResult};
pub use distance::{DistanceMetric, DistanceOutput};
pub use filter::{GroupQuorumFilter, KmerFilter, passes_all};
pub use index::KmerIndex;
pub use merge::MergeMode;
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;
@@ -3,9 +3,9 @@ use obilayeredmap::{LayeredStore, open_data};
use obiskio::SKResult;
use crate::common::{load_meta, olm_to_sk};
use crate::partition::KmerPartitions;
use crate::index::KmerIndex;
impl KmerPartitions {
impl KmerIndex {
/// Open all count matrices for partition `part`, one per layer.
/// Layers without a `counts/` directory are skipped.
pub fn count_store(&self, part: usize) -> SKResult<LayeredStore<PersistentCompactIntMatrix>> {
+5 -6
View File
@@ -13,7 +13,7 @@ use crate::index::KmerIndex;
use crate::meta::{GenomeInfo, IndexMeta};
use crate::state::{IndexState, SENTINEL_INDEXED};
pub use obikpartitionner::MergeMode;
pub use crate::merge_layer::MergeMode;
// ── per-partition diagnostic record ──────────────────────────────────────────
@@ -189,13 +189,12 @@ impl KmerIndex {
let t = Stage::start("merge_partitions");
let pb = progress_bar("merge", n_partitions as u64, "partitions");
let dst_partition = &dst.partition;
let block_bits = dst.meta.config.block_bits;
// Pre-build source list once (avoid rebuilding per partition)
let srcs: Vec<(&obikpartitionner::KmerPartitions, usize)> = remaining_sources
let srcs: Vec<(&KmerIndex, usize)> = remaining_sources
.iter()
.map(|s| (&s.partition, s.meta.genomes.len()))
.map(|s| (*s, s.meta.genomes.len()))
.collect();
// Per-partition unitig byte sizes across remaining sources (stat() only)
@@ -225,7 +224,7 @@ impl KmerIndex {
.run(
&order,
|i| {
dst_partition.merge_partition(
dst.merge_partition(
i,
srcs,
mode,
@@ -418,7 +417,7 @@ fn is_trivial(src: &KmerIndex, mode: MergeMode) -> bool {
}
fn index_unitig_size(src: &KmerIndex) -> u64 {
let n = src.partition.n_partitions();
let n = src.n_partitions();
(0..n).map(|i| partition_unitig_bytes(src, i)).sum()
}
@@ -25,7 +25,7 @@ use obiskio::{SKError, SKResult, UnitigFileReader};
use crate::common::{ColBuilder, load_meta, olm_to_sk};
use crate::graph_pipeline::{build_graph, materialize_layer};
use crate::partition::KmerPartitions;
use crate::index::KmerIndex;
mod src_layer;
@@ -198,7 +198,7 @@ mod matrix_builder_tests {
// ── KmerPartition::merge_partition ────────────────────────────────────────────
impl KmerPartitions {
impl KmerIndex {
/// Merge `sources` into destination partition `i`.
///
/// Each entry in `sources` is `(partition, n_genomes)` where `n_genomes` is
@@ -211,7 +211,7 @@ impl KmerPartitions {
pub fn merge_partition(
&self,
i: usize,
sources: &[(&KmerPartitions, usize)],
sources: &[(&KmerIndex, usize)],
mode: MergeMode,
n_dst_genomes: usize,
block_bits: u8,
+1 -1
View File
@@ -1,6 +1,6 @@
use std::collections::HashMap;
use obikpartitionner::GroupQuorumFilter;
use crate::GroupQuorumFilter;
use obitaxonomy::{TaxPath, TaxPattern};
use crate::meta::{GenomeInfo, IndexMeta};
@@ -6,7 +6,7 @@ use obikseq::CanonicalKmer;
use obilayeredmap::{IndexMode, MphfLayer, OLMError};
use obiskio::{SKError, SKResult};
use crate::partition::KmerPartitions;
use crate::index::KmerIndex;
fn olm_to_sk(e: OLMError) -> SKError {
match e {
@@ -143,7 +143,7 @@ pub enum QueryHit<'a> {
// ── KmerPartition::query_partition_with ──────────────────────────────────────
impl KmerPartitions {
impl KmerIndex {
/// Query a single partition for a pre-deduplicated map of canonical
/// k-mers → their occurrences (`seq_idx`, `pos`) in the query batch.
///
+3 -4
View File
@@ -1,6 +1,6 @@
use std::path::Path;
use obikpartitionner::{KmerFilter, MergeMode};
use crate::{KmerFilter, MergeMode};
use obisys::{Reporter, Stage, progress_bar};
use tracing::info;
@@ -46,7 +46,7 @@ impl KmerIndex {
meta.genomes = src.meta.genomes.clone();
let n_genomes = src.meta.genomes.len();
let n_partitions = src.partition.n_partitions();
let n_partitions = src.n_partitions();
// ── Create an empty destination KmerPartition ─────────────────────────
let dst_partition = KmerIndex::create_skeleton(output, &meta)?;
@@ -59,14 +59,13 @@ impl KmerIndex {
let t = Stage::start("rebuild");
let pb = progress_bar("rebuild", n_partitions as u64, "partitions");
let src_partition = &src.partition;
let block_bits = meta.config.block_bits;
let order: Vec<usize> = (0..n_partitions).collect();
let runner = crate::numa::PartitionRunner::new();
runner.run(
&order,
|i| dst_partition.rebuild_partition(src_partition, i, filters, mode, n_genomes, block_bits),
|i| dst_partition.rebuild_partition(src, i, filters, mode, n_genomes, block_bits),
|_, _, _| { pb.inc(1); },
).map_err(OKIError::Partition)?;
@@ -14,7 +14,7 @@ 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::partition::KmerPartitions;
use crate::index::KmerIndex;
// ── Builders — pair matrix builder + column builders for one mode ─────────────
@@ -193,7 +193,7 @@ fn iter_src_layers(
// ── KmerPartition::rebuild_partition ─────────────────────────────────────────
impl KmerPartitions {
impl KmerIndex {
/// Rebuild partition `i` from `src` into `self` (an empty destination partition).
///
/// Only k-mers whose per-genome row passes all `filters` are written.
@@ -203,7 +203,7 @@ impl KmerPartitions {
/// `n_genomes` is the number of genome columns in the source (and destination).
pub fn rebuild_partition(
&self,
src: &KmerPartitions,
src: &KmerIndex,
i: usize,
filters: &[Box<dyn KmerFilter>],
mode: MergeMode,
+6 -7
View File
@@ -1,4 +1,3 @@
use obikpartitionner::KmerPartitions;
use obilayeredmap::{IndexMode, layer::Layer};
use obisys::{Reporter, Stage, progress_bar};
use std::fs;
@@ -35,7 +34,7 @@ impl KmerIndex {
return Err(OKIError::NotIndexed(self.root_path.clone()));
}
let n = self.partition.n_partitions();
let n = self.n_partitions();
info!(
"reindex {} partition(s): {:?} → {:?}",
n, self.meta.config.evidence, target,
@@ -49,7 +48,7 @@ impl KmerIndex {
runner.run(
&order,
|i| {
reindex_partition(&self.partition, i, &target, block_bits)
reindex_partition(self, i, &target, block_bits)
.map_err(|e| OKIError::InvalidInput(format!("partition {i}: {e}")))
},
|_, _, _| {
@@ -71,19 +70,19 @@ impl KmerIndex {
/// Process all layers of one partition's index directory.
fn reindex_partition(
partition: &KmerPartitions,
index: &KmerIndex,
i: usize,
target: &IndexMode,
block_bits: u8,
) -> OKIResult<()> {
if !partition.index_dir(i).exists() {
if !index.index_dir(i).exists() {
return Ok(());
}
let n_layers = partition
let n_layers = index
.n_layers(i)
.map_err(|e| OKIError::InvalidInput(e.to_string()))?;
for layer_idx in 0..n_layers {
reindex_layer(&partition.layer_dir(i, layer_idx), target, block_bits)?;
reindex_layer(&index.layer_dir(i, layer_idx), target, block_bits)?;
}
Ok(())
}
+6 -15
View File
@@ -1,6 +1,6 @@
use std::path::Path;
use obikpartitionner::{KmerPartitions, OutputCol};
use crate::OutputCol;
use obisys::{Reporter, Stage, progress_bar};
use tracing::info;
@@ -41,7 +41,7 @@ impl KmerIndex {
.collect();
let n_src_genomes = src.meta.genomes.len();
let n_partitions = src.partition.n_partitions();
let n_partitions = src.n_partitions();
let dst_partition = KmerIndex::create_skeleton(output, &meta)?;
@@ -54,7 +54,6 @@ impl KmerIndex {
let t = Stage::start("select");
let pb = progress_bar("select", n_partitions as u64, "partitions");
let src_partition = &src.partition;
let order: Vec<usize> = (0..n_partitions).collect();
let runner = crate::numa::PartitionRunner::new();
@@ -63,7 +62,7 @@ impl KmerIndex {
&order,
|i| {
dst_partition.select_partition(
src_partition,
src,
i,
specs,
n_src_genomes,
@@ -99,14 +98,7 @@ impl KmerIndex {
}
let n_src_genomes = self.meta.genomes.len();
let n_partitions = self.partition.n_partitions();
let src_partition = KmerPartitions::open_with_config(
&self.root_path,
self.meta.config.kmer_size,
self.meta.config.minimizer_size,
self.meta.config.n_bits,
)?;
let n_partitions = self.n_partitions();
info!(
"select (in-place): {} partition(s), {} source genome(s) → {} output column(s)",
@@ -118,15 +110,14 @@ impl KmerIndex {
let t = Stage::start("select");
let pb = progress_bar("select", n_partitions as u64, "partitions");
let partition = &self.partition;
let order: Vec<usize> = (0..n_partitions).collect();
let runner = crate::numa::PartitionRunner::new();
runner
.run(
&order,
|i| {
partition.select_partition(
&src_partition,
self.select_partition(
self,
i,
specs,
n_src_genomes,
@@ -9,7 +9,7 @@ use obicompactvec::{
use obilayeredmap::OLMError;
use obiskio::{SKError, SKResult};
use crate::partition::KmerPartitions;
use crate::index::KmerIndex;
// ── AggOp ─────────────────────────────────────────────────────────────────────
@@ -170,7 +170,7 @@ fn fill_builders(
// ── KmerPartition::select_partition ──────────────────────────────────────────
impl KmerPartitions {
impl KmerIndex {
/// Rewrite the data matrices of partition `i` in `src` into `self`.
///
/// `specs` defines the output columns (projection/aggregation).
@@ -178,7 +178,7 @@ impl KmerPartitions {
/// `in_place` — `self` and `src` share the same root; write to temp dirs then swap.
pub fn select_partition(
&self,
src: &KmerPartitions,
src: &KmerIndex,
i: usize,
specs: &[OutputCol],
_n_src_genomes: usize,
+6 -6
View File
@@ -89,13 +89,13 @@ impl KmerIndex {
let (n_kmers, mphf_b, evidence_b, matrix_b) = (0..n)
.into_par_iter()
.map(|i| {
let index_dir = self.partition.index_dir(i);
let index_dir = self.index_dir(i);
if !index_dir.exists() { return (0usize, 0u64, 0u64, 0u64); }
let n_layers = self.partition.n_layers(i).unwrap_or(0);
let n_layers = self.n_layers(i).unwrap_or(0);
(0..n_layers).fold((0usize, 0u64, 0u64, 0u64), |acc, l| {
let lb = layer_bytes(&self.partition.layer_dir(i, l));
let lb = layer_bytes(&self.layer_dir(i, l));
(acc.0 + lb.n_kmers, acc.1 + lb.mphf, acc.2 + lb.evidence, acc.3 + lb.matrix)
})
})
@@ -138,13 +138,13 @@ impl KmerIndex {
let mut counts = vec![0u64; n_genomes];
let mut n_kmers = 0usize;
let index_dir = self.partition.index_dir(i);
let index_dir = self.index_dir(i);
if !index_dir.exists() { return (0, counts); }
let n_layers = self.partition.n_layers(i).unwrap_or(0);
let n_layers = self.n_layers(i).unwrap_or(0);
for l in 0..n_layers {
let this_layer_dir = self.partition.layer_dir(i, l);
let this_layer_dir = self.layer_dir(i, l);
if !this_layer_dir.exists() { continue; }
n_kmers += LayerMeta::load(&this_layer_dir).map(|m| m.n).unwrap_or(0);
@@ -1,4 +1,5 @@
use super::*;
use crate::meta::IndexConfig;
// ── QueryStats::AddAssign ───────────────────────────────────────────────────
@@ -44,8 +45,15 @@ fn query_stats_default_is_zero() {
#[test]
fn query_partition_with_missing_index_dir_returns_default_stats() {
let tmp = tempfile::tempdir().expect("tempdir");
let partition =
KmerPartitions::create(tmp.path().join("idx"), 2, 21, 9, false).expect("create partition");
let config = IndexConfig {
kmer_size: 21,
minimizer_size: 9,
n_bits: 2,
with_counts: false,
evidence: obilayeredmap::IndexMode::Exact,
block_bits: 0,
};
let index = KmerIndex::create(tmp.path().join("idx"), config, None, false).expect("create index");
let mut kmers: HashMap<CanonicalKmer, Vec<KmerDesc>> = HashMap::new();
// Any well-formed canonical k-mer works here — the call must return
@@ -53,7 +61,7 @@ fn query_partition_with_missing_index_dir_returns_default_stats() {
let kmer = CanonicalKmer::from_raw_unchecked(0u64);
kmers.insert(kmer, vec![KmerDesc { seq_idx: 0, pos: 0 }]);
let stats = partition
let stats = index
.query_partition_with(0, &kmers, 1, false, |_event| {
panic!("on_event must not be called: no index was built");
})
@@ -65,11 +73,18 @@ fn query_partition_with_missing_index_dir_returns_default_stats() {
#[test]
fn query_partition_with_empty_kmers_is_a_noop() {
let tmp = tempfile::tempdir().expect("tempdir");
let partition =
KmerPartitions::create(tmp.path().join("idx"), 2, 21, 9, false).expect("create partition");
let config = IndexConfig {
kmer_size: 21,
minimizer_size: 9,
n_bits: 2,
with_counts: false,
evidence: obilayeredmap::IndexMode::Exact,
block_bits: 0,
};
let index = KmerIndex::create(tmp.path().join("idx"), config, None, false).expect("create index");
let kmers: HashMap<CanonicalKmer, Vec<KmerDesc>> = HashMap::new();
let stats = partition
let stats = index
.query_partition_with(0, &kmers, 1, false, |_event| {
panic!("on_event must not be called on an empty kmer map");
})
+1 -1
View File
@@ -2,7 +2,7 @@ use std::path::PathBuf;
use clap::Args;
use obikindex::{KmerIndex, MergeMode};
use obikpartitionner::filter::{MaxTotalCount, MinComplexity, MinTotalCount};
use obikindex::filter::{MaxTotalCount, MinComplexity, MinTotalCount};
use obisys::Reporter;
use tracing::info;
+5 -1
View File
@@ -240,7 +240,11 @@ pub fn run(args: IndexArgs) {
let n_workers = args.common.threads.max(1);
let max_open = args.common.effective_max_open();
scatter(idx.partition_mut(), args.common.seqfile_paths(), k, level_max, theta, n_workers, max_open, &mut rep);
let mut router = idx.partition_router().unwrap_or_else(|e| {
eprintln!("error opening partition router: {e}");
std::process::exit(1);
});
scatter(&mut router, args.common.seqfile_paths(), k, level_max, theta, n_workers, max_open, &mut rep);
idx.mark_scattered().unwrap_or_else(|e| {
eprintln!("error marking scatter done: {e}");
+1 -1
View File
@@ -1,6 +1,6 @@
use clap::Args;
use obikindex::{GroupFilterParams, IndexMeta, MetaPred};
use obikpartitionner::KmerFilter;
use obikindex::KmerFilter;
/// CLI args for ingroup/outgroup filtering — embeddable in any command via `#[command(flatten)]`.
#[derive(Args)]
+1 -1
View File
@@ -1,6 +1,6 @@
use std::collections::HashMap;
use obikpartitionner::KmerDesc;
use obikindex::KmerDesc;
use obikseq::CanonicalKmer;
use obiread::record::SeqRecord;
use obiskbuilder::SuperKmerIter;
+2 -2
View File
@@ -1,7 +1,7 @@
use std::time::Instant;
use obikindex::KmerIndex;
use obikpartitionner::{KmerDesc, QueryHit, QueryStats};
use obikindex::{KmerDesc, QueryHit, QueryStats};
use obikrope::Rope;
use obikseq::CanonicalKmer;
use obiread::record::parse_chunk;
@@ -110,7 +110,7 @@ pub(super) fn process_chunk(
continue;
}
let stats = idx.partition()
let stats = idx
.query_partition_with(
part_idx,
kmers,
+1 -1
View File
@@ -3,7 +3,7 @@ use std::path::PathBuf;
use clap::{Args, ValueEnum};
use obikindex::{IndexMeta, KmerIndex};
use obikpartitionner::{AggOp, OutputCol};
use obikindex::{AggOp, OutputCol};
use obisys::Reporter;
use tracing::info;
+1 -2
View File
@@ -36,7 +36,6 @@ pub fn run(args: UnitigArgs) {
info!("unitig: building de Bruijn graph from {n} partition(s) (k={k})");
let filters = args.filter.build_filters(idx.meta());
let partition = idx.partition();
let mut rep = Reporter::new();
// ── Phase 1 : collect filtered kmers in parallel ──────────────────────────
@@ -45,7 +44,7 @@ pub fn run(args: UnitigArgs) {
let g = (0..n)
.into_par_iter()
.fold(GraphDeBruijn::new, |mut local_g, i| {
partition
idx
.iter_partition_kmers(i, use_counts, n_genomes, &filters, |kmer, _row| {
local_g.push(kmer);
true
+2 -2
View File
@@ -3,7 +3,7 @@ use std::sync::Arc;
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
use std::time::Instant;
use obikpartitionner::KmerPartitions;
use obikpartitionner::PartitionRouter;
use obipipeline::{ThrottleGuard, Throttled, throttle};
use obiread::NucPage;
use obisys::spinner;
@@ -38,7 +38,7 @@ impl Drop for GuardedIter {
/// Run scatter: normalise → build superkmers → route to partition → close.
/// Reports the "scatter" stage to `rep`.
pub fn scatter(
kp: &mut KmerPartitions,
kp: &mut PartitionRouter,
path_source: impl Iterator<Item = PathBuf> + Send + 'static,
k: usize,
level_max: usize,
-5
View File
@@ -13,11 +13,8 @@ obikrope = { path = "../obikrope" }
niffler = "3.0.0"
remove_dir_all = "1.0"
obikseq = { path = "../obikseq" }
obikentropy = { path = "../obikentropy" }
obiskbuilder = { path = "../obiskbuilder" }
obiskio = { path = "../obiskio" }
obidebruinj = { path = "../obidebruinj" }
obilayeredmap = { path = "../obilayeredmap" }
rayon = "1"
sysinfo = "0.39"
serde = { version = "1", features = ["derive"] }
@@ -28,6 +25,4 @@ epserde = "0.8"
memmap2 = "0.9.10"
obicompactvec = { path = "../obicompactvec" }
ptr_hash = "1.1"
indicatif = "0.18"
obisys = { path = "../obisys" }
obipipeline = { path = "../obipipeline" }
+1 -15
View File
@@ -1,18 +1,4 @@
mod common;
mod distance;
mod dump_layer;
pub mod filter;
mod graph_pipeline;
mod index_layer;
mod kmer_sort;
mod merge_layer;
mod partition;
mod query_layer;
mod rebuild_layer;
mod select_layer;
pub use filter::{GroupQuorumFilter, KmerFilter, passes_all};
pub use merge_layer::MergeMode;
pub use partition::{KmerPartitions, KmerSpectrum, PARTITIONS_SUBDIR};
pub use query_layer::{KmerDesc, QueryHit, QueryStats};
pub use select_layer::{AggOp, OutputCol};
pub use partition::{partition_dir, KmerSpectrum, PartitionRouter, PARTITIONS_SUBDIR};
+6 -6
View File
@@ -1,19 +1,19 @@
//! K-mer partitioning: routing super-kmers into per-partition files,
//! deduplicating them, and counting unique canonical k-mers.
//!
//! Submodules: [`kmer_partition`] (`KmerPartition`, `KmerSpectrum` and the
//! routing/lifecycle API), [`dereplicate`] (two-phase split+merge
//! deduplication), [`count`] (unique-kmer enumeration, MPHF, abundance
//! counting).
//! Submodules: [`router`] (`PartitionRouter`, `KmerSpectrum`, the
//! `partition_dir` naming convention, and the routing/lifecycle API),
//! [`dereplicate`] (two-phase split+merge deduplication), [`count`]
//! (unique-kmer enumeration, MPHF, abundance counting).
mod count;
mod dereplicate;
mod kmer_partition;
mod router;
#[cfg(test)]
mod tests;
pub use kmer_partition::{KmerPartitions, KmerSpectrum};
pub use router::{partition_dir, KmerSpectrum, PartitionRouter};
const SK_EXT: &str = "skmer.zst";
pub const PARTITIONS_SUBDIR: &str = "partitions";
@@ -7,8 +7,6 @@ use std::time::Instant;
use obisys::progress_bar;
use obikseq::RoutableSuperKmer;
use obilayeredmap::IndexMode;
use obilayeredmap::meta::PartitionMeta;
use obiskio::SKResult;
use rayon::prelude::*;
use remove_dir_all::remove_dir_all;
@@ -18,19 +16,22 @@ use niffler::Level;
use niffler::send::compression::Format;
use obiskio::SKFileWriter;
use crate::common::load_meta;
use crate::kmer_sort::chunk_size_from_ram;
use super::count::count_partition;
use super::dereplicate::{dereplicate_partition, optimal_buckets};
use super::{PARTITIONS_SUBDIR, SK_EXT};
/// Name of a partition's layered-index subdirectory — the single source of
/// truth `index_dir`/`layer_dir` build on, replacing what used to be a
/// `const INDEX_SUBDIR: &str = "index";` (or a bare `"index"` literal)
/// redefined independently in every module that needed a partition's index
/// path.
const INDEX_SUBDIR: &str = "index";
/// Path of partition `i`'s directory under `root` — the single source of
/// truth for the `part_{i:05}` on-disk naming convention. Shared by
/// `PartitionRouter` (which writes here) and, one crate up, `KmerIndex`
/// (which builds `index_dir`/`layer_dir` on top of this same root, once
/// `PartitionRouter` has finished writing) — the only piece of
/// partition-directory knowledge that genuinely needs to cross the
/// crate boundary, since both sides must agree on where a partition lives.
pub fn partition_dir(root: &Path, i: usize) -> PathBuf {
root.join(PARTITIONS_SUBDIR).join(format!("part_{i:05}"))
}
pub struct KmerSpectrum {
pub f0: u64,
@@ -38,78 +39,51 @@ pub struct KmerSpectrum {
pub counts: BTreeMap<u32, u64>,
}
pub struct KmerPartitions {
/// Routes raw super-kmers into per-partition files, then dereplicates and
/// counts them — this crate's entire job now that layer/query/merge/select/
/// rebuild/dump/distance concerns have moved to `obikindex` (they operate on
/// built layers, which don't exist yet at this stage — see
/// `DevDocMD/implementation/partition_layer_cache.md`). Transient: it makes
/// sense only while raw partition files are being written or processed, a
/// phase that always precedes any `Layer`.
pub struct PartitionRouter {
root_path: PathBuf,
n_partitions: usize,
partitions_mask: u64,
kmer_size: usize,
minimizer_size: usize,
writers: Vec<Option<SKFileWriter>>,
level: Level,
closed: bool,
}
impl KmerPartitions {
pub fn create<P: AsRef<Path>>(
path: P,
n_bits: usize,
kmer_size: usize,
minimizer_size: usize,
force: bool,
) -> SKResult<Self> {
Self::create_with(path, n_bits, kmer_size, minimizer_size, Level::One, force)
}
pub fn create_with<P: AsRef<Path>>(
path: P,
n_bits: usize,
kmer_size: usize,
minimizer_size: usize,
level: Level,
force: bool,
) -> SKResult<Self> {
let root_path = path.as_ref().to_owned();
impl PartitionRouter {
/// Create a fresh partition layout at `root_path` for `n_partitions =
/// 2^n_bits` partitions.
pub fn create(root_path: &Path, n_bits: usize, force: bool) -> SKResult<Self> {
// `root_path` itself may already exist as a bare directory: callers
// typically hold an index-level lock file there before creating the
// partition layout. What actually signals a pre-existing partition
// set is the `PARTITIONS_SUBDIR` subdirectory, not the root itself.
if root_path.join(PARTITIONS_SUBDIR).exists() {
if force {
remove_dir_all(&root_path)?;
remove_dir_all(root_path)?;
} else {
return Err(io::Error::new(
io::ErrorKind::AlreadyExists,
format!(
"{}: partition directory already exists",
root_path.display()
),
format!("{}: partition directory already exists", root_path.display()),
)
.into());
}
}
fs::create_dir_all(root_path.join(PARTITIONS_SUBDIR))?;
let n_partitions = 1usize << n_bits;
let writers = (0..n_partitions).map(|_| None).collect();
let partition = Self {
root_path,
n_partitions,
partitions_mask: (1u64 << n_bits) - 1,
kmer_size,
minimizer_size,
writers,
level,
closed: false,
};
Ok(partition)
Self::new(root_path, n_bits)
}
pub fn open_with_config<P: AsRef<Path>>(
path: P,
kmer_size: usize,
minimizer_size: usize,
n_bits: usize,
) -> SKResult<Self> {
let root_path = path.as_ref().to_owned();
/// Reopen the partition layout at `root_path` for further routing (or
/// for `dereplicate`/`count_kmer`, which don't need `writers` but reuse
/// this same handle for consistency). Every caller of `open` wants to
/// write or process, never just read paths (those live on `KmerIndex`
/// directly), so this always starts open too.
pub fn open(root_path: &Path, n_bits: usize) -> SKResult<Self> {
if !root_path.exists() {
return Err(io::Error::new(
io::ErrorKind::NotFound,
@@ -117,17 +91,19 @@ impl KmerPartitions {
)
.into());
}
Self::new(root_path, n_bits)
}
fn new(root_path: &Path, n_bits: usize) -> SKResult<Self> {
let n_partitions = 1usize << n_bits;
let writers = (0..n_partitions).map(|_| None).collect();
Ok(Self {
root_path,
root_path: root_path.to_owned(),
n_partitions,
partitions_mask: (1u64 << n_bits) - 1,
kmer_size,
minimizer_size,
writers,
level: Level::One,
closed: true,
closed: false,
})
}
@@ -173,61 +149,6 @@ impl KmerPartitions {
!self.closed
}
pub fn path(&self) -> &Path {
&self.root_path
}
/// Path of partition `i` directory.
pub fn partition_dir(&self, i: usize) -> PathBuf {
self.root_path
.join(PARTITIONS_SUBDIR)
.join(format!("part_{i:05}"))
}
/// Path of partition `i`'s layered-index directory (`<partition>/index`).
pub fn index_dir(&self, i: usize) -> PathBuf {
self.partition_dir(i).join(INDEX_SUBDIR)
}
/// Path of layer `l` within partition `i`'s layered index — composes
/// [`index_dir`](Self::index_dir) with `obilayeredmap`'s own
/// `layer_N` naming convention rather than reimplementing it.
pub fn layer_dir(&self, i: usize, l: usize) -> PathBuf {
obilayeredmap::layer_dir(&self.index_dir(i), l)
}
/// Partition `i`'s metadata (layer count, evidence mode) — the single
/// entry point for this, so that callers outside `obikpartitionner`
/// never need to know it's a `meta.json` loaded via
/// `obilayeredmap::meta::PartitionMeta`, nor handle its own recovery
/// path for indexes built before that file existed (see
/// [`crate::common::load_meta`]).
pub fn partition_meta(&self, i: usize) -> SKResult<PartitionMeta> {
load_meta(&self.index_dir(i), "partition_meta")
}
/// Number of layers in partition `i` — see [`partition_meta`](Self::partition_meta).
pub fn n_layers(&self, i: usize) -> SKResult<usize> {
self.partition_meta(i).map(|m| m.n_layers)
}
/// Evidence mode of partition `i` — see [`partition_meta`](Self::partition_meta).
pub fn index_mode(&self, i: usize) -> SKResult<IndexMode> {
self.partition_meta(i).map(|m| m.mode)
}
pub fn kmer_size(&self) -> usize {
self.kmer_size
}
pub fn minimizer_size(&self) -> usize {
self.minimizer_size
}
pub fn n_partitions(&self) -> usize {
self.n_partitions
}
/// Deduplicate all `raw.{ext}` files in parallel, replacing each with a
/// `dereplicated.{ext}` file where identical canonical sequences are merged
/// and their counts summed.
@@ -243,10 +164,6 @@ impl KmerPartitions {
///
/// If a merged count exceeds the 24-bit header limit, the sequence is
/// emitted as multiple records whose counts sum to the true total.
///
/// `temp_bits` controls the split fan-out (`2^temp_bits` temp files per
/// partition). Higher values reduce per-temp-file memory at the cost of
/// more temporary file descriptors — all managed by the global fd pool.
pub fn dereplicate(&self) -> SKResult<()> {
let level = self.level;
let sys = System::new_all();
@@ -264,7 +181,7 @@ impl KmerPartitions {
let results: Vec<SKResult<()>> = (0..self.n_partitions)
.into_par_iter()
.map(|i| {
let dir = self.partition_dir(i);
let dir = partition_dir(&self.root_path, i);
if !dir.exists() {
pb.inc(1);
return Ok(());
@@ -295,9 +212,6 @@ impl KmerPartitions {
///
/// Returns the aggregated `KmerSpectrum`. Per-partition spectrum files are
/// deleted after aggregation unless `keep_partial` is true.
///
/// Partitions are processed in parallel via Rayon (one task per thread).
/// Peak memory per partition is ~80 MB, so n_threads partitions run simultaneously.
pub fn count_kmer(&self, keep_partial: bool) -> SKResult<KmerSpectrum> {
let sys = System::new_all();
let available = match sys.available_memory() {
@@ -312,7 +226,7 @@ impl KmerPartitions {
let results: Vec<SKResult<()>> = (0..self.n_partitions)
.into_par_iter()
.map(|i| {
let dir = self.partition_dir(i);
let dir = partition_dir(&self.root_path, i);
let dedup_path = dir.join(format!("dereplicated.{SK_EXT}"));
if !dedup_path.exists() {
pb.inc(1);
@@ -337,7 +251,7 @@ impl KmerPartitions {
let mut f1: u64 = 0;
for i in 0..self.n_partitions {
let path = self.partition_dir(i).join("kmer_spectrum_raw.json");
let path = partition_dir(&self.root_path, i).join("kmer_spectrum_raw.json");
if !path.exists() {
continue;
}
@@ -364,7 +278,7 @@ impl KmerPartitions {
fn check_not_closed(&self) -> SKResult<()> {
if self.closed {
Err(io::Error::new(io::ErrorKind::BrokenPipe, "write to closed KmerPartition").into())
Err(io::Error::new(io::ErrorKind::BrokenPipe, "write to closed PartitionRouter").into())
} else {
Ok(())
}
@@ -372,7 +286,7 @@ impl KmerPartitions {
fn ensure_writer(&mut self, partition: usize) -> SKResult<&mut SKFileWriter> {
if self.writers[partition].is_none() {
let dir = self.partition_dir(partition);
let dir = partition_dir(&self.root_path, partition);
fs::create_dir_all(&dir)?;
let file_path = dir.join(format!("raw.{SK_EXT}"));
let writer = SKFileWriter::create_with(file_path, Format::Zstd, self.level)?;
@@ -382,7 +296,7 @@ impl KmerPartitions {
}
}
impl Drop for KmerPartitions {
impl Drop for PartitionRouter {
fn drop(&mut self) {
let _ = self.close();
}
+2 -2
View File
@@ -6,7 +6,7 @@ use obikseq::SuperKmer;
use obiskbuilder::build_superkmers;
use super::count::count_partition;
use super::{KmerPartitions, PARTITIONS_SUBDIR};
use super::{PartitionRouter, PARTITIONS_SUBDIR};
const K: usize = 11;
const M: usize = 5;
@@ -46,7 +46,7 @@ fn pipeline_counts(seqs: &[&[u8]]) -> (u64, u64) {
let superkmers: Vec<_> = build_superkmers(rope, K, 1, 0.0);
let dir = tempfile::tempdir().unwrap();
let mut kp = KmerPartitions::create(dir.path(), 0, K, M, true).unwrap();
let mut kp = PartitionRouter::create(dir.path(), 0, true).unwrap();
kp.write_batch(superkmers).unwrap();
kp.close().unwrap();
kp.dereplicate().unwrap();
+2 -12
View File
@@ -1,10 +1,9 @@
use std::sync::Arc;
use obikpartitionner::KmerPartitions;
use obisys::progress_bar;
use obikindex::KmerIndex;
use obikindex::{OKIError, OKIResult};
use obikindex::OKIResult;
use super::cache::PartitionCache;
use super::family_scan::{Selection, scan_layer_families};
@@ -87,16 +86,7 @@ impl SnpAlignmentExt for KmerIndex {
let n_genomes = self.meta().genomes.len();
let with_counts = self.meta().config.with_counts;
let k = self.kmer_size();
let n_bits = n_parts.trailing_zeros() as usize;
let partition = KmerPartitions::open_with_config(
self.root_path(),
self.kmer_size(),
self.minimizer_size(),
n_bits,
)
.map_err(OKIError::Partition)?;
let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?);
let cache = Arc::new(PartitionCache::build(self, n_parts, with_counts)?);
let layer_dirs = super::family_scan::sibling_layer_dirs(self)?;
let selections =
super::subsample::compute_selections(self, &layer_dirs, subsample, entropy_bias)?;
+5 -15
View File
@@ -4,7 +4,6 @@ use std::sync::atomic::Ordering;
use rayon::prelude::*;
use obikpartitionner::KmerPartitions;
use obikseq::CanonicalKmer;
use obilayeredmap::MphfLayer;
use obilayeredmap::meta::IndexMode;
@@ -12,7 +11,7 @@ use obipipeline::ThrottleGuard;
use obisys::progress_bar;
use obikindex::KmerIndex;
use obikindex::{OKIError, OKIResult};
use obikindex::OKIResult;
use super::cache::PartitionCache;
use super::helpers::{central_base, is_minorant};
@@ -76,19 +75,10 @@ pub trait SiblingAnnexBuildExt {
impl SiblingAnnexBuildExt for KmerIndex {
fn build_sibling_annex(&self) -> OKIResult<()> {
let n_parts = self.n_partitions();
let n_bits = n_parts.trailing_zeros() as usize;
let partition = KmerPartitions::open_with_config(
self.root_path(),
self.kmer_size(),
self.minimizer_size(),
n_bits,
)
.map_err(OKIError::Partition)?;
tracing::info!("opening {n_parts} partition(s) for the sibling-annex sweep");
let cache = Arc::new(PartitionCache::build(
&partition,
self,
n_parts,
self.meta().config.with_counts,
)?);
@@ -96,18 +86,18 @@ impl SiblingAnnexBuildExt for KmerIndex {
let pb = progress_bar("sibling_annex", n_parts as u64, "partitions");
let mut total_slots: u64 = 0;
for part in 0..n_parts {
let index_dir = self.partition().index_dir(part);
let index_dir = self.index_dir(part);
if !index_dir.exists() {
pb.inc(1);
continue;
}
let meta = self.partition().partition_meta(part)?;
let meta = self.partition_meta(part)?;
let mut part_slots: u64 = 0;
for l in 0..meta.n_layers {
part_slots += build_layer_sibling_annex(
self,
&self.partition().layer_dir(part, l),
&self.layer_dir(part, l),
&meta.mode,
n_parts,
l,
+5 -6
View File
@@ -3,13 +3,12 @@ use rayon::prelude::*;
use std::path::Path;
use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix};
use obikpartitionner::KmerPartitions;
use obikseq::CanonicalKmer;
use obilayeredmap::meta::IndexMode;
use obilayeredmap::{Layer, OLMResult};
use obisys::progress_bar;
use obikindex::OKIResult;
use obikindex::{KmerIndex, OKIResult};
use super::SiblingAnnex;
use super::iter::SiblingLayerExt;
@@ -159,7 +158,7 @@ pub(super) struct PartitionCache {
impl PartitionCache {
pub(super) fn build(
partition: &KmerPartitions,
index: &KmerIndex,
n_parts: usize,
with_counts: bool,
) -> OKIResult<Self> {
@@ -167,15 +166,15 @@ impl PartitionCache {
let built: Vec<(Vec<Mat>, usize)> = (0..n_parts)
.into_par_iter()
.map(|part| -> OKIResult<(Vec<Mat>, usize)> {
let index_dir = partition.index_dir(part);
let index_dir = index.index_dir(part);
if !index_dir.exists() {
pb.inc(1);
return Ok((Vec::new(), 0));
}
let meta = partition.partition_meta(part)?;
let meta = index.partition_meta(part)?;
let mut mats = Vec::with_capacity(meta.n_layers);
for l in 0..meta.n_layers {
let Ok(mat) = Mat::open(&partition.layer_dir(part, l), &meta.mode, with_counts)
let Ok(mat) = Mat::open(&index.layer_dir(part, l), &meta.mode, with_counts)
else {
continue;
};
+2 -12
View File
@@ -2,11 +2,10 @@ use std::sync::Arc;
use ndarray::Array2;
use obikpartitionner::KmerPartitions;
use obisys::progress_bar;
use obikindex::KmerIndex;
use obikindex::{OKIError, OKIResult};
use obikindex::OKIResult;
use super::cache::PartitionCache;
use super::distance::RawSnpDistanceOutput;
@@ -69,8 +68,6 @@ impl CardinalityExt for KmerIndex {
let n_genomes = self.meta().genomes.len();
let with_counts = self.meta().config.with_counts;
let k = self.kmer_size();
let n_bits = n_parts.trailing_zeros() as usize;
let included = Array2::from_shape_fn((n_genomes, n_genomes), |(i, j)| {
if i == j {
return false;
@@ -80,14 +77,7 @@ impl CardinalityExt for KmerIndex {
total > 0 && (snp as f64 / total as f64) <= ratio_ceiling
});
let partition = KmerPartitions::open_with_config(
self.root_path(),
self.kmer_size(),
self.minimizer_size(),
n_bits,
)
.map_err(OKIError::Partition)?;
let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?);
let cache = Arc::new(PartitionCache::build(self, n_parts, with_counts)?);
let layer_dirs = super::family_scan::sibling_layer_dirs(self)?;
let pb = progress_bar("cardinality_tally", layer_dirs.len() as u64, "layers");
+2 -12
View File
@@ -2,11 +2,10 @@ use std::sync::Arc;
use ndarray::Array2;
use obikpartitionner::KmerPartitions;
use obisys::progress_bar;
use obikindex::KmerIndex;
use obikindex::{OKIError, OKIResult};
use obikindex::OKIResult;
use super::cache::PartitionCache;
use super::family_scan::{Selection, scan_layer_families};
@@ -70,16 +69,7 @@ where
let n_genomes = index.meta().genomes.len();
let with_counts = index.meta().config.with_counts;
let k = index.kmer_size();
let n_bits = n_parts.trailing_zeros() as usize;
let partition = KmerPartitions::open_with_config(
index.root_path(),
index.kmer_size(),
index.minimizer_size(),
n_bits,
)
.map_err(OKIError::Partition)?;
let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?);
let cache = Arc::new(PartitionCache::build(index, n_parts, with_counts)?);
let layer_dirs = super::family_scan::sibling_layer_dirs(index)?;
let pb = progress_bar(label, layer_dirs.len() as u64, "layers");
+2 -21
View File
@@ -13,7 +13,6 @@ use std::io::{BufWriter, Write};
use std::path::{Path, PathBuf};
use std::sync::Arc;
use obikpartitionner::KmerPartitions;
use obisys::progress_bar;
use obikindex::KmerIndex;
@@ -139,16 +138,7 @@ impl ShannonEntropyExt for KmerIndex {
let n_genomes = self.meta().genomes.len();
let with_counts = self.meta().config.with_counts;
let k = self.kmer_size();
let n_bits = n_parts.trailing_zeros() as usize;
let partition = KmerPartitions::open_with_config(
self.root_path(),
self.kmer_size(),
self.minimizer_size(),
n_bits,
)
.map_err(OKIError::Partition)?;
let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?);
let cache = Arc::new(PartitionCache::build(self, n_parts, with_counts)?);
let layer_dirs = super::family_scan::sibling_layer_dirs(self)?;
let selections =
super::subsample::compute_selections(self, &layer_dirs, subsample, entropy_bias)?;
@@ -240,16 +230,7 @@ pub(super) fn ensure_entropy_annexes(index: &KmerIndex, layer_dirs: &[PathBuf])
let n_genomes = index.meta().genomes.len();
let with_counts = index.meta().config.with_counts;
let k = index.kmer_size();
let n_bits = n_parts.trailing_zeros() as usize;
let partition = KmerPartitions::open_with_config(
index.root_path(),
index.kmer_size(),
index.minimizer_size(),
n_bits,
)
.map_err(OKIError::Partition)?;
let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?);
let cache = Arc::new(PartitionCache::build(index, n_parts, with_counts)?);
let minorant_counts = super::subsample::minorant_counts(layer_dirs)?;
let pb = progress_bar("entropy_annex_build", missing.len() as u64, "layers");
+3 -3
View File
@@ -105,13 +105,13 @@ pub(crate) fn sibling_layer_dirs(index: &KmerIndex) -> OKIResult<Vec<PathBuf>> {
let n_parts = index.n_partitions();
let mut layer_dirs = Vec::new();
for part in 0..n_parts {
let index_dir = index.partition().index_dir(part);
let index_dir = index.index_dir(part);
if !index_dir.exists() {
continue;
}
let n_layers = index.partition().n_layers(part)?;
let n_layers = index.n_layers(part)?;
for l in 0..n_layers {
let this_layer_dir = index.partition().layer_dir(part, l);
let this_layer_dir = index.layer_dir(part, l);
let annex_path = this_layer_dir.join(ANNEX_FILE_NAME);
if !annex_path.exists() {
return Err(OKIError::InvalidInput(format!(
+2 -12
View File
@@ -29,11 +29,10 @@ use std::sync::Arc;
use ndarray::Array2;
use obikpartitionner::KmerPartitions;
use obisys::progress_bar;
use obikindex::KmerIndex;
use obikindex::{OKIError, OKIResult};
use obikindex::OKIResult;
use super::alignment::{SnpAlignment, iupac_code};
use super::cache::PartitionCache;
@@ -83,16 +82,7 @@ impl SankoffBundleExt for KmerIndex {
let n_genomes = self.meta().genomes.len();
let with_counts = self.meta().config.with_counts;
let k = self.kmer_size();
let n_bits = n_parts.trailing_zeros() as usize;
let partition = KmerPartitions::open_with_config(
self.root_path(),
self.kmer_size(),
self.minimizer_size(),
n_bits,
)
.map_err(OKIError::Partition)?;
let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?);
let cache = Arc::new(PartitionCache::build(self, n_parts, with_counts)?);
let layer_dirs = sibling_layer_dirs(self)?;
// Computed once — every pass below iterates the exact same
// families, in the exact same layers, per the shared-selection
+2 -12
View File
@@ -2,11 +2,10 @@ use std::sync::Arc;
use rayon::prelude::*;
use obikpartitionner::KmerPartitions;
use obisys::progress_bar;
use obikindex::KmerIndex;
use obikindex::{OKIError, OKIResult};
use obikindex::OKIResult;
use super::ANNEX_FILE_NAME;
use super::SiblingAnnex;
@@ -109,19 +108,10 @@ impl SiblingStatsExt for KmerIndex {
let n_genomes = self.meta().genomes.len();
let with_counts = self.meta().config.with_counts;
let k = self.kmer_size();
let n_bits = n_parts.trailing_zeros() as usize;
// Same whole-run cache as `build_sibling_annex` — see its docs for
// why re-opening per lookup (or per call to a batching helper) is
// not good enough on a real index.
let partition = KmerPartitions::open_with_config(
self.root_path(),
self.kmer_size(),
self.minimizer_size(),
n_bits,
)
.map_err(OKIError::Partition)?;
let cache = Arc::new(PartitionCache::build(&partition, n_parts, with_counts)?);
let cache = Arc::new(PartitionCache::build(self, n_parts, with_counts)?);
let layer_dirs = super::family_scan::sibling_layer_dirs(self)?;
// One layer at a time, not parallelised across layers — see
+9 -8
View File
@@ -69,11 +69,12 @@ fn build_single_genome_index(dir: &Path, label: &str, seq: &[u8]) -> KmerIndex {
let mut rep = Reporter::new();
let stream = obiread::open_nuc_stream(fasta_path.to_str().unwrap(), K).expect("open fasta");
let mut router = idx.partition_router().expect("partition_router");
for page in stream {
let batch = obiskbuilder::build_superkmers_page(page, K, /* level_max */ 1, /* theta */ 0.0);
idx.partition_mut().write_batch(batch).expect("write_batch");
router.write_batch(batch).expect("write_batch");
}
idx.partition_mut().close().expect("close partition writers");
router.close().expect("close partition writers");
idx.mark_scattered().expect("mark_scattered");
idx.dereplicate_and_count(false, &mut rep).expect("dereplicate_and_count");
idx.build_layers(1, None, false, &mut rep).expect("build_layers");
@@ -87,9 +88,9 @@ fn canonical(ascii: &[u8]) -> CanonicalKmer {
/// Read back the annex entry for a given canonical k-mer from the merged
/// index's (single) partition/layer, asserting it was found at all.
fn annex_info_for(idx: &KmerIndex, kmer: CanonicalKmer) -> FamilyMask {
let meta = idx.partition().partition_meta(0).unwrap();
let meta = idx.partition_meta(0).unwrap();
for l in 0..meta.n_layers {
let layer_dir = idx.partition().layer_dir(0, l);
let layer_dir = idx.layer_dir(0, l);
let mphf = MphfLayer::open(&layer_dir, &meta.mode).unwrap();
if let Some(slot) = mphf.find(kmer) {
let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME)).unwrap();
@@ -165,7 +166,7 @@ fn sibling_annex_works_after_pack_sparse() {
let merged = merge_two(dir.path(), &g1, &g2);
merged.pack_matrices(true).expect("pack_matrices(sparse)");
let index_dir = merged.partition().index_dir(0);
let index_dir = merged.index_dir(0);
assert!(
index_dir.join("layer_0").join("presence").join("is_multi.prsb").exists(),
"pack_matrices(true) must leave the sparse marker file behind"
@@ -308,9 +309,9 @@ fn sibling_annex_no_empty_masks_after_build() {
let g1 = build_single_genome_index(dir.path(), "g1", &seq);
g1.build_sibling_annex().expect("build_sibling_annex");
let n_layers = g1.partition().n_layers(0).expect("partition meta");
let n_layers = g1.n_layers(0).expect("partition meta");
for l in 0..n_layers {
let layer_dir = g1.partition().layer_dir(0, l);
let layer_dir = g1.layer_dir(0, l);
let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME)).expect("annex open");
for slot in 0..annex.len() {
let mask = annex.get(slot).expect("slot must have an entry");
@@ -358,7 +359,7 @@ fn sibling_annex_records_the_real_layer_of_each_family_member() {
let merged = merge_two(dir.path(), &g1, &g2);
merged.build_sibling_annex().expect("build_sibling_annex");
let n_layers = merged.partition().n_layers(0).unwrap();
let n_layers = merged.n_layers(0).unwrap();
assert_eq!(n_layers, 2, "fixture assumption: one merge, one new layer");
let g1_kmer = canonical(b"AACCGCTTAAG"); // own base C (1), sibling base G (2) — lives in layer 1