Push zunrplorkwkt #70

Merged
coissac merged 93 commits from push-zunrplorkwkt into main 2026-08-28 23:15:38 +00:00
49 changed files with 357 additions and 427 deletions
Showing only changes of commit 6c860f120f - Show all commits
@@ -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