OBIKmers/obikmer
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 34 Wiki Activity

rename KmerPartition to KmerPartitions and update Mat enum

Browse Source 
Rename the KmerPartition type to KmerPartitions across obikindex, obikpartitionner, and obikphylo/siblings to reflect an updated data model. Update the Mat enum in siblings/cache.rs to add a SparsePresence variant and simplify opening logic by delegating sparse versus dense detection to PersistentBitMatrix. Apply consistent code formatting, import reordering, and multi-line refactoring throughout the affected modules.
This commit is contained in:
Eric Coissac
2026-08-20 15:57:33 +02:00
parent 76cbd3a886
commit f9ef6b8391
26 changed files with 1290 additions and 733 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/obikindex/src/index.rs
+102 -40
View File
@@ -2,7 +2,7 @@ use std::collections::BTreeMap;
use std::fs;
use std::path::{Path, PathBuf};
use obikpartitionner::{KmerPartition, KmerSpectrum, PARTITIONS_SUBDIR};
use obikpartitionner::{KmerPartitions, KmerSpectrum, PARTITIONS_SUBDIR};
use obisys::{Reporter, Stage, progress_bar};
use rayon::prelude::*;
use tracing::info;
@@ -16,7 +16,7 @@ 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: KmerPartition,
pub(crate) partition: KmerPartitions,
}
impl KmerIndex {
@@ -31,7 +31,7 @@ impl KmerIndex {
force: bool,
) -> OKIResult<Self> {
let root_path = path.as_ref().to_owned();
let partition = KmerPartition::create(
let partition = KmerPartitions::create(
&root_path,
config.n_bits,
config.kmer_size,
@@ -45,7 +45,11 @@ impl KmerIndex {
meta.genomes.push(info);
}
meta.write(&root_path)?;
Ok(Self { root_path, meta, partition })
Ok(Self {
root_path,
meta,
partition,
})
}
pub fn open<P: AsRef<Path>>(path: P) -> OKIResult<Self> {
@@ -53,13 +57,17 @@ 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 = KmerPartition::open_with_config(
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,
partition,
})
}
/// Return `true` if `path` contains an `index.meta` file.
@@ -96,12 +104,12 @@ impl KmerIndex {
pub(crate) fn create_skeleton<P: AsRef<Path>>(
output: P,
meta: &IndexMeta,
) -> OKIResult<KmerPartition> {
) -> OKIResult<KmerPartitions> {
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(KmerPartition::open_with_config(
Ok(KmerPartitions::open_with_config(
output,
meta.config.kmer_size,
meta.config.minimizer_size,
@@ -134,12 +142,24 @@ impl KmerIndex {
/// 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.
pub fn root_path(&self) -> &Path { &self.root_path }
pub fn meta(&self) -> &IndexMeta { &self.meta }
pub fn meta_mut(&mut self) -> &mut IndexMeta { &mut self.meta }
pub fn kmer_size(&self) -> usize { self.meta.config.kmer_size }
pub fn minimizer_size(&self) -> usize { self.meta.config.minimizer_size }
pub fn n_partitions(&self) -> usize { self.partition.n_partitions() }
pub fn root_path(&self) -> &Path {
&self.root_path
}
pub fn meta(&self) -> &IndexMeta {
&self.meta
}
pub fn meta_mut(&mut self) -> &mut IndexMeta {
&mut self.meta
}
pub fn kmer_size(&self) -> usize {
self.meta.config.kmer_size
}
pub fn minimizer_size(&self) -> usize {
self.meta.config.minimizer_size
}
pub fn n_partitions(&self) -> usize {
self.partition.n_partitions()
}
/// Number of layers per partition.
///
@@ -152,7 +172,7 @@ impl KmerIndex {
/// 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 KmerPartition {
pub fn partition_mut(&mut self) -> &mut KmerPartitions {
&mut self.partition
}
@@ -174,7 +194,11 @@ impl KmerIndex {
///
/// Writes `spectrums/{label}.json` and touches `count.done` upon completion.
/// Per-partition spectrum files are removed unless `keep_intermediate` is true.
pub fn dereplicate_and_count(&self, keep_intermediate: bool, rep: &mut Reporter) -> OKIResult<()> {
pub fn dereplicate_and_count(
&self,
keep_intermediate: bool,
rep: &mut Reporter,
) -> OKIResult<()> {
let t = Stage::start("dereplicate");
self.partition.dereplicate()?;
rep.push(t.stop());
@@ -189,11 +213,17 @@ impl KmerIndex {
}
fn write_spectrum(&self, sp: &KmerSpectrum) -> OKIResult<()> {
let label = self.meta.genomes.first().map(|g| g.label.as_str()).unwrap_or("unknown");
let label = self
.meta
.genomes
.first()
.map(|g| g.label.as_str())
.unwrap_or("unknown");
let spectrums_dir = self.root_path.join("spectrums");
fs::create_dir_all(&spectrums_dir)?;
let path = spectrums_dir.join(format!("{label}.json"));
let spectrum_map: BTreeMap<String, u64> = sp.counts
let spectrum_map: BTreeMap<String, u64> = sp
.counts
.iter()
.map(|(&c, &f)| (format!("{c:010}"), f))
.collect();
@@ -226,17 +256,28 @@ impl KmerIndex {
let order: Vec<usize> = (0..n).collect();
let runner = crate::numa::PartitionRunner::new();
runner.run(
&order,
|i| self.partition.build_index_layer(i, min_ab, max_ab, with_counts, &evidence, block_bits),
|i, n_kmers, _| {
if n_kmers > 0 {
total_kmers += n_kmers;
pb.inc(1);
pb.set_message(format!("{i}: {n_kmers} kmers"));
}
},
).map_err(OKIError::Partition)?;
runner
.run(
&order,
|i| {
self.partition.build_index_layer(
i,
min_ab,
max_ab,
with_counts,
&evidence,
block_bits,
)
},
|i, n_kmers, _| {
if n_kmers > 0 {
total_kmers += n_kmers;
pb.inc(1);
pb.set_message(format!("{i}: {n_kmers} kmers"));
}
},
)
.map_err(OKIError::Partition)?;
pb.finish_and_clear();
info!("done — {} total kmers indexed", total_kmers);
@@ -253,7 +294,7 @@ impl KmerIndex {
}
/// Borrow the inner partition for direct superkmer-level queries.
pub fn partition(&self) -> &KmerPartition {
pub fn partition(&self) -> &KmerPartitions {
&self.partition
}
@@ -282,12 +323,14 @@ impl KmerIndex {
&order,
|i| -> OKIResult<()> {
let index_dir = self.partition.index_dir(i);
if !index_dir.exists() { return Ok(()); }
if !index_dir.exists() {
return Ok(());
}
let n_layers = self.partition.n_layers(i)?;
for l in 0..n_layers {
let layer_dir = self.partition.layer_dir(i, l);
let presence_dir = layer_dir.join("presence");
let counts_dir = layer_dir.join("counts");
let counts_dir = layer_dir.join("counts");
if presence_dir.exists() {
if sparse {
pack_sparse_bit_matrix(&presence_dir).map_err(OKIError::Io)?;
@@ -295,11 +338,15 @@ impl KmerIndex {
pack_bit_matrix(&presence_dir).map_err(OKIError::Io)?;
}
}
if counts_dir.exists() { pack_compact_int_matrix(&counts_dir).map_err(OKIError::Io)?; }
if counts_dir.exists() {
pack_compact_int_matrix(&counts_dir).map_err(OKIError::Io)?;
}
}
Ok(())
},
|_, _, _| { pb.inc(1); },
|_, _, _| {
pb.inc(1);
},
)?;
pb.finish_and_clear();
Ok(())
@@ -318,18 +365,27 @@ impl KmerIndex {
.into_par_iter()
.filter_map(|i| {
let index_dir = self.partition.index_dir(i);
if !index_dir.exists() { return None; }
if !index_dir.exists() {
return None;
}
let n_layers = match self.partition.n_layers(i) {
Ok(n) => n,
Err(e) => return Some(OKIError::Io(std::io::Error::new(std::io::ErrorKind::Other, e.to_string()))),
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 = self.partition.layer_dir(i, l);
let meta_path = layer_dir.join(LayerMeta::FILENAME);
if meta_path.exists() { continue; }
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(),
Ok(r) => r.n_kmers(),
Err(e) => return Some(OKIError::Partition(e)),
};
if let Err(e) = LayerMeta::save(&layer_dir, n_kmers) {
@@ -340,7 +396,9 @@ impl KmerIndex {
})
.collect();
if let Some(e) = errors.into_iter().next() { return Err(e); }
if let Some(e) = errors.into_iter().next() {
return Err(e);
}
Ok(())
}
}
@@ -355,7 +413,11 @@ fn label_from_path(path: &Path) -> String {
while let Some(pos) = s.rfind('.') {
s.truncate(pos);
}
if s.is_empty() { "unknown".to_string() } else { s }
if s.is_empty() {
"unknown".to_string()
} else {
s
}
}
fn touch(path: &Path) -> Result<(), std::io::Error> {
src/obikindex/src/merge.rs
+29 -14
View File
@@ -193,7 +193,7 @@ impl KmerIndex {
let block_bits = dst.meta.config.block_bits;
// Pre-build source list once (avoid rebuilding per partition)
let srcs: Vec<(&obikpartitionner::KmerPartition, usize)> = remaining_sources
let srcs: Vec<(&obikpartitionner::KmerPartitions, usize)> = remaining_sources
.iter()
.map(|s| (&s.partition, s.meta.genomes.len()))
.collect();
@@ -221,19 +221,34 @@ impl KmerIndex {
let runner = crate::numa::PartitionRunner::new();
let mut part_stats: Vec<PartStat> = Vec::with_capacity(n_partitions);
runner.run(
&order,
|i| dst_partition.merge_partition(i, srcs, mode, n_dst_genomes, block_bits, evidence),
|i, g_len, dur| {
pb.inc(1);
debug!(
"partition {i}: done in {:.1}s — {} new kmers",
dur.as_secs_f64(),
g_len,
);
part_stats.push(PartStat { id: i, unitig_bytes: partition_sizes[i], g_len });
},
).map_err(OKIError::Partition)?;
runner
.run(
&order,
|i| {
dst_partition.merge_partition(
i,
srcs,
mode,
n_dst_genomes,
block_bits,
evidence,
)
},
|i, g_len, dur| {
pb.inc(1);
debug!(
"partition {i}: done in {:.1}s — {} new kmers",
dur.as_secs_f64(),
g_len,
);
part_stats.push(PartStat {
id: i,
unitig_bytes: partition_sizes[i],
g_len,
});
},
)
.map_err(OKIError::Partition)?;
pb.finish_and_clear();
src/obikindex/src/reindex.rs
+21 -10
View File
@@ -1,17 +1,17 @@
use std::fs;
use std::path::Path;
use obikpartitionner::KmerPartition;
use obikpartitionner::KmerPartitions;
use obilayeredmap::{IndexMode, layer::Layer};
use obisys::{Reporter, Stage, progress_bar};
use std::fs;
use std::path::Path;
use tracing::info;
use crate::error::{OKIError, OKIResult};
use crate::index::KmerIndex;
use crate::state::IndexState;
const EVIDENCE_FILE: &str = "evidence.bin";
const EVIDENCE_FILE: &str = "evidence.bin";
const FINGERPRINT_FILE: &str = "fingerprint.bin";
const UNITIG_IDX_FILE: &str = "unitigs.bin.idx";
const UNITIG_IDX_FILE: &str = "unitigs.bin.idx";
fn olm_to_oki(e: obilayeredmap::OLMError) -> OKIError {
OKIError::InvalidInput(e.to_string())
@@ -48,9 +48,13 @@ impl KmerIndex {
let runner = crate::numa::PartitionRunner::new();
runner.run(
&order,
|i| reindex_partition(&self.partition, i, &target, block_bits)
.map_err(|e| OKIError::InvalidInput(format!("partition {i}: {e}"))),
|_, _, _| { pb.inc(1); },
|i| {
reindex_partition(&self.partition, i, &target, block_bits)
.map_err(|e| OKIError::InvalidInput(format!("partition {i}: {e}")))
},
|_, _, _| {
pb.inc(1);
},
)?;
pb.finish_and_clear();
@@ -66,11 +70,18 @@ impl KmerIndex {
}
/// Process all layers of one partition's index directory.
fn reindex_partition(partition: &KmerPartition, i: usize, target: &IndexMode, block_bits: u8) -> OKIResult<()> {
fn reindex_partition(
partition: &KmerPartitions
, i: usize
, target: &IndexMode
, block_bits: u,
8) -> OKIResult<()> {
if !partition.index_dir(i).exists() {
return Ok(());
}
let n_layers = partition.n_layers(i).map_err(|e| OKIError::InvalidInput(e.to_string()))?;
let n_layers = partition
.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)?;
}
src/obikindex/src/select.rs
+56 -22
View File
@@ -1,6 +1,6 @@
use std::path::Path;
use obikpartitionner::{KmerPartition, OutputCol};
use obikpartitionner::{KmerPartitions, OutputCol};
use obisys::{Reporter, Stage, progress_bar};
use tracing::info;
@@ -35,31 +35,48 @@ impl KmerIndex {
let mut meta = IndexMeta::new(src.meta.config.clone());
meta.config.with_counts = !output_presence;
meta.genomes = specs.iter()
meta.genomes = specs
.iter()
.map(|s| GenomeInfo::new(s.label.clone()))
.collect();
let n_src_genomes = src.meta.genomes.len();
let n_partitions = src.partition.n_partitions();
let n_src_genomes = src.meta.genomes.len();
let n_partitions = src.partition.n_partitions();
let dst_partition = KmerIndex::create_skeleton(output, &meta)?;
info!(
"select: {} partition(s), {} source genome(s) → {} output column(s)",
n_partitions, n_src_genomes, specs.len(),
n_partitions,
n_src_genomes,
specs.len(),
);
let t = Stage::start("select");
let pb = progress_bar("select", n_partitions as u64, "partitions");
let t = Stage::start("select");
let pb = progress_bar("select", n_partitions as u64, "partitions");
let src_partition = &src.partition;
let order: Vec<usize> = (0..n_partitions).collect();
let runner = crate::numa::PartitionRunner::new();
runner.run(
&order,
|i| dst_partition.select_partition(src_partition, i, specs, n_src_genomes, threshold, output_presence, false),
|_, _, _| { pb.inc(1); },
).map_err(OKIError::Partition)?;
runner
.run(
&order,
|i| {
dst_partition.select_partition(
src_partition,
i,
specs,
n_src_genomes,
threshold,
output_presence,
false,
)
},
|_, _, _| {
pb.inc(1);
},
)
.map_err(OKIError::Partition)?;
pb.finish_and_clear();
rep.push(t.stop());
@@ -82,9 +99,9 @@ impl KmerIndex {
}
let n_src_genomes = self.meta.genomes.len();
let n_partitions = self.partition.n_partitions();
let n_partitions = self.partition.n_partitions();
let src_partition = KmerPartition::open_with_config(
let src_partition = KmerPartitions::open_with_config(
&self.root_path,
self.meta.config.kmer_size,
self.meta.config.minimizer_size,
@@ -93,26 +110,43 @@ impl KmerIndex {
info!(
"select (in-place): {} partition(s), {} source genome(s) → {} output column(s)",
n_partitions, n_src_genomes, specs.len(),
n_partitions,
n_src_genomes,
specs.len(),
);
let t = Stage::start("select");
let t = Stage::start("select");
let pb = progress_bar("select", n_partitions as u64, "partitions");
let partition = &self.partition;
let order: Vec<usize> = (0..n_partitions).collect();
let runner = crate::numa::PartitionRunner::new();
runner.run(
&order,
|i| partition.select_partition(&src_partition, i, specs, n_src_genomes, threshold, output_presence, true),
|_, _, _| { pb.inc(1); },
).map_err(OKIError::Partition)?;
runner
.run(
&order,
|i| {
partition.select_partition(
&src_partition,
i,
specs,
n_src_genomes,
threshold,
output_presence,
true,
)
},
|_, _, _| {
pb.inc(1);
},
)
.map_err(OKIError::Partition)?;
pb.finish_and_clear();
rep.push(t.stop());
self.meta.config.with_counts = !output_presence;
self.meta.genomes = specs.iter()
self.meta.genomes = specs
.iter()
.map(|s| GenomeInfo::new(s.label.clone()))
.collect();
self.meta.write(&self.root_path)?;