Extract index modules into specialized workspace subcrates
This commit partitions the obikindex crate into multiple focused subcrates (obikfilter, obikmerge, obikquery, obikrebuild, obikselect, obikstats, obikdump, and obikidxcache) to reduce coupling and clarify module boundaries. It standardizes error handling across the workspace using OKIError and OKIResult, updates index APIs to support lazy, disk-backed partition access, and migrates NUMA system utilities to a new obisys crate. All modifications are structural, focusing on dependency graph expansion, import path updates, and API surface reorganization without altering core runtime behavior.
This commit is contained in:
@@ -0,0 +1,10 @@
|
||||
[package]
|
||||
name = "obikselect"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
|
||||
[dependencies]
|
||||
obikindex = { path = "../obikindex" }
|
||||
obicompactvec = { path = "../obicompactvec" }
|
||||
obisys = { path = "../obisys" }
|
||||
tracing = "0.1.44"
|
||||
@@ -0,0 +1,11 @@
|
||||
//! Genome-column selection/aggregation on an `obikindex::KmerIndex`:
|
||||
//! projecting/aggregating genome columns (`AggOp`) into new output columns
|
||||
//! (`OutputCol`), either into a new index (`select`) or in place
|
||||
//! (`select_in_place`). Orthogonal to per-kmer filtering (`obikfilter`),
|
||||
//! which operates on rows, not columns. Kept as a separate crate — not a
|
||||
//! module of `obikindex` — so the dependency runs one way only.
|
||||
|
||||
mod select;
|
||||
mod select_layer;
|
||||
|
||||
pub use select_layer::{AggOp, OutputCol};
|
||||
@@ -0,0 +1,154 @@
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
|
||||
use obikindex::IndexBuilder;
|
||||
use crate::OutputCol;
|
||||
use obisys::{Reporter, Stage, progress_bar, PartitionRunner};
|
||||
use tracing::info;
|
||||
|
||||
use obikindex::{OKIError, OKIResult};
|
||||
use obikindex::KmerIndex;
|
||||
use obikindex::{GenomeInfo, IndexMeta};
|
||||
use obikindex::IndexState;
|
||||
|
||||
impl KmerIndex {
|
||||
/// Create a new index at `output` by projecting/aggregating the genome columns
|
||||
/// of `src` according to `specs`.
|
||||
///
|
||||
/// `output_presence` — if true, output uses bit matrices (0/1), regardless of
|
||||
/// whether the source stores counts. The caller is responsible for ensuring all
|
||||
/// specs use logical operators when `output_presence=true` on a count source.
|
||||
pub fn select<P: AsRef<Path>>(
|
||||
output: P,
|
||||
src: &KmerIndex,
|
||||
specs: &[OutputCol],
|
||||
threshold: u32,
|
||||
output_presence: bool,
|
||||
force: bool,
|
||||
rep: &mut Reporter,
|
||||
) -> OKIResult<Self> {
|
||||
let output = output.as_ref();
|
||||
|
||||
if src.state()? != IndexState::Indexed {
|
||||
return Err(OKIError::NotIndexed(src.root_path.clone()));
|
||||
}
|
||||
|
||||
KmerIndex::clear_output_for_create(output, force)?;
|
||||
|
||||
let mut config = src.meta.config.clone();
|
||||
config.with_counts = !output_presence;
|
||||
let genomes: Vec<GenomeInfo> = specs
|
||||
.iter()
|
||||
.map(|s| GenomeInfo::new(s.label.clone()))
|
||||
.collect();
|
||||
|
||||
let n_src_genomes = src.meta.genomes().map_err(OKIError::Io)?.len();
|
||||
let n_partitions = src.n_partitions();
|
||||
|
||||
let dst_partition = KmerIndex::create_skeleton(output, config, genomes)?;
|
||||
|
||||
info!(
|
||||
"select: {} partition(s), {} source genome(s) → {} output column(s)",
|
||||
n_partitions,
|
||||
n_src_genomes,
|
||||
specs.len(),
|
||||
);
|
||||
|
||||
let t = Stage::start("select");
|
||||
let pb = progress_bar("select", n_partitions as u64, "partitions");
|
||||
|
||||
let order: Vec<usize> = (0..n_partitions).collect();
|
||||
let runner = PartitionRunner::new();
|
||||
runner
|
||||
.run(
|
||||
&order,
|
||||
|i| {
|
||||
dst_partition.select_partition(
|
||||
src,
|
||||
i,
|
||||
specs,
|
||||
n_src_genomes,
|
||||
threshold,
|
||||
output_presence,
|
||||
false,
|
||||
)
|
||||
},
|
||||
|_, _, _| {
|
||||
pb.inc(1);
|
||||
},
|
||||
)
|
||||
.map_err(OKIError::Partition)?;
|
||||
|
||||
pb.finish_and_clear();
|
||||
rep.push(t.stop());
|
||||
|
||||
KmerIndex::finalize_indexed(output, rep)
|
||||
}
|
||||
|
||||
/// Rewrite the genome columns of this index in-place according to `specs`.
|
||||
///
|
||||
/// The MPHF and unitig files are unchanged; only data matrices are rewritten.
|
||||
pub fn select_in_place(
|
||||
&mut self,
|
||||
specs: &[OutputCol],
|
||||
threshold: u32,
|
||||
output_presence: bool,
|
||||
rep: &mut Reporter,
|
||||
) -> OKIResult<()> {
|
||||
if self.state()? != IndexState::Indexed {
|
||||
return Err(OKIError::NotIndexed(self.root_path.clone()));
|
||||
}
|
||||
|
||||
let n_src_genomes = self.meta.genomes().map_err(OKIError::Io)?.len();
|
||||
let n_partitions = self.n_partitions();
|
||||
|
||||
info!(
|
||||
"select (in-place): {} partition(s), {} source genome(s) → {} output column(s)",
|
||||
n_partitions,
|
||||
n_src_genomes,
|
||||
specs.len(),
|
||||
);
|
||||
|
||||
let t = Stage::start("select");
|
||||
let pb = progress_bar("select", n_partitions as u64, "partitions");
|
||||
|
||||
let order: Vec<usize> = (0..n_partitions).collect();
|
||||
let runner = PartitionRunner::new();
|
||||
runner
|
||||
.run(
|
||||
&order,
|
||||
|i| {
|
||||
self.select_partition(
|
||||
self,
|
||||
i,
|
||||
specs,
|
||||
n_src_genomes,
|
||||
threshold,
|
||||
output_presence,
|
||||
true,
|
||||
)
|
||||
},
|
||||
|_, _, _| {
|
||||
pb.inc(1);
|
||||
},
|
||||
)
|
||||
.map_err(OKIError::Partition)?;
|
||||
|
||||
pb.finish_and_clear();
|
||||
rep.push(t.stop());
|
||||
|
||||
let mut config = self.meta.config.clone();
|
||||
config.with_counts = !output_presence;
|
||||
let genomes: Vec<GenomeInfo> = specs
|
||||
.iter()
|
||||
.map(|s| GenomeInfo::new(s.label.clone()))
|
||||
.collect();
|
||||
self.meta.rewrite_config(config, genomes).map_err(OKIError::Io)?;
|
||||
self.meta = Arc::new(IndexMeta::open(self).map_err(OKIError::Io)?);
|
||||
|
||||
let t_pack = Stage::start("pack");
|
||||
self.pack_matrices(false)?;
|
||||
rep.push(t_pack.stop());
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,297 @@
|
||||
use std::fs;
|
||||
use std::io;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use obicompactvec::{
|
||||
ColGroup, MatrixGroupOps, PersistentBitMatrix, PersistentBitMatrixBuilder,
|
||||
PersistentCompactIntMatrix, PersistentCompactIntMatrixBuilder,
|
||||
};
|
||||
use obikindex::{OKIError, OKIResult};
|
||||
|
||||
use obikindex::KmerIndex;
|
||||
|
||||
// ── AggOp ─────────────────────────────────────────────────────────────────────
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum AggOp {
|
||||
Any,
|
||||
All,
|
||||
None,
|
||||
Sum,
|
||||
Min,
|
||||
Max,
|
||||
}
|
||||
|
||||
impl AggOp {
|
||||
pub fn is_logical(self) -> bool {
|
||||
matches!(self, AggOp::Any | AggOp::All | AggOp::None)
|
||||
}
|
||||
}
|
||||
|
||||
// ── OutputCol ─────────────────────────────────────────────────────────────────
|
||||
|
||||
pub struct OutputCol {
|
||||
pub label: String,
|
||||
pub indices: Vec<usize>,
|
||||
pub op: AggOp,
|
||||
}
|
||||
|
||||
// ── Helpers ───────────────────────────────────────────────────────────────────
|
||||
|
||||
/// Copy all plain files (not subdirectories) from `src_dir` to `dst_dir`.
|
||||
fn copy_layer_files(src_dir: &Path, dst_dir: &Path) -> io::Result<()> {
|
||||
for entry in fs::read_dir(src_dir)? {
|
||||
let entry = entry?;
|
||||
let path = entry.path();
|
||||
if path.is_file() {
|
||||
fs::copy(&path, dst_dir.join(entry.file_name()))?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// ── fill_builders ─────────────────────────────────────────────────────────────
|
||||
|
||||
fn fill_builders(
|
||||
specs: &[OutputCol],
|
||||
src_layer_dir: &Path,
|
||||
src_is_count: bool,
|
||||
threshold: u32,
|
||||
output_presence: bool,
|
||||
mut dst_bit: Option<&mut PersistentBitMatrixBuilder>,
|
||||
mut dst_int: Option<&mut PersistentCompactIntMatrixBuilder>,
|
||||
) -> OKIResult<()> {
|
||||
if src_is_count {
|
||||
let mat = PersistentCompactIntMatrix::open(src_layer_dir).map_err(OKIError::Io)?;
|
||||
for spec in specs {
|
||||
let g = ColGroup::new(&spec.label, spec.indices.clone());
|
||||
if output_presence {
|
||||
let b = dst_bit.as_deref_mut().unwrap();
|
||||
match spec.op {
|
||||
AggOp::Any => {
|
||||
b.add_col_from(&mat.partial_group_any(&g, threshold).map_err(OKIError::Io)?)
|
||||
}
|
||||
AggOp::All => {
|
||||
b.add_col_from(&mat.partial_group_all(&g, threshold).map_err(OKIError::Io)?)
|
||||
}
|
||||
AggOp::None => {
|
||||
b.add_col_from(&mat.partial_group_none(&g, threshold).map_err(OKIError::Io)?)
|
||||
}
|
||||
AggOp::Sum => {
|
||||
b.add_col_from_int(&mat.partial_group_sum(&g).map_err(OKIError::Io)?)
|
||||
}
|
||||
AggOp::Min => {
|
||||
b.add_col_from_int(&mat.partial_group_min(&g).map_err(OKIError::Io)?)
|
||||
}
|
||||
AggOp::Max => {
|
||||
b.add_col_from_int(&mat.partial_group_max(&g).map_err(OKIError::Io)?)
|
||||
}
|
||||
}
|
||||
.map_err(OKIError::Io)?;
|
||||
} else {
|
||||
let b = dst_int.as_deref_mut().unwrap();
|
||||
match spec.op {
|
||||
AggOp::Sum => b.add_col_from(&mat.partial_group_sum(&g).map_err(OKIError::Io)?),
|
||||
AggOp::Min => b.add_col_from(&mat.partial_group_min(&g).map_err(OKIError::Io)?),
|
||||
AggOp::Max => b.add_col_from(&mat.partial_group_max(&g).map_err(OKIError::Io)?),
|
||||
AggOp::Any => b.add_col_from_bit(
|
||||
&mat.partial_group_any(&g, threshold).map_err(OKIError::Io)?,
|
||||
),
|
||||
AggOp::All => b.add_col_from_bit(
|
||||
&mat.partial_group_all(&g, threshold).map_err(OKIError::Io)?,
|
||||
),
|
||||
AggOp::None => b.add_col_from_bit(
|
||||
&mat.partial_group_none(&g, threshold).map_err(OKIError::Io)?,
|
||||
),
|
||||
}
|
||||
.map_err(OKIError::Io)?;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
let mat = PersistentBitMatrix::open(src_layer_dir).map_err(OKIError::Io)?;
|
||||
for spec in specs {
|
||||
let g = ColGroup::new(&spec.label, spec.indices.clone());
|
||||
if output_presence {
|
||||
let b = dst_bit.as_deref_mut().unwrap();
|
||||
match spec.op {
|
||||
AggOp::Any => {
|
||||
b.add_col_from(&mat.partial_group_any(&g, 1).map_err(OKIError::Io)?)
|
||||
}
|
||||
AggOp::All => {
|
||||
b.add_col_from(&mat.partial_group_all(&g, 1).map_err(OKIError::Io)?)
|
||||
}
|
||||
AggOp::None => {
|
||||
b.add_col_from(&mat.partial_group_none(&g, 1).map_err(OKIError::Io)?)
|
||||
}
|
||||
AggOp::Sum => {
|
||||
b.add_col_from_int(&mat.partial_group_sum(&g).map_err(OKIError::Io)?)
|
||||
}
|
||||
AggOp::Min => {
|
||||
b.add_col_from_int(&mat.partial_group_min(&g).map_err(OKIError::Io)?)
|
||||
}
|
||||
AggOp::Max => {
|
||||
b.add_col_from_int(&mat.partial_group_max(&g).map_err(OKIError::Io)?)
|
||||
}
|
||||
}
|
||||
.map_err(OKIError::Io)?;
|
||||
} else {
|
||||
let b = dst_int.as_deref_mut().unwrap();
|
||||
match spec.op {
|
||||
AggOp::Sum => b.add_col_from(&mat.partial_group_sum(&g).map_err(OKIError::Io)?),
|
||||
AggOp::Min => b.add_col_from(&mat.partial_group_min(&g).map_err(OKIError::Io)?),
|
||||
AggOp::Max => b.add_col_from(&mat.partial_group_max(&g).map_err(OKIError::Io)?),
|
||||
AggOp::Any => {
|
||||
b.add_col_from_bit(&mat.partial_group_any(&g, 1).map_err(OKIError::Io)?)
|
||||
}
|
||||
AggOp::All => {
|
||||
b.add_col_from_bit(&mat.partial_group_all(&g, 1).map_err(OKIError::Io)?)
|
||||
}
|
||||
AggOp::None => {
|
||||
b.add_col_from_bit(&mat.partial_group_none(&g, 1).map_err(OKIError::Io)?)
|
||||
}
|
||||
}
|
||||
.map_err(OKIError::Io)?;
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// ── KmerPartition::select_partition ──────────────────────────────────────────
|
||||
|
||||
impl KmerIndex {
|
||||
/// Rewrite the data matrices of partition `i` in `src` into `self`.
|
||||
///
|
||||
/// `specs` defines the output columns (projection/aggregation).
|
||||
/// `output_presence` — if true, all output builders use bit (0/1) format.
|
||||
/// `in_place` — `self` and `src` share the same root; write to temp dirs then swap.
|
||||
pub fn select_partition(
|
||||
&self,
|
||||
src: &KmerIndex,
|
||||
i: usize,
|
||||
specs: &[OutputCol],
|
||||
_n_src_genomes: usize,
|
||||
threshold: u32,
|
||||
output_presence: bool,
|
||||
in_place: bool,
|
||||
) -> OKIResult<()> {
|
||||
let src_index_dir = src.index_dir(i);
|
||||
if !src_index_dir.exists() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let n_src_layers = src.n_layers(i)?;
|
||||
if n_src_layers == 0 {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let dst_index_dir = self.index_dir(i);
|
||||
if !in_place {
|
||||
fs::create_dir_all(&dst_index_dir)?;
|
||||
}
|
||||
|
||||
let data_subdir = if output_presence {
|
||||
"presence"
|
||||
} else {
|
||||
"counts"
|
||||
};
|
||||
|
||||
for l in 0..n_src_layers {
|
||||
let src_layer_dir = src.layer_dir(i, l);
|
||||
if !src_layer_dir.exists() {
|
||||
continue;
|
||||
}
|
||||
|
||||
let dst_layer_dir = self.layer_dir(i, l);
|
||||
|
||||
let counts_dir = src_layer_dir.join("counts");
|
||||
let presence_dir = src_layer_dir.join("presence");
|
||||
let src_is_count = counts_dir.exists() && !presence_dir.exists();
|
||||
|
||||
// Determine number of slots and detect implicit layers.
|
||||
let n = if counts_dir.exists() {
|
||||
PersistentCompactIntMatrix::open(&src_layer_dir)
|
||||
.map_err(OKIError::Io)?
|
||||
.n()
|
||||
} else if presence_dir.exists() {
|
||||
PersistentBitMatrix::open(&src_layer_dir)
|
||||
.map_err(OKIError::Io)?
|
||||
.n()
|
||||
} else {
|
||||
// Implicit single-genome layer: no data matrix needed in output either.
|
||||
if !in_place {
|
||||
fs::create_dir_all(&dst_layer_dir)?;
|
||||
copy_layer_files(&src_layer_dir, &dst_layer_dir)?;
|
||||
}
|
||||
continue;
|
||||
};
|
||||
|
||||
// Choose the output data directory (temp name for in-place).
|
||||
let (dst_data_dir, final_data_dir): (PathBuf, PathBuf) = if in_place {
|
||||
let tmp = dst_layer_dir.join(format!("{data_subdir}_new"));
|
||||
let perm = dst_layer_dir.join(data_subdir);
|
||||
(tmp, perm)
|
||||
} else {
|
||||
let perm = dst_layer_dir.join(data_subdir);
|
||||
(perm.clone(), perm)
|
||||
};
|
||||
|
||||
if !in_place {
|
||||
fs::create_dir_all(&dst_layer_dir)?;
|
||||
copy_layer_files(&src_layer_dir, &dst_layer_dir)?;
|
||||
}
|
||||
fs::create_dir_all(&dst_data_dir)?;
|
||||
|
||||
let (mut dst_bit, mut dst_int) = if output_presence {
|
||||
(
|
||||
Some(PersistentBitMatrixBuilder::new(n, &dst_data_dir).map_err(OKIError::Io)?),
|
||||
None,
|
||||
)
|
||||
} else {
|
||||
(
|
||||
None,
|
||||
Some(
|
||||
PersistentCompactIntMatrixBuilder::new(n, &dst_data_dir)
|
||||
.map_err(OKIError::Io)?,
|
||||
),
|
||||
)
|
||||
};
|
||||
|
||||
fill_builders(
|
||||
specs,
|
||||
&src_layer_dir,
|
||||
src_is_count,
|
||||
threshold,
|
||||
output_presence,
|
||||
dst_bit.as_mut(),
|
||||
dst_int.as_mut(),
|
||||
)?;
|
||||
|
||||
if output_presence {
|
||||
dst_bit.unwrap().close().map_err(OKIError::Io)?;
|
||||
} else {
|
||||
dst_int.unwrap().close().map_err(OKIError::Io)?;
|
||||
}
|
||||
|
||||
// In-place: swap old data dir for new.
|
||||
if in_place {
|
||||
let old_data_dir = if src_is_count {
|
||||
dst_layer_dir.join("counts")
|
||||
} else {
|
||||
dst_layer_dir.join("presence")
|
||||
};
|
||||
if old_data_dir.exists() {
|
||||
fs::remove_dir_all(&old_data_dir)?;
|
||||
}
|
||||
fs::rename(&dst_data_dir, &final_data_dir)?;
|
||||
}
|
||||
}
|
||||
|
||||
if !in_place {
|
||||
src.partition_meta(i)?
|
||||
.save(&dst_index_dir)?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user