From 6c860f120fdadf9b15cb07c0554842e2ca8cc8ef Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Thu, 20 Aug 2026 13:49:06 +0200 Subject: [PATCH] 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. --- .../implementation/partition_layer_cache.md | 73 ++++++++ src/Cargo.lock | 13 +- src/obikindex/Cargo.toml | 8 + src/obikindex/examples/compare_sparse.rs | 8 +- .../src/common.rs | 0 src/obikindex/src/distance.rs | 4 +- src/obikindex/src/dump.rs | 10 +- .../src/dump_layer.rs | 8 +- .../src/filter.rs | 0 .../src/graph_pipeline.rs | 0 src/obikindex/src/index.rs | 141 +++++++------- .../src/index_layer.rs | 4 +- src/obikindex/src/lib.rs | 15 +- .../src/matrix_store.rs} | 4 +- src/obikindex/src/merge.rs | 11 +- .../src/merge_layer/mod.rs | 6 +- .../src/merge_layer/src_layer.rs | 0 src/obikindex/src/predicate.rs | 2 +- .../src/query_layer.rs | 4 +- src/obikindex/src/rebuild.rs | 7 +- .../src/rebuild_layer.rs | 6 +- src/obikindex/src/reindex.rs | 13 +- src/obikindex/src/select.rs | 21 +-- .../src/select_layer.rs | 6 +- src/obikindex/src/stats.rs | 12 +- .../src/tests/query_layer.rs | 27 ++- src/obikmer/src/cmd/filter/mod.rs | 2 +- src/obikmer/src/cmd/index/mod.rs | 6 +- src/obikmer/src/cmd/predicate.rs | 2 +- src/obikmer/src/cmd/query/batch.rs | 2 +- src/obikmer/src/cmd/query/chunk.rs | 4 +- src/obikmer/src/cmd/select/mod.rs | 2 +- src/obikmer/src/cmd/unitig/mod.rs | 3 +- src/obikmer/src/steps/scatter.rs | 4 +- src/obikpartitionner/Cargo.toml | 5 - src/obikpartitionner/src/lib.rs | 16 +- src/obikpartitionner/src/partition/mod.rs | 12 +- .../{kmer_partition.rs => router.rs} | 172 +++++------------- src/obikpartitionner/src/partition/tests.rs | 4 +- src/obikphylo/src/siblings/alignment.rs | 14 +- src/obikphylo/src/siblings/build.rs | 20 +- src/obikphylo/src/siblings/cache.rs | 11 +- src/obikphylo/src/siblings/cardinality.rs | 14 +- src/obikphylo/src/siblings/distance.rs | 14 +- src/obikphylo/src/siblings/entropy.rs | 23 +-- src/obikphylo/src/siblings/family_scan.rs | 6 +- src/obikphylo/src/siblings/sankoff_bundle.rs | 14 +- src/obikphylo/src/siblings/stats.rs | 14 +- src/obikphylo/src/siblings/tests.rs | 17 +- 49 files changed, 357 insertions(+), 427 deletions(-) rename src/{obikpartitionner => obikindex}/src/common.rs (100%) rename src/{obikpartitionner => obikindex}/src/dump_layer.rs (97%) rename src/{obikpartitionner => obikindex}/src/filter.rs (100%) rename src/{obikpartitionner => obikindex}/src/graph_pipeline.rs (100%) rename src/{obikpartitionner => obikindex}/src/index_layer.rs (98%) rename src/{obikpartitionner/src/distance.rs => obikindex/src/matrix_store.rs} (96%) rename src/{obikpartitionner => obikindex}/src/merge_layer/mod.rs (99%) rename src/{obikpartitionner => obikindex}/src/merge_layer/src_layer.rs (100%) rename src/{obikpartitionner => obikindex}/src/query_layer.rs (99%) rename src/{obikpartitionner => obikindex}/src/rebuild_layer.rs (99%) rename src/{obikpartitionner => obikindex}/src/select_layer.rs (99%) rename src/{obikpartitionner => obikindex}/src/tests/query_layer.rs (79%) rename src/obikpartitionner/src/partition/{kmer_partition.rs => router.rs} (66%) diff --git a/DevDocMD/implementation/partition_layer_cache.md b/DevDocMD/implementation/partition_layer_cache.md index 528cfd3b..7b3e087c 100644 --- a/DevDocMD/implementation/partition_layer_cache.md +++ b/DevDocMD/implementation/partition_layer_cache.md @@ -69,6 +69,79 @@ let a per-partition type hold layers of mixed `D` (today `LayeredMap` 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` 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 diff --git a/src/Cargo.lock b/src/Cargo.lock index b90d56a3..e7ba19b5 100644 --- a/src/Cargo.lock +++ b/src/Cargo.lock @@ -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", diff --git a/src/obikindex/Cargo.toml b/src/obikindex/Cargo.toml index e0707278..4edeb066 100644 --- a/src/obikindex/Cargo.toml +++ b/src/obikindex/Cargo.toml @@ -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" diff --git a/src/obikindex/examples/compare_sparse.rs b/src/obikindex/examples/compare_sparse.rs index 51376fad..a772b82b 100644 --- a/src/obikindex/examples/compare_sparse.rs +++ b/src/obikindex/examples/compare_sparse.rs @@ -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; diff --git a/src/obikpartitionner/src/common.rs b/src/obikindex/src/common.rs similarity index 100% rename from src/obikpartitionner/src/common.rs rename to src/obikindex/src/common.rs diff --git a/src/obikindex/src/distance.rs b/src/obikindex/src/distance.rs index f5d63be5..2b651c1b 100644 --- a/src/obikindex/src/distance.rs +++ b/src/obikindex/src/distance.rs @@ -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::>()?; 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::>()?; let global = LayeredStore::new(stores); diff --git a/src/obikindex/src/dump.rs b/src/obikindex/src/dump.rs index 756243fb..5c9499dc 100644 --- a/src/obikindex/src/dump.rs +++ b/src/obikindex/src/dump.rs @@ -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::::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); diff --git a/src/obikpartitionner/src/dump_layer.rs b/src/obikindex/src/dump_layer.rs similarity index 97% rename from src/obikpartitionner/src/dump_layer.rs rename to src/obikindex/src/dump_layer.rs index 119c71be..d14b07f4 100644 --- a/src/obikpartitionner/src/dump_layer.rs +++ b/src/obikindex/src/dump_layer.rs @@ -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 { diff --git a/src/obikpartitionner/src/filter.rs b/src/obikindex/src/filter.rs similarity index 100% rename from src/obikpartitionner/src/filter.rs rename to src/obikindex/src/filter.rs diff --git a/src/obikpartitionner/src/graph_pipeline.rs b/src/obikindex/src/graph_pipeline.rs similarity index 100% rename from src/obikpartitionner/src/graph_pipeline.rs rename to src/obikindex/src/graph_pipeline.rs diff --git a/src/obikindex/src/index.rs b/src/obikindex/src/index.rs index b8e973ad..160eaa25 100644 --- a/src/obikindex/src/index.rs +++ b/src/obikindex/src/index.rs @@ -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 { 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>(path: P) -> OKIResult { @@ -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>( - output: P, - meta: &IndexMeta, - ) -> OKIResult { + pub(crate) fn create_skeleton>(output: P, meta: &IndexMeta) -> OKIResult { 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 (`/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`). + pub fn partition_meta(&self, i: usize) -> obiskio::SKResult { + 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 { + 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 { + 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 { - 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 { + 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; diff --git a/src/obikpartitionner/src/index_layer.rs b/src/obikindex/src/index_layer.rs similarity index 98% rename from src/obikpartitionner/src/index_layer.rs rename to src/obikindex/src/index_layer.rs index 8c6a77cc..46b8d842 100644 --- a/src/obikpartitionner/src/index_layer.rs +++ b/src/obikindex/src/index_layer.rs @@ -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>, Xx64, Vec>; @@ -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 diff --git a/src/obikindex/src/lib.rs b/src/obikindex/src/lib.rs index 85f4a0ed..f8f33005 100644 --- a/src/obikindex/src/lib.rs +++ b/src/obikindex/src/lib.rs @@ -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; diff --git a/src/obikpartitionner/src/distance.rs b/src/obikindex/src/matrix_store.rs similarity index 96% rename from src/obikpartitionner/src/distance.rs rename to src/obikindex/src/matrix_store.rs index 291f3b05..33ab9f83 100644 --- a/src/obikpartitionner/src/distance.rs +++ b/src/obikindex/src/matrix_store.rs @@ -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> { diff --git a/src/obikindex/src/merge.rs b/src/obikindex/src/merge.rs index 444dc487..b00e6f71 100644 --- a/src/obikindex/src/merge.rs +++ b/src/obikindex/src/merge.rs @@ -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() } diff --git a/src/obikpartitionner/src/merge_layer/mod.rs b/src/obikindex/src/merge_layer/mod.rs similarity index 99% rename from src/obikpartitionner/src/merge_layer/mod.rs rename to src/obikindex/src/merge_layer/mod.rs index 77e68e29..9f96f73d 100644 --- a/src/obikpartitionner/src/merge_layer/mod.rs +++ b/src/obikindex/src/merge_layer/mod.rs @@ -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, diff --git a/src/obikpartitionner/src/merge_layer/src_layer.rs b/src/obikindex/src/merge_layer/src_layer.rs similarity index 100% rename from src/obikpartitionner/src/merge_layer/src_layer.rs rename to src/obikindex/src/merge_layer/src_layer.rs diff --git a/src/obikindex/src/predicate.rs b/src/obikindex/src/predicate.rs index b98286ef..dc808cd8 100644 --- a/src/obikindex/src/predicate.rs +++ b/src/obikindex/src/predicate.rs @@ -1,6 +1,6 @@ use std::collections::HashMap; -use obikpartitionner::GroupQuorumFilter; +use crate::GroupQuorumFilter; use obitaxonomy::{TaxPath, TaxPattern}; use crate::meta::{GenomeInfo, IndexMeta}; diff --git a/src/obikpartitionner/src/query_layer.rs b/src/obikindex/src/query_layer.rs similarity index 99% rename from src/obikpartitionner/src/query_layer.rs rename to src/obikindex/src/query_layer.rs index 8c4b1588..41ff61b1 100644 --- a/src/obikpartitionner/src/query_layer.rs +++ b/src/obikindex/src/query_layer.rs @@ -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. /// diff --git a/src/obikindex/src/rebuild.rs b/src/obikindex/src/rebuild.rs index 3db37f2f..9c7ffb63 100644 --- a/src/obikindex/src/rebuild.rs +++ b/src/obikindex/src/rebuild.rs @@ -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 = (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)?; diff --git a/src/obikpartitionner/src/rebuild_layer.rs b/src/obikindex/src/rebuild_layer.rs similarity index 99% rename from src/obikpartitionner/src/rebuild_layer.rs rename to src/obikindex/src/rebuild_layer.rs index 9873ce84..32d74bcb 100644 --- a/src/obikpartitionner/src/rebuild_layer.rs +++ b/src/obikindex/src/rebuild_layer.rs @@ -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], mode: MergeMode, diff --git a/src/obikindex/src/reindex.rs b/src/obikindex/src/reindex.rs index a58f6dee..6265e665 100644 --- a/src/obikindex/src/reindex.rs +++ b/src/obikindex/src/reindex.rs @@ -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(()) } diff --git a/src/obikindex/src/select.rs b/src/obikindex/src/select.rs index 1f6591f5..72f95fa3 100644 --- a/src/obikindex/src/select.rs +++ b/src/obikindex/src/select.rs @@ -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 = (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 = (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, diff --git a/src/obikpartitionner/src/select_layer.rs b/src/obikindex/src/select_layer.rs similarity index 99% rename from src/obikpartitionner/src/select_layer.rs rename to src/obikindex/src/select_layer.rs index bd66883e..cc54f937 100644 --- a/src/obikpartitionner/src/select_layer.rs +++ b/src/obikindex/src/select_layer.rs @@ -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, diff --git a/src/obikindex/src/stats.rs b/src/obikindex/src/stats.rs index 995b7c20..822db9e5 100644 --- a/src/obikindex/src/stats.rs +++ b/src/obikindex/src/stats.rs @@ -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); diff --git a/src/obikpartitionner/src/tests/query_layer.rs b/src/obikindex/src/tests/query_layer.rs similarity index 79% rename from src/obikpartitionner/src/tests/query_layer.rs rename to src/obikindex/src/tests/query_layer.rs index f1d2a9e7..037c3b86 100644 --- a/src/obikpartitionner/src/tests/query_layer.rs +++ b/src/obikindex/src/tests/query_layer.rs @@ -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> = 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> = 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"); }) diff --git a/src/obikmer/src/cmd/filter/mod.rs b/src/obikmer/src/cmd/filter/mod.rs index b2ca626a..e60cf8fb 100644 --- a/src/obikmer/src/cmd/filter/mod.rs +++ b/src/obikmer/src/cmd/filter/mod.rs @@ -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; diff --git a/src/obikmer/src/cmd/index/mod.rs b/src/obikmer/src/cmd/index/mod.rs index ca9fe857..c4d3dd6b 100644 --- a/src/obikmer/src/cmd/index/mod.rs +++ b/src/obikmer/src/cmd/index/mod.rs @@ -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}"); diff --git a/src/obikmer/src/cmd/predicate.rs b/src/obikmer/src/cmd/predicate.rs index 6d0f953e..87e0e653 100644 --- a/src/obikmer/src/cmd/predicate.rs +++ b/src/obikmer/src/cmd/predicate.rs @@ -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)] diff --git a/src/obikmer/src/cmd/query/batch.rs b/src/obikmer/src/cmd/query/batch.rs index 55cd9452..54d1c9db 100644 --- a/src/obikmer/src/cmd/query/batch.rs +++ b/src/obikmer/src/cmd/query/batch.rs @@ -1,6 +1,6 @@ use std::collections::HashMap; -use obikpartitionner::KmerDesc; +use obikindex::KmerDesc; use obikseq::CanonicalKmer; use obiread::record::SeqRecord; use obiskbuilder::SuperKmerIter; diff --git a/src/obikmer/src/cmd/query/chunk.rs b/src/obikmer/src/cmd/query/chunk.rs index 509d84b1..54debabf 100644 --- a/src/obikmer/src/cmd/query/chunk.rs +++ b/src/obikmer/src/cmd/query/chunk.rs @@ -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, diff --git a/src/obikmer/src/cmd/select/mod.rs b/src/obikmer/src/cmd/select/mod.rs index 5f4180ad..e50a4e60 100644 --- a/src/obikmer/src/cmd/select/mod.rs +++ b/src/obikmer/src/cmd/select/mod.rs @@ -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; diff --git a/src/obikmer/src/cmd/unitig/mod.rs b/src/obikmer/src/cmd/unitig/mod.rs index 62ad39b9..30af4a21 100644 --- a/src/obikmer/src/cmd/unitig/mod.rs +++ b/src/obikmer/src/cmd/unitig/mod.rs @@ -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 diff --git a/src/obikmer/src/steps/scatter.rs b/src/obikmer/src/steps/scatter.rs index a6b3b2f2..2929435f 100644 --- a/src/obikmer/src/steps/scatter.rs +++ b/src/obikmer/src/steps/scatter.rs @@ -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 + Send + 'static, k: usize, level_max: usize, diff --git a/src/obikpartitionner/Cargo.toml b/src/obikpartitionner/Cargo.toml index 310b746d..672436ab 100644 --- a/src/obikpartitionner/Cargo.toml +++ b/src/obikpartitionner/Cargo.toml @@ -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" } diff --git a/src/obikpartitionner/src/lib.rs b/src/obikpartitionner/src/lib.rs index 7c59f8ce..6f8793a0 100644 --- a/src/obikpartitionner/src/lib.rs +++ b/src/obikpartitionner/src/lib.rs @@ -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}; diff --git a/src/obikpartitionner/src/partition/mod.rs b/src/obikpartitionner/src/partition/mod.rs index 883c83a3..917868ed 100644 --- a/src/obikpartitionner/src/partition/mod.rs +++ b/src/obikpartitionner/src/partition/mod.rs @@ -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"; diff --git a/src/obikpartitionner/src/partition/kmer_partition.rs b/src/obikpartitionner/src/partition/router.rs similarity index 66% rename from src/obikpartitionner/src/partition/kmer_partition.rs rename to src/obikpartitionner/src/partition/router.rs index aad39b4f..d4f2beda 100644 --- a/src/obikpartitionner/src/partition/kmer_partition.rs +++ b/src/obikpartitionner/src/partition/router.rs @@ -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, } -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>, level: Level, closed: bool, } -impl KmerPartitions { - pub fn create>( - path: P, - n_bits: usize, - kmer_size: usize, - minimizer_size: usize, - force: bool, - ) -> SKResult { - Self::create_with(path, n_bits, kmer_size, minimizer_size, Level::One, force) - } - - pub fn create_with>( - path: P, - n_bits: usize, - kmer_size: usize, - minimizer_size: usize, - level: Level, - force: bool, - ) -> SKResult { - 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 { // `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>( - path: P, - kmer_size: usize, - minimizer_size: usize, - n_bits: usize, - ) -> SKResult { - 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 { 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 { 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 (`/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 { - 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 { - 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 { - 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> = (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 { let sys = System::new_all(); let available = match sys.available_memory() { @@ -312,7 +226,7 @@ impl KmerPartitions { let results: Vec> = (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(); } diff --git a/src/obikpartitionner/src/partition/tests.rs b/src/obikpartitionner/src/partition/tests.rs index 38d44728..6af36738 100644 --- a/src/obikpartitionner/src/partition/tests.rs +++ b/src/obikpartitionner/src/partition/tests.rs @@ -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(); diff --git a/src/obikphylo/src/siblings/alignment.rs b/src/obikphylo/src/siblings/alignment.rs index 754fd577..74d2d76f 100644 --- a/src/obikphylo/src/siblings/alignment.rs +++ b/src/obikphylo/src/siblings/alignment.rs @@ -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)?; diff --git a/src/obikphylo/src/siblings/build.rs b/src/obikphylo/src/siblings/build.rs index 6d125e19..e400ac8a 100644 --- a/src/obikphylo/src/siblings/build.rs +++ b/src/obikphylo/src/siblings/build.rs @@ -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, diff --git a/src/obikphylo/src/siblings/cache.rs b/src/obikphylo/src/siblings/cache.rs index 8bc8f4f2..06c29676 100644 --- a/src/obikphylo/src/siblings/cache.rs +++ b/src/obikphylo/src/siblings/cache.rs @@ -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 { @@ -167,15 +166,15 @@ impl PartitionCache { let built: Vec<(Vec, usize)> = (0..n_parts) .into_par_iter() .map(|part| -> OKIResult<(Vec, 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; }; diff --git a/src/obikphylo/src/siblings/cardinality.rs b/src/obikphylo/src/siblings/cardinality.rs index 211a1f63..4225bae7 100644 --- a/src/obikphylo/src/siblings/cardinality.rs +++ b/src/obikphylo/src/siblings/cardinality.rs @@ -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"); diff --git a/src/obikphylo/src/siblings/distance.rs b/src/obikphylo/src/siblings/distance.rs index 8b8b15e6..523dca4d 100644 --- a/src/obikphylo/src/siblings/distance.rs +++ b/src/obikphylo/src/siblings/distance.rs @@ -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"); diff --git a/src/obikphylo/src/siblings/entropy.rs b/src/obikphylo/src/siblings/entropy.rs index 16aef5e3..e2caa01a 100644 --- a/src/obikphylo/src/siblings/entropy.rs +++ b/src/obikphylo/src/siblings/entropy.rs @@ -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"); diff --git a/src/obikphylo/src/siblings/family_scan.rs b/src/obikphylo/src/siblings/family_scan.rs index f36c3f73..5a45b6c9 100644 --- a/src/obikphylo/src/siblings/family_scan.rs +++ b/src/obikphylo/src/siblings/family_scan.rs @@ -105,13 +105,13 @@ pub(crate) fn sibling_layer_dirs(index: &KmerIndex) -> OKIResult> { 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!( diff --git a/src/obikphylo/src/siblings/sankoff_bundle.rs b/src/obikphylo/src/siblings/sankoff_bundle.rs index b360721b..0c93928a 100644 --- a/src/obikphylo/src/siblings/sankoff_bundle.rs +++ b/src/obikphylo/src/siblings/sankoff_bundle.rs @@ -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 diff --git a/src/obikphylo/src/siblings/stats.rs b/src/obikphylo/src/siblings/stats.rs index 2a636187..7d38c9f2 100644 --- a/src/obikphylo/src/siblings/stats.rs +++ b/src/obikphylo/src/siblings/stats.rs @@ -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 diff --git a/src/obikphylo/src/siblings/tests.rs b/src/obikphylo/src/siblings/tests.rs index 929f49bf..2faba0cc 100644 --- a/src/obikphylo/src/siblings/tests.rs +++ b/src/obikphylo/src/siblings/tests.rs @@ -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