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:
Eric Coissac
@@ -0,0 +1,7 @@
|
||||
[package]
|
||||
name = "obikquery"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
|
||||
[dependencies]
|
||||
obikindex = { path = "../obikindex" }
|
||||
@@ -0,0 +1,10 @@
|
||||
//! Query-side operations on an `obikindex::KmerIndex`: staging ground for
|
||||
//! code currently living in `obikindex::index`'s query path, to be migrated
|
||||
//! here to lighten that crate — see `DevDocMD/` for the rationale. Kept as
|
||||
//! a separate crate — not a module of `obikindex` — so the dependency runs
|
||||
//! one way only (query code depends on the data model, never the reverse),
|
||||
//! same pattern as `obikindexer` for the build side.
|
||||
|
||||
mod query_layer;
|
||||
|
||||
pub use query_layer::{KmerDesc, QueryHit, QueryStats};
|
||||
@@ -0,0 +1,230 @@
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
|
||||
use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix};
|
||||
use obikseq::CanonicalKmer;
|
||||
use obikindex::layer::MphfLayer;
|
||||
use obikindex::{OKIError, OKIResult};
|
||||
|
||||
use obikindex::KmerIndex;
|
||||
|
||||
// ── per-layer query handle ────────────────────────────────────────────────────
|
||||
|
||||
enum QueryLayer {
|
||||
Presence(MphfLayer, PersistentBitMatrix),
|
||||
Count(MphfLayer, PersistentCompactIntMatrix),
|
||||
}
|
||||
|
||||
impl QueryLayer {
|
||||
fn open(layer_dir: &Path, with_counts: bool) -> OKIResult<Self> {
|
||||
let mphf = MphfLayer::open(layer_dir)?;
|
||||
let counts_dir = layer_dir.join("counts");
|
||||
let presence_dir = layer_dir.join("presence");
|
||||
|
||||
if with_counts && counts_dir.exists() {
|
||||
let mat = PersistentCompactIntMatrix::open(layer_dir).map_err(OKIError::Io)?;
|
||||
Ok(QueryLayer::Count(mphf, mat))
|
||||
} else if presence_dir.exists() || !counts_dir.exists() {
|
||||
// presence mode, or no matrix at all → Implicit handled inside open()
|
||||
let mat = PersistentBitMatrix::open(layer_dir).map_err(OKIError::Io)?;
|
||||
Ok(QueryLayer::Presence(mphf, mat))
|
||||
} else {
|
||||
// counts exist but not presence — count layer, no presence requested
|
||||
let mat = PersistentCompactIntMatrix::open(layer_dir).map_err(OKIError::Io)?;
|
||||
Ok(QueryLayer::Count(mphf, mat))
|
||||
}
|
||||
}
|
||||
|
||||
/// MPHF lookup only — no matrix access. `Some(slot)` on hit.
|
||||
fn find_slot(&self, kmer: CanonicalKmer) -> Option<usize> {
|
||||
match self {
|
||||
QueryLayer::Presence(mphf, _) | QueryLayer::Count(mphf, _) => mphf.find(kmer),
|
||||
}
|
||||
}
|
||||
|
||||
/// Number of genome columns this layer's matrix actually has. Bounds
|
||||
/// column-major iteration — usually equal to the index's `n_genomes`, but
|
||||
/// `PersistentBitMatrix::Implicit` (the documented mono-genome fast path)
|
||||
/// always reports exactly `1`, regardless of the index's real genome
|
||||
/// count, so callers must use this rather than assuming `n_genomes`.
|
||||
fn n_cols(&self) -> usize {
|
||||
match self {
|
||||
QueryLayer::Presence(_, mat) => mat.n_cols(),
|
||||
QueryLayer::Count(_, mat) => mat.n_cols(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Every nonzero `(idx into slots, col, value)` triple among `slots`.
|
||||
/// Format-agnostic: each matrix picks its own natural traversal
|
||||
/// (`PersistentBitMatrix::nonzero_iter` dispatches to a genuinely
|
||||
/// row-major decode on `Sparse`, not a column-major point-probe loop —
|
||||
/// see `DevDocMD/architecture/siblings.md`, "`query` never benefits
|
||||
/// from sparse row-major access"). Replaces the old per-`(genome,
|
||||
/// slot)` `col_value` point lookup, which this layer's `Sparse`
|
||||
/// presence matrices paid for badly: each such lookup rebuilt the
|
||||
/// entire row just to return one cell.
|
||||
fn nonzero_iter<'a>(
|
||||
&'a self,
|
||||
slots: &'a [usize],
|
||||
) -> Box<dyn Iterator<Item = (usize, usize, u32)> + 'a> {
|
||||
match self {
|
||||
QueryLayer::Presence(_, mat) => mat.nonzero_iter(slots),
|
||||
QueryLayer::Count(_, mat) => Box::new(mat.nonzero_iter(slots)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── KmerDesc — one occurrence of a k-mer in the query batch ──────────────────
|
||||
|
||||
/// Describes one occurrence of a (deduplicated) k-mer in the query batch:
|
||||
/// which sequence it came from, and its absolute s-mer position within it.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct KmerDesc {
|
||||
pub seq_idx: u32,
|
||||
pub pos: u32,
|
||||
}
|
||||
|
||||
/// Aggregate counters for one `query_partition_with` call — feeds the
|
||||
/// dedup-ratio and column-scan logging in `obikmer::cmd::query` (occurrences
|
||||
/// vs. unique k-mers is the whole justification for k-mer-level
|
||||
/// dereplication; columns scanned / `get()` calls quantify the column-major
|
||||
/// fetch's locality claim).
|
||||
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct QueryStats {
|
||||
/// Distinct canonical k-mers queried in this partition.
|
||||
pub n_unique_kmers: usize,
|
||||
/// Total `MphfLayer::find` calls issued (a k-mer tried against more than
|
||||
/// one layer before a hit, or against all layers on a miss, counts once
|
||||
/// per layer attempted).
|
||||
pub n_mphf_calls: usize,
|
||||
/// Distinct canonical k-mers that matched some layer.
|
||||
pub n_hits: usize,
|
||||
/// Total genome columns scanned across all hit layers (sum of
|
||||
/// `layer.n_cols()` over layers with at least one hit).
|
||||
pub n_columns_scanned: usize,
|
||||
/// Total `col_value` calls issued during the column-major fetch pass
|
||||
/// (`n_columns_scanned` × hits-per-layer, summed over layers).
|
||||
pub n_col_get_calls: usize,
|
||||
}
|
||||
|
||||
impl std::ops::AddAssign for QueryStats {
|
||||
fn add_assign(&mut self, other: Self) {
|
||||
self.n_unique_kmers += other.n_unique_kmers;
|
||||
self.n_mphf_calls += other.n_mphf_calls;
|
||||
self.n_hits += other.n_hits;
|
||||
self.n_columns_scanned += other.n_columns_scanned;
|
||||
self.n_col_get_calls += other.n_col_get_calls;
|
||||
}
|
||||
}
|
||||
|
||||
// ── QueryHit — one event delivered to query_partition_with's callback ───────
|
||||
|
||||
/// One event from [`KmerPartition::query_partition_with`]'s two-stage query:
|
||||
/// a `Found` event once per hit k-mer (stage 1, MPHF-only — mark the k-mer as
|
||||
/// indexed regardless of any genome's value), then a `Value` event per
|
||||
/// `(hit k-mer, genome)` pair with a nonzero matrix value (stage 2,
|
||||
/// column-major fetch). Carried as one enum, not two separate callbacks, so
|
||||
/// the caller only needs one `FnMut` closure — passing two closures that each
|
||||
/// need to mutably borrow the same accumulator does not borrow-check.
|
||||
pub enum QueryHit<'a> {
|
||||
Found(&'a [KmerDesc]),
|
||||
Value(&'a [KmerDesc], usize, u32),
|
||||
}
|
||||
|
||||
// ── KmerPartition::query_partition_with ──────────────────────────────────────
|
||||
|
||||
impl KmerIndex {
|
||||
/// Query a single partition for a pre-deduplicated map of canonical
|
||||
/// k-mers → their occurrences (`seq_idx`, `pos`) in the query batch.
|
||||
///
|
||||
/// Two stages:
|
||||
/// 1. **MPHF-only pass**: for each unique k-mer, try each layer's MPHF in
|
||||
/// turn (stopping at the first hit) and bucket confirmed hits by
|
||||
/// `(layer, slot)`. Emits one `QueryHit::Found` per hit k-mer. This
|
||||
/// stage's cost is independent of the index's genome count.
|
||||
/// 2. **Column-major fetch**: for each layer with at least one hit, walk
|
||||
/// its matrix **column by column** (genome by genome) — for each
|
||||
/// genome, scan the slots bucketed in stage 1 and look up their value.
|
||||
/// Emits one `QueryHit::Value` per nonzero `(k-mer, genome)` pair.
|
||||
/// Total lookups are the same as a row-major pass (`n_hits × n_cols`
|
||||
/// in the worst case); the win is memory locality — both persistent
|
||||
/// matrix formats are column-oriented on disk (one `mmap`'d region per
|
||||
/// genome), so scanning one column at a time touches far fewer
|
||||
/// distinct mmap regions than fetching one full row per hit.
|
||||
pub fn query_partition_with<F>(
|
||||
&self,
|
||||
part_idx: usize,
|
||||
kmers: &HashMap<CanonicalKmer, Vec<KmerDesc>>,
|
||||
n_genomes: usize,
|
||||
with_counts: bool,
|
||||
mut on_event: F,
|
||||
) -> OKIResult<QueryStats>
|
||||
where
|
||||
F: FnMut(QueryHit),
|
||||
{
|
||||
let mut stats = QueryStats::default();
|
||||
|
||||
if kmers.is_empty() {
|
||||
return Ok(stats);
|
||||
}
|
||||
|
||||
let index_dir = self.index_dir(part_idx);
|
||||
if !index_dir.exists() {
|
||||
return Ok(stats);
|
||||
}
|
||||
|
||||
let meta = self.partition_meta(part_idx)?;
|
||||
let layers: Vec<QueryLayer> = (0..meta.n_layers)
|
||||
.map(|i| QueryLayer::open(&self.layer_dir(part_idx, i), with_counts))
|
||||
.collect::<OKIResult<_>>()?;
|
||||
|
||||
// ── Stage 1: MPHF-only pass, bucket hits by (layer_idx, slot) ────────
|
||||
let mut by_layer: Vec<HashMap<usize, &Vec<KmerDesc>>> =
|
||||
(0..layers.len()).map(|_| HashMap::new()).collect();
|
||||
|
||||
for (kmer, descs) in kmers {
|
||||
stats.n_unique_kmers += 1;
|
||||
for (layer_idx, layer) in layers.iter().enumerate() {
|
||||
stats.n_mphf_calls += 1;
|
||||
if let Some(slot) = layer.find_slot(*kmer) {
|
||||
by_layer[layer_idx].insert(slot, descs);
|
||||
on_event(QueryHit::Found(descs));
|
||||
stats.n_hits += 1;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── Stage 2: nonzero-cell fetch, per layer ────────────────────────────
|
||||
// Format-agnostic — see `QueryLayer::nonzero_iter`. `n_cols` still
|
||||
// bounds accepted genome columns (Implicit reports fewer than
|
||||
// `n_genomes`; see `n_cols`'s doc), cells beyond it are dropped
|
||||
// rather than ever produced, since `nonzero_iter` only knows the
|
||||
// matrix's own column count, not the caller's `n_genomes`.
|
||||
for (layer_idx, slots) in by_layer.iter().enumerate() {
|
||||
if slots.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let layer = &layers[layer_idx];
|
||||
let n_cols = layer.n_cols().min(n_genomes);
|
||||
stats.n_columns_scanned += n_cols;
|
||||
|
||||
let slot_list: Vec<usize> = slots.keys().copied().collect();
|
||||
for (idx, g, v) in layer.nonzero_iter(&slot_list) {
|
||||
if g >= n_cols {
|
||||
continue;
|
||||
}
|
||||
stats.n_col_get_calls += 1;
|
||||
debug_assert_ne!(v, 0, "nonzero_iter must not yield zero-valued cells");
|
||||
let descs = slots[&slot_list[idx]];
|
||||
on_event(QueryHit::Value(descs, g, v));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(stats)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "tests/query_layer.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,94 @@
|
||||
use super::*;
|
||||
use obikindex::IndexConfig;
|
||||
|
||||
// ── QueryStats::AddAssign ───────────────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn query_stats_add_assign_sums_fields() {
|
||||
let mut total = QueryStats {
|
||||
n_unique_kmers: 3,
|
||||
n_mphf_calls: 5,
|
||||
n_hits: 2,
|
||||
n_columns_scanned: 1,
|
||||
n_col_get_calls: 7,
|
||||
};
|
||||
total += QueryStats {
|
||||
n_unique_kmers: 1,
|
||||
n_mphf_calls: 4,
|
||||
n_hits: 1,
|
||||
n_columns_scanned: 2,
|
||||
n_col_get_calls: 3,
|
||||
};
|
||||
assert_eq!(total.n_unique_kmers, 4);
|
||||
assert_eq!(total.n_mphf_calls, 9);
|
||||
assert_eq!(total.n_hits, 3);
|
||||
assert_eq!(total.n_columns_scanned, 3);
|
||||
assert_eq!(total.n_col_get_calls, 10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn query_stats_default_is_zero() {
|
||||
let s = QueryStats::default();
|
||||
assert_eq!(s.n_unique_kmers, 0);
|
||||
assert_eq!(s.n_mphf_calls, 0);
|
||||
assert_eq!(s.n_hits, 0);
|
||||
assert_eq!(s.n_columns_scanned, 0);
|
||||
assert_eq!(s.n_col_get_calls, 0);
|
||||
}
|
||||
|
||||
// ── query_partition_with on a not-yet-indexed partition ─────────────────────
|
||||
|
||||
/// A `KmerPartition` created but never taken through `build_layers` has no
|
||||
/// `index/` subdirectory under any partition — `query_partition_with` must
|
||||
/// recognise this and return default (all-zero) stats rather than erroring,
|
||||
/// exactly like an empty `kmers` map.
|
||||
#[test]
|
||||
fn query_partition_with_missing_index_dir_returns_default_stats() {
|
||||
let tmp = tempfile::tempdir().expect("tempdir");
|
||||
let config = IndexConfig {
|
||||
kmer_size: 21,
|
||||
minimizer_size: 9,
|
||||
n_bits: 2,
|
||||
with_counts: false,
|
||||
evidence: obikindex::layer::IndexMode::Exact,
|
||||
block_bits: 0,
|
||||
};
|
||||
let index = KmerIndex::create(tmp.path().join("idx"), config, None).expect("create index");
|
||||
|
||||
let mut kmers: HashMap<CanonicalKmer, Vec<KmerDesc>> = HashMap::new();
|
||||
// Any well-formed canonical k-mer works here — the call must return
|
||||
// before ever attempting an MPHF lookup, since `index/` doesn't exist.
|
||||
let kmer = CanonicalKmer::from_raw_unchecked(0u64);
|
||||
kmers.insert(kmer, vec![KmerDesc { seq_idx: 0, pos: 0 }]);
|
||||
|
||||
let stats = index
|
||||
.query_partition_with(0, &kmers, 1, false, |_event| {
|
||||
panic!("on_event must not be called: no index was built");
|
||||
})
|
||||
.expect("query_partition_with should not error on a missing index dir");
|
||||
|
||||
assert_eq!(stats, QueryStats::default());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn query_partition_with_empty_kmers_is_a_noop() {
|
||||
let tmp = tempfile::tempdir().expect("tempdir");
|
||||
let config = IndexConfig {
|
||||
kmer_size: 21,
|
||||
minimizer_size: 9,
|
||||
n_bits: 2,
|
||||
with_counts: false,
|
||||
evidence: obikindex::layer::IndexMode::Exact,
|
||||
block_bits: 0,
|
||||
};
|
||||
let index = KmerIndex::create(tmp.path().join("idx"), config, None).expect("create index");
|
||||
|
||||
let kmers: HashMap<CanonicalKmer, Vec<KmerDesc>> = HashMap::new();
|
||||
let stats = index
|
||||
.query_partition_with(0, &kmers, 1, false, |_event| {
|
||||
panic!("on_event must not be called on an empty kmer map");
|
||||
})
|
||||
.expect("query_partition_with on an empty map should not error");
|
||||
|
||||
assert_eq!(stats, QueryStats::default());
|
||||
}
|
||||
Reference in New Issue
Block a user