refactor: shift indexing algorithms to layer-centric partition APIs

Replaces manual filesystem path handling and parallel iteration with dedicated internal utilities. Introduces `par_over_layer0` and `layer0()` to abstract partition layer access, enabling partition-driven initialization and centralized progress tracking. Removes disabled or internal methods (`rewrite_config`, `open_data`) no longer part of the active interface. Updates error propagation and metadata persistence to align with the new layer-centric workflow.
This commit is contained in:
Eric Coissac
2026-08-22 13:51:42 +02:00
parent fc4464a0ef
commit 2419a6c21d
10 changed files with 234 additions and 199 deletions
+67 -67
View File
@@ -5,7 +5,7 @@ use std::sync::Arc;
use crate::layer::KmerLayer;
use crate::partition::KmerPartition;
use obisys::progress_bar;
use rayon::prelude::*;
// use rayon::prelude::*;
use obikseq::{set_k, set_m};
@@ -112,16 +112,16 @@ impl KmerIndex {
1usize << self.meta.config.n_bits
}
/// Path of partition `i`'s raw directory (`partitions/part_{i:05}`) —
/// delegates to `crate::partition`, the Partition tier's own naming
/// primitive (mirrors `layer_dir` delegating to `crate::layer`).
/// `obikindexer::algorithms::partitionner::PartitionRouter` reaches this
/// same directory only indirectly, through this method — `obikindexer`
/// depends on `KmerIndex`, not the other way around — see
/// `DevDocMD/implementation/partition_layer_cache.md`.
pub(crate) fn partition_dir(&self, i: usize) -> OKIResult<PathBuf> {
Ok(self.partition(i)?.dir().to_path_buf())
}
// /// Path of partition `i`'s raw directory (`partitions/part_{i:05}`) —
// /// delegates to `crate::partition`, the Partition tier's own naming
// /// primitive (mirrors `layer_dir` delegating to `crate::layer`).
// /// `obikindexer::algorithms::partitionner::PartitionRouter` reaches this
// /// same directory only indirectly, through this method — `obikindexer`
// /// depends on `KmerIndex`, not the other way around — see
// /// `DevDocMD/implementation/partition_layer_cache.md`.
// pub(crate) fn partition_dir(&self, i: usize) -> OKIResult<PathBuf> {
// Ok(self.partition(i)?.dir().to_path_buf())
// }
pub fn partition(&self, i: usize) -> OKIResult<KmerPartition> {
let n = self.n_partitions();
@@ -227,62 +227,62 @@ impl KmerIndex {
Ok(())
}
/// Write a `layer_meta.json` in any layer directory that is missing one.
///
/// Old indexes were built before this file was required. The number of
/// kmers is recovered from `unitigs.bin`, which is always present.
/// TODO: Should not exist anymore
pub(crate) fn upgrade_layer_meta(&self) -> OKIResult<()> {
use obicompactvec::LayerMeta;
use obiskio::UnitigFileReader;
// /// Write a `layer_meta.json` in any layer directory that is missing one.
// ///
// /// Old indexes were built before this file was required. The number of
// /// kmers is recovered from `unitigs.bin`, which is always present.
// /// TODO: Should not exist anymore
// pub(crate) fn upgrade_layer_meta(&self) -> OKIResult<()> {
// use obicompactvec::LayerMeta;
// use obiskio::UnitigFileReader;
let n = self.n_partitions();
let errors: Vec<_> = (0..n)
.into_par_iter()
.filter_map(|i| {
let index_dir = self.index_dir(i);
if !index_dir.exists() {
return None;
}
let n_layers = match self.n_layers(i) {
Ok(n) => n,
Err(e) => {
return Some(OKIError::Io(std::io::Error::new(
std::io::ErrorKind::Other,
e.to_string(),
)));
}
};
for l in 0..n_layers {
let layer_dir = match self.layer_dir(i, l) {
Ok(d) => d,
Err(e) => {
return Some(OKIError::Io(std::io::Error::new(
std::io::ErrorKind::Other,
e.to_string(),
)));
}
};
let meta_path = layer_dir.join(LayerMeta::FILENAME);
if meta_path.exists() {
continue;
}
let unitigs_path = layer_dir.join("unitigs.bin");
let n_kmers = match UnitigFileReader::open_sequential(&unitigs_path) {
Ok(r) => r.n_kmers(),
Err(e) => return Some(OKIError::Partition(e)),
};
if let Err(e) = LayerMeta::save(&layer_dir, n_kmers) {
return Some(OKIError::Io(e));
}
}
None
})
.collect();
// let n = self.n_partitions();
// let errors: Vec<_> = (0..n)
// .into_par_iter()
// .filter_map(|i| {
// let index_dir = self.index_dir(i);
// if !index_dir.exists() {
// return None;
// }
// let n_layers = match self.n_layers(i) {
// Ok(n) => n,
// Err(e) => {
// return Some(OKIError::Io(std::io::Error::new(
// std::io::ErrorKind::Other,
// e.to_string(),
// )));
// }
// };
// for l in 0..n_layers {
// let layer_dir = match self.layer_dir(i, l) {
// Ok(d) => d,
// Err(e) => {
// return Some(OKIError::Io(std::io::Error::new(
// std::io::ErrorKind::Other,
// e.to_string(),
// )));
// }
// };
// let meta_path = layer_dir.join(LayerMeta::FILENAME);
// if meta_path.exists() {
// continue;
// }
// let unitigs_path = layer_dir.join("unitigs.bin");
// let n_kmers = match UnitigFileReader::open_sequential(&unitigs_path) {
// Ok(r) => r.n_kmers(),
// Err(e) => return Some(OKIError::Partition(e)),
// };
// if let Err(e) = LayerMeta::save(&layer_dir, n_kmers) {
// return Some(OKIError::Io(e));
// }
// }
// None
// })
// .collect();
if let Some(e) = errors.into_iter().next() {
return Err(e);
}
Ok(())
}
// if let Some(e) = errors.into_iter().next() {
// return Err(e);
// }
// Ok(())
// }
}
+23 -23
View File
@@ -218,29 +218,29 @@ impl IndexMeta {
self.write_full(&on_disk)
}
/// Overwrite `config` and the whole `genomes` list at once, preserving
/// `state`. `config` otherwise never changes once an index exists —
/// this is the deliberate, rare exception for construction paths that
/// legitimately rewrite it in place (`select_in_place`, `reindex`).
/// The caller must refresh its own cached `Arc<IndexMeta>` afterward
/// (e.g. `self.meta = Arc::new(IndexMeta::open(self)?)`) — this method
/// only updates the file, it has no way to reach back into whatever
/// `KmerIndex` holds it.
/// TODO: This methode is strange
pub(crate) fn rewrite_config(
&self,
config: IndexConfig,
genomes: Vec<GenomeInfo>,
) -> io::Result<()> {
let _guard = self.lock.write().unwrap();
let state = Self::read_full(&self.root_path)?.state;
self.write_full(&IndexMetadata {
version: self.version,
config,
genomes,
state,
})
}
// /// Overwrite `config` and the whole `genomes` list at once, preserving
// /// `state`. `config` otherwise never changes once an index exists —
// /// this is the deliberate, rare exception for construction paths that
// /// legitimately rewrite it in place (`select_in_place`, `reindex`).
// /// The caller must refresh its own cached `Arc<IndexMeta>` afterward
// /// (e.g. `self.meta = Arc::new(IndexMeta::open(self)?)`) — this method
// /// only updates the file, it has no way to reach back into whatever
// /// `KmerIndex` holds it.
// /// TODO: This methode is strange
// pub(crate) fn rewrite_config(
// &self,
// config: IndexConfig,
// genomes: Vec<GenomeInfo>,
// ) -> io::Result<()> {
// let _guard = self.lock.write().unwrap();
// let state = Self::read_full(&self.root_path)?.state;
// self.write_full(&IndexMetadata {
// version: self.version,
// config,
// genomes,
// state,
// })
// }
pub fn set_state(&self, state: IndexState) -> io::Result<()> {
let _guard = self.lock.write().unwrap();
+15 -15
View File
@@ -1,4 +1,4 @@
use crate::layer::utils::layer_dir;
// use crate::layer::utils::layer_dir;
use obicompactvec::{
BinaryMatrix, PersistentBitMatrix, PersistentBitMatrixBuilder, PersistentCompactIntMatrix,
PersistentCompactIntMatrixBuilder, PersistentSparseBitMatrix,
@@ -24,20 +24,20 @@ pub trait LayerData: Sized {
fn read(&self, slot: usize) -> Self::Item;
}
/// Opens layer `i`'s data only, skipping the MPHF — for callers that only
/// need matrix-level operations (distance traits, column weights, group
/// filters, sub-matrix extraction) and never look up a kmer for this layer.
/// `TypedLayer<D>::open` always pays for the MPHF too (and doesn't expose `data`
/// once open), so it's the wrong tool for these; this is the other half of
/// the same `D::open(layer_dir)` call, without the MPHF alongside it.
///
/// Takes `(root, i)`, not a pre-built path: the caller names a layer
/// *number* within a partition it already knows the root of, the same
/// vocabulary as [`layer_dir`] and `LayeredMap` — never the `layer_N`
/// naming convention itself, which stays private to this crate.
pub(crate) fn open_data<D: LayerData>(root: &Path, i: usize) -> OKIResult<D> {
D::open(&layer_dir(root, i))
}
// /// Opens layer `i`'s data only, skipping the MPHF — for callers that only
// /// need matrix-level operations (distance traits, column weights, group
// /// filters, sub-matrix extraction) and never look up a kmer for this layer.
// /// `TypedLayer<D>::open` always pays for the MPHF too (and doesn't expose `data`
// /// once open), so it's the wrong tool for these; this is the other half of
// /// the same `D::open(layer_dir)` call, without the MPHF alongside it.
// ///
// /// Takes `(root, i)`, not a pre-built path: the caller names a layer
// /// *number* within a partition it already knows the root of, the same
// /// vocabulary as [`layer_dir`] and `LayeredMap` — never the `layer_N`
// /// naming convention itself, which stays private to this crate.
// pub(crate) fn open_data<D: LayerData>(root: &Path, i: usize) -> OKIResult<D> {
// D::open(&layer_dir(root, i))
// }
impl LayerData for () {
type Item = ();
+12 -31
View File
@@ -8,16 +8,14 @@
mod count;
mod kmer_sort;
use crate::algorithms::{existing_layer0, par_over_layer0};
use std::collections::BTreeMap;
use std::fs;
use std::io;
use std::sync::atomic::{AtomicU64, Ordering};
use obikalgorithm::Algorithm;
use obikindex::KmerIndex;
use obiskio::SKResult;
use obisys::Progress;
use rayon::prelude::*;
use sysinfo::System;
use count::count_partition;
@@ -86,7 +84,7 @@ impl Algorithm for Counter<'_> {
fn run(&mut self) -> obikalgorithm::Result<KmerSpectrum> {
let index = self.index;
let n_partitions = self.n_partitions;
let on_progress = &self.on_progress;
let on_progress = self.on_progress.as_deref();
let sys = System::new_all();
// available_memory() can return 0 on macOS when the compressor page count exceeds
// free+inactive+purgeable pages (sysinfo saturating_sub). Fall back to half of total.
@@ -96,40 +94,23 @@ impl Algorithm for Counter<'_> {
};
let n_threads = rayon::current_num_threads().max(1) as u64;
let chunk_kmers = chunk_size_from_ram(available / n_threads);
let done = AtomicU64::new(0);
let results: Vec<SKResult<()>> = (0..n_partitions)
.into_par_iter()
.map(|i| {
let dir = index.layer_dir(i, 0);
let dedup_path = obikindex::layer::dereplicated_superkmers_path(&dir);
let result = if dedup_path.exists() {
count_partition(&dir, &dedup_path, chunk_kmers)
} else {
Ok(())
};
if let Some(cb) = on_progress {
let pos = done.fetch_add(1, Ordering::Relaxed) + 1;
cb(Progress {
position: pos,
total: Some(n_partitions as u64),
});
}
result
})
.collect();
for r in results {
r?;
}
par_over_layer0(index, n_partitions, on_progress, |layer| {
let dedup_path = layer.dereplicated_superkmers_path();
if dedup_path.exists() {
count_partition(layer.dir(), &dedup_path, chunk_kmers)
} else {
Ok(())
}
})?;
// Aggregate per-partition spectra.
let mut counts: BTreeMap<u32, u64> = BTreeMap::new();
let mut f0: u64 = 0;
let mut f1: u64 = 0;
for i in 0..n_partitions {
let path = index.layer_dir(i, 0).join("kmer_spectrum_raw.json");
for layer in existing_layer0(index, n_partitions) {
let path = layer.dir().join("kmer_spectrum_raw.json");
if !path.exists() {
continue;
}
@@ -14,7 +14,7 @@ use niffler::send::compression::Format;
use niffler::Level;
use obikseq::superkmer::SuperKmer;
use obikseq::Sequence;
use obikindex::layer::{dereplicated_superkmers_path, raw_superkmers_path};
use obikindex::layer::KmerLayer;
use obiskio::{SKFileMeta, SKFileReader, SKFileWriter, SKResult};
/// Scratch-file extension for this algorithm's own intermediate split
@@ -72,15 +72,16 @@ fn remove_skmer_file(path: &Path) -> SKResult<()> {
/// Maximum value that fits in the 24-bit COUNT field of a SuperKmer header.
const MAX_SK_COUNT: u64 = (1 << 24) - 1;
/// Deduplicate one partition's layer-0 directory in place (two-phase split
/// + merge): `raw_superkmers_path(dir)` -> `dereplicated_superkmers_path(dir)`.
pub(crate) fn dereplicate_partition(dir: &Path, level: Level, n_temp: usize) -> SKResult<()> {
let raw_path = raw_superkmers_path(dir);
/// Deduplicate one partition's layer-0 in place (two-phase split
/// + merge): `layer.raw_superkmers_path()` -> `layer.dereplicated_superkmers_path()`.
pub(crate) fn dereplicate_partition(layer: &KmerLayer, level: Level, n_temp: usize) -> SKResult<()> {
let raw_path = layer.raw_superkmers_path();
if !raw_path.exists() {
return Ok(());
}
let dir = layer.dir();
let out_path = dereplicated_superkmers_path(dir);
let out_path = layer.dereplicated_superkmers_path();
let mut writer = SKFileWriter::create_with(&out_path, Format::Zstd, level)?;
if n_temp == 1 {
@@ -7,12 +7,9 @@
mod dereplicate;
use std::sync::atomic::{AtomicU64, Ordering};
use crate::algorithms::par_over_layer0;
use niffler::Level;
use obiskio::SKResult;
use obisys::Progress;
use rayon::prelude::*;
use sysinfo::System;
use obikalgorithm::Algorithm;
@@ -79,7 +76,7 @@ impl Algorithm for Dereplicator<'_> {
let level = self.level;
let index = self.index;
let n_partitions = self.n_partitions;
let on_progress = &self.on_progress;
let on_progress = self.on_progress.as_deref();
let sys = System::new_all();
// available_memory() can return 0 on macOS when the compressor page count exceeds
// free+inactive+purgeable pages (sysinfo saturating_sub). Fall back to half of total.
@@ -89,30 +86,12 @@ impl Algorithm for Dereplicator<'_> {
};
let n_threads = rayon::current_num_threads().max(1) as u64;
let available_per_thread = available / n_threads;
let done = AtomicU64::new(0);
let results: Vec<SKResult<()>> = (0..n_partitions)
.into_par_iter()
.map(|i| {
let dir = index.layer_dir(i, 0);
let result = if dir.exists() {
let raw_path = obikindex::layer::raw_superkmers_path(&dir);
let n_buckets = optimal_buckets(&raw_path, available_per_thread);
dereplicate_partition(&dir, level, n_buckets)
} else {
Ok(())
};
if let Some(cb) = on_progress {
let pos = done.fetch_add(1, Ordering::Relaxed) + 1;
cb(Progress { position: pos, total: Some(n_partitions as u64) });
}
result
})
.collect();
for r in results {
r?;
}
par_over_layer0(index, n_partitions, on_progress, |layer| {
let raw_path = layer.raw_superkmers_path();
let n_buckets = optimal_buckets(&raw_path, available_per_thread);
dereplicate_partition(&layer, level, n_buckets)
})?;
Ok(())
}
}
@@ -5,8 +5,8 @@
//! `DevDocMD/implementation/partition_layer_cache.md`.
use obikalgorithm::Algorithm;
use obikindex::{KmerIndex, PartitionRunner};
use obisys::Progress;
use obikindex::KmerIndex;
use obisys::{PartitionRunner, Progress};
use crate::extensions::PrivateBuilder;
+63
View File
@@ -12,3 +12,66 @@ pub mod counter;
pub mod dereplicator;
pub mod layer_builder;
pub mod partitionner;
use std::sync::atomic::{AtomicU64, Ordering};
use obikindex::KmerIndex;
use obikindex::layer::KmerLayer;
use obiskio::SKResult;
use obisys::Progress;
use rayon::prelude::*;
use crate::extensions::PrivateBuilder;
/// Runs `f` over every partition's layer 0 in parallel, skipping partitions
/// whose layer 0 doesn't exist yet (e.g. no data routed to them) — the
/// per-partition-in-parallel shape shared by
/// [`dereplicator::Dereplicator::run`] and [`counter::Counter::run`]:
/// resolve layer 0 via [`PrivateBuilder::layer0`], apply `f`, report one
/// progress tick per partition regardless of outcome, then propagate the
/// first error once every partition has run.
pub(crate) fn par_over_layer0<F>(
index: &KmerIndex,
n_partitions: usize,
on_progress: Option<&(dyn Fn(Progress) + Sync)>,
f: F,
) -> SKResult<()>
where
F: Fn(KmerLayer) -> SKResult<()> + Sync,
{
let done = AtomicU64::new(0);
let results: Vec<SKResult<()>> = (0..n_partitions)
.into_par_iter()
.map(|i| {
let result = match index.layer0(i) {
Ok(layer) => f(layer),
Err(_) => Ok(()),
};
if let Some(cb) = on_progress {
let pos = done.fetch_add(1, Ordering::Relaxed) + 1;
cb(Progress {
position: pos,
total: Some(n_partitions as u64),
});
}
result
})
.collect();
for r in results {
r?;
}
Ok(())
}
/// Sequential counterpart to [`par_over_layer0`]: every partition's layer 0
/// that exists, in order, skipping those that don't — for callers that
/// aggregate into shared state (e.g. [`counter::Counter::run`]'s spectrum
/// aggregation) rather than run independent per-partition work in parallel.
pub(crate) fn existing_layer0(
index: &KmerIndex,
n_partitions: usize,
) -> impl Iterator<Item = KmerLayer> + '_ {
(0..n_partitions).filter_map(move |i| index.layer0(i).ok())
}
@@ -212,14 +212,6 @@ impl<'a> PartitionRouter<'a> {
// ── private ───────────────────────────────────────────────────────────────
/// Directory of partition `i`'s layer 0 — every raw/dereplicated
/// superkmer file and provisional `mphf1.bin`/`counts1.bin` this router
/// produces lives here, alongside where `build_index_layer`
/// (`obikindex`) will later turn it into the real layer 0.
fn layer0_dir(&self, i: usize) -> PathBuf {
obikindex::layer::layer_dir(&self.index.index_dir(i), 0)
}
fn check_not_closed(&self) -> SKResult<()> {
if self.closed {
Err(io::Error::new(io::ErrorKind::BrokenPipe, "write to closed PartitionRouter").into())
@@ -230,9 +222,14 @@ impl<'a> PartitionRouter<'a> {
fn ensure_writer(&mut self, partition: usize) -> SKResult<&mut SKFileWriter> {
if self.writers[partition].is_none() {
let dir = self.layer0_dir(partition);
KmerLayer::create(&dir).map_err(|e| io::Error::other(e.to_string()))?;
let file_path = obikindex::layer::raw_superkmers_path(&dir);
let part = self
.index
.partition(partition)
.map_err(|e| io::Error::other(e.to_string()))?;
let layer = KmerLayer::new(&part, 0)
.create()
.map_err(|e| io::Error::other(e.to_string()))?;
let file_path = layer.raw_superkmers_path();
let writer = SKFileWriter::create_with(file_path, Format::Zstd, self.level)?;
self.writers[partition] = Some(writer);
}
+30 -16
View File
@@ -34,7 +34,7 @@ use epserde::prelude::*;
use obicompactvec::{PersistentCompactIntMatrix, PersistentCompactIntVec};
use obidebruinj::GraphDeBruijn;
use obikindex::layer::IndexMode;
use obikindex::layer::{TypedLayer, meta::PartitionMeta};
use obikindex::layer::{KmerLayer, TypedLayer};
use obikindex::{KmerIndex, OKIError, OKIResult};
use crate::graph_pipeline::{materialize_layer, write_graph_as_unitigs};
use obiskio::{SKError, SKFileMeta, SKFileReader};
@@ -54,6 +54,12 @@ pub(crate) trait PrivateBuilder {
/// Mark dereplicate+count as complete (`IndexState::Counted`).
fn mark_counted(&self) -> OKIResult<()>;
/// Partition `i`'s layer 0 — the one and only layer every pipeline
/// algorithm in this crate (`partitionner`/`dereplicator`/`counter`/
/// `layer_builder`) reads from or writes to; later layers only exist
/// after a `merge`, which lives in `obikindex`, not here.
fn layer0(&self, i: usize) -> OKIResult<KmerLayer>;
/// Mark layer construction as complete (`IndexState::Indexed`).
fn mark_indexed(&self) -> OKIResult<()>;
@@ -98,6 +104,10 @@ impl PrivateBuilder for KmerIndex {
self.meta().mark_counted().map_err(OKIError::Io)
}
fn layer0(&self, i: usize) -> OKIResult<KmerLayer> {
self.partition(i)?.layer(0)
}
fn mark_indexed(&self) -> OKIResult<()> {
self.meta().mark_indexed().map_err(OKIError::Io)
}
@@ -134,8 +144,9 @@ impl PrivateBuilder for KmerIndex {
mode: &IndexMode,
block_bits: u8,
) -> Result<usize, SKError> {
let layer0_dir = self.layer_dir(i, 0);
let dedup_path = obikindex::layer::dereplicated_superkmers_path(&layer0_dir);
let layer0 = self.layer0(i).map_err(|e| io::Error::other(e.to_string()))?;
let layer0_dir = layer0.dir();
let dedup_path = layer0.dereplicated_superkmers_path();
if !dedup_path.exists() {
return Ok(0);
}
@@ -185,34 +196,37 @@ impl PrivateBuilder for KmerIndex {
}
let n_kmers = if with_counts {
let n = write_graph_as_unitigs(g, &layer0_dir)?;
let n = write_graph_as_unitigs(g, layer0_dir)
.map_err(|e| io::Error::other(e.to_string()))?;
TypedLayer::<PersistentCompactIntMatrix>::build(
&layer0_dir,
layer0_dir,
block_bits,
mode,
|kmer| match (&mphf1_opt, &counts1_opt) {
(Some(mphf), Some(counts)) => counts.get(mphf.index(&kmer.raw())),
_ => 1,
},
)?;
)
.map_err(|e| io::Error::other(e.to_string()))?;
n
} else {
materialize_layer(g, &layer0_dir, block_bits, mode)?
materialize_layer(g, layer0_dir, block_bits, mode)
.map_err(|e| io::Error::other(e.to_string()))?
};
let index_dir = layer0_dir.parent().expect("layer_dir has a parent");
PartitionMeta {
n_layers: 1,
mode: mode.clone(),
}
.save(index_dir)?;
Ok(n_kmers)
}
fn remove_build_artifacts(&self, i: usize) {
let layer0_dir = self.layer_dir(i, 0);
let dedup = obikindex::layer::dereplicated_superkmers_path(&layer0_dir);
let layer0 = match self.layer0(i) {
Ok(layer) => layer,
Err(e) => {
eprintln!("warning: could not locate layer 0 of partition {i}: {e}");
return;
}
};
let layer0_dir = layer0.dir();
let dedup = layer0.dereplicated_superkmers_path();
remove_if_exists(&SKFileMeta::sidecar_path(&dedup));
remove_if_exists(&dedup);
remove_if_exists(&layer0_dir.join("mphf1.bin"));