refactor: Shift KmerIndex ownership to Arc for thread-safe sharing
Replaces lifetime-bound references with runtime reference counting across multiple crates. This enables safe concurrent access across parallel workers without explicit cloning or manual lifetime management. Introduces the `query` and `utils` CLI commands in obikmer2, along with supporting modules for batch processing, sparse indexing, sliding-window findere logic, and output formatting. Updates dependency manifests and aligns test suites with the new ownership model.
This commit is contained in:
Eric Coissac
@@ -4,4 +4,15 @@ version = "0.1.0"
|
||||
edition = "2024"
|
||||
|
||||
[dependencies]
|
||||
obikindex = { path = "../obikindex" }
|
||||
obikindex = { path = "../obikindex" }
|
||||
obikidxcache = { path = "../obikidxcache" }
|
||||
obikseq = { path = "../obikseq" }
|
||||
obikrope = { path = "../obikrope" }
|
||||
obiread = { path = "../obiread" }
|
||||
obiskbuilder = { path = "../obiskbuilder" }
|
||||
serde_json = "1"
|
||||
tracing = "0.1.44"
|
||||
|
||||
[dev-dependencies]
|
||||
obikseq = { path = "../obikseq", features = ["test-utils"] }
|
||||
tempfile = "3"
|
||||
|
||||
@@ -0,0 +1,79 @@
|
||||
use std::collections::HashMap;
|
||||
|
||||
use obikseq::CanonicalKmer;
|
||||
use obiread::record::SeqRecord;
|
||||
use obiskbuilder::SuperKmerIter;
|
||||
|
||||
use crate::query_layer::KmerDesc;
|
||||
|
||||
/// A batch of query sequences, with k-mers deduplicated directly (not just at
|
||||
/// the superkmer level) and pre-split by partition.
|
||||
///
|
||||
/// Superkmer *construction* (`SuperKmerIter`) is still required — it's the
|
||||
/// mechanism that computes minimizers and partition routing — but the dedup
|
||||
/// key is the canonical k-mer, not the superkmer: two different superkmers
|
||||
/// that happen to share a k-mer (read overlaps, repeats, a SNP splitting an
|
||||
/// otherwise-identical run) are deduplicated too, not just identical whole
|
||||
/// superkmers. This also means each unique k-mer triggers at most one MPHF
|
||||
/// lookup, not one per occurrence.
|
||||
pub(crate) struct QueryBatch {
|
||||
/// Sequence ids in batch order.
|
||||
pub(crate) ids: Vec<String>,
|
||||
/// Raw sequence bytes (for output), in batch order.
|
||||
pub(crate) seqs: Vec<Vec<u8>>,
|
||||
/// Total kmer count per sequence (used for `--detail` coverage allocation).
|
||||
pub(crate) n_kmers: Vec<u32>,
|
||||
/// Deduplicated k-mer occurrences, one map per partition.
|
||||
pub(crate) by_partition: Vec<HashMap<CanonicalKmer, Vec<KmerDesc>>>,
|
||||
}
|
||||
|
||||
impl QueryBatch {
|
||||
/// Build a batch from a vec of parsed sequence records, deduplicating
|
||||
/// k-mers and routing them to partitions in the same pass.
|
||||
pub(crate) fn from_records(
|
||||
records: Vec<SeqRecord>,
|
||||
k: usize,
|
||||
level_max: usize,
|
||||
theta: f64,
|
||||
n_partitions: usize,
|
||||
) -> Self {
|
||||
let mut ids = Vec::with_capacity(records.len());
|
||||
let mut seqs = Vec::with_capacity(records.len());
|
||||
let mut n_kmers = Vec::with_capacity(records.len());
|
||||
let mask = (n_partitions as u64) - 1;
|
||||
let mut by_partition: Vec<HashMap<CanonicalKmer, Vec<KmerDesc>>> =
|
||||
(0..n_partitions).map(|_| HashMap::new()).collect();
|
||||
|
||||
for (seq_idx, record) in records.into_iter().enumerate() {
|
||||
let mut kmer_offset = 0u32;
|
||||
|
||||
for rsk in SuperKmerIter::new(&record.normalized, k, level_max, theta) {
|
||||
let part_idx = (rsk.minimizer().seq_hash() & mask) as usize;
|
||||
let map = &mut by_partition[part_idx];
|
||||
for (j, kmer) in rsk.superkmer().iter_canonical_kmers().enumerate() {
|
||||
map.entry(kmer).or_default().push(KmerDesc {
|
||||
seq_idx: seq_idx as u32,
|
||||
pos: kmer_offset + j as u32,
|
||||
});
|
||||
}
|
||||
let n = (rsk.seql() - k + 1) as u32;
|
||||
kmer_offset += n;
|
||||
}
|
||||
|
||||
ids.push(record.id);
|
||||
seqs.push(record.sequence);
|
||||
n_kmers.push(kmer_offset);
|
||||
}
|
||||
|
||||
Self {
|
||||
ids,
|
||||
seqs,
|
||||
n_kmers,
|
||||
by_partition,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "tests/batch.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,279 @@
|
||||
use std::time::Instant;
|
||||
|
||||
use obikidxcache::index_cache::IndexCache;
|
||||
use obikindex::GenomeInfo;
|
||||
use obikrope::Rope;
|
||||
use obikseq::CanonicalKmer;
|
||||
use obiread::record::parse_chunk;
|
||||
use tracing::debug;
|
||||
|
||||
use crate::batch::QueryBatch;
|
||||
use crate::findere::{ConfirmedHit, sparse_findere_for_genome};
|
||||
use crate::output::emit_batch;
|
||||
use crate::query_layer::{QueryHit, QueryPartition, QueryStats};
|
||||
use crate::smer_index::SmerIndex;
|
||||
|
||||
pub(crate) struct SeqAcc {
|
||||
pub(crate) kmer_count: u32,
|
||||
pub(crate) kmer_missing: u32,
|
||||
pub(crate) genome_totals: Vec<u32>,
|
||||
}
|
||||
|
||||
impl SeqAcc {
|
||||
fn new(n_genomes: usize) -> Self {
|
||||
Self {
|
||||
kmer_count: 0,
|
||||
kmer_missing: 0,
|
||||
genome_totals: vec![0u32; n_genomes],
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Turn one chunk of raw sequence bytes into an obitools-style annotated
|
||||
/// FASTA byte buffer, matching its k-mers against `cache`'s index.
|
||||
///
|
||||
/// `cache` is expected to already hold every partition/layer this chunk
|
||||
/// might touch (built once, up front, for the whole run — see
|
||||
/// `obikidxcache::IndexCache::new`), so every lookup here is a plain
|
||||
/// in-memory operation, never a disk open.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn process_chunk(
|
||||
cache: &IndexCache,
|
||||
rope: Rope,
|
||||
k: usize,
|
||||
n_genomes: usize,
|
||||
n_partitions: usize,
|
||||
with_counts: bool,
|
||||
effective_z: usize,
|
||||
detail: bool,
|
||||
count_missing: bool,
|
||||
force_presence: bool,
|
||||
presence_threshold: u32,
|
||||
genomes: &[GenomeInfo],
|
||||
) -> Vec<u8> {
|
||||
let chunk_start = Instant::now();
|
||||
let chunk_bytes = rope.len();
|
||||
|
||||
let records = parse_chunk(&rope, k);
|
||||
if records.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
let batch = QueryBatch::from_records(records, k, 6, 0.7, n_partitions);
|
||||
let n_seqs = batch.ids.len();
|
||||
|
||||
// Estimate QueryBatch::by_partition's actual memory footprint: the
|
||||
// k-mer-level dedup map — one HashMap<CanonicalKmer, Vec<KmerDesc>> per
|
||||
// partition, sized by *unique* k-mers, not shrunk by dedup. On real
|
||||
// workloads with a low intra-chunk duplication rate this can dwarf every
|
||||
// other per-chunk structure, including the sparse Findere ones logged
|
||||
// further down. Measured by allocated capacity, not logical length, to
|
||||
// reflect real memory pressure (HashMap/Vec growth slack) — `by_partition`
|
||||
// is alive for the entire process_chunk call (never drained, only
|
||||
// iterated by reference), so this is its footprint for the whole chunk
|
||||
// lifetime, not a transient.
|
||||
let hashmap_slot_bytes = (std::mem::size_of::<CanonicalKmer>()
|
||||
+ std::mem::size_of::<Vec<crate::query_layer::KmerDesc>>()
|
||||
+ 1) as u64; // +1 ≈ hashbrown control byte per slot
|
||||
let by_partition_map_bytes: u64 = batch
|
||||
.by_partition
|
||||
.iter()
|
||||
.map(|m| m.capacity() as u64 * hashmap_slot_bytes)
|
||||
.sum();
|
||||
let by_partition_desc_bytes: u64 = batch
|
||||
.by_partition
|
||||
.iter()
|
||||
.flat_map(|m| m.values())
|
||||
.map(|v| v.capacity() as u64 * std::mem::size_of::<crate::query_layer::KmerDesc>() as u64)
|
||||
.sum();
|
||||
let by_partition_bytes = by_partition_map_bytes + by_partition_desc_bytes;
|
||||
|
||||
debug!(
|
||||
n_unique_kmers_total = batch.by_partition.iter().map(|m| m.len() as u64).sum::<u64>(),
|
||||
by_partition_map_bytes,
|
||||
by_partition_desc_bytes,
|
||||
by_partition_bytes,
|
||||
chunk_bytes,
|
||||
"by_partition memory retained"
|
||||
);
|
||||
|
||||
// Sparse bookkeeping for the whole chunk:
|
||||
// - smer_index: O(total_smers) — is this s-mer in the index at all.
|
||||
// - by_genome[g]: raw (seq_idx, pos_smer, value) hits for genome g, only
|
||||
// ever containing nonzero entries (query_partition_with never emits a
|
||||
// QueryHit::Value for a zero value) — empty for every genome this chunk
|
||||
// never matched, which is the common case for unrelated queries.
|
||||
let mut smer_index = SmerIndex::new(&batch.n_kmers);
|
||||
let mut by_genome: Vec<Vec<(u32, u32, u32)>> = (0..n_genomes).map(|_| Vec::new()).collect();
|
||||
|
||||
// Dedup-ratio bookkeeping: occurrences (from batch.n_kmers, computed
|
||||
// before dedup) vs. unique k-mers actually queried (query_stats) — the
|
||||
// entire justification for k-mer-level dereplication. If this ratio
|
||||
// stays close to 1.0 on real data, dereplication isn't paying for
|
||||
// itself and that should show up here.
|
||||
let n_occurrences: u64 = batch.n_kmers.iter().map(|&n| n as u64).sum();
|
||||
let mut query_stats = QueryStats::default();
|
||||
|
||||
for (part_idx, kmers) in batch.by_partition.iter().enumerate() {
|
||||
if kmers.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
let stats = cache.query_partition_with(part_idx, kmers, n_genomes, |event| match event {
|
||||
QueryHit::Found(descs) => {
|
||||
for desc in descs {
|
||||
smer_index.mark_found(desc.seq_idx as usize, desc.pos as usize);
|
||||
}
|
||||
}
|
||||
QueryHit::Value(descs, g, v) => {
|
||||
for desc in descs {
|
||||
by_genome[g].push((desc.seq_idx, desc.pos, v));
|
||||
}
|
||||
}
|
||||
});
|
||||
query_stats += stats;
|
||||
}
|
||||
|
||||
debug!(
|
||||
n_occurrences,
|
||||
n_unique_kmers = query_stats.n_unique_kmers,
|
||||
n_mphf_calls = query_stats.n_mphf_calls,
|
||||
n_hits = query_stats.n_hits,
|
||||
n_columns_scanned = query_stats.n_columns_scanned,
|
||||
n_col_get_calls = query_stats.n_col_get_calls,
|
||||
"k-mer dedup + column-major fetch"
|
||||
);
|
||||
|
||||
// ── Sparse Findere: per-genome run detection + sliding-window minimum ────
|
||||
//
|
||||
// Confirmed z-windows, per genome, replace the dense win_min matrix:
|
||||
// total retained memory is O(actual hits), not O(total_smers × n_genomes)
|
||||
// — the whole point of this pass. See sparse_findere_for_genome's doc for
|
||||
// why run detection is equivalent to the dense scan's semantics.
|
||||
let presence = force_presence || !with_counts;
|
||||
let threshold = presence_threshold;
|
||||
let z = effective_z;
|
||||
|
||||
let n_kmers_out: Vec<usize> = batch
|
||||
.n_kmers
|
||||
.iter()
|
||||
.map(|&n| {
|
||||
let n = n as usize;
|
||||
if n >= z { n - z + 1 } else { 0 }
|
||||
})
|
||||
.collect();
|
||||
let mut out_offsets = Vec::with_capacity(n_seqs + 1);
|
||||
{
|
||||
let mut total = 0usize;
|
||||
out_offsets.push(0);
|
||||
for &n in &n_kmers_out {
|
||||
total += n;
|
||||
out_offsets.push(total);
|
||||
}
|
||||
}
|
||||
let total_out = *out_offsets.last().unwrap_or(&0);
|
||||
|
||||
let n_dense_would_be = n_occurrences as u64 * n_genomes as u64;
|
||||
let mut n_sparse_entries = 0u64;
|
||||
let mut n_runs_total = 0usize;
|
||||
let mut run_len_total = 0usize;
|
||||
|
||||
let mut confirmed_by_genome: Vec<Vec<ConfirmedHit>> = Vec::with_capacity(n_genomes);
|
||||
for hits in &mut by_genome {
|
||||
n_sparse_entries += hits.len() as u64;
|
||||
let (confirmed, n_runs, run_len) = sparse_findere_for_genome(hits, z, presence, threshold);
|
||||
n_runs_total += n_runs;
|
||||
run_len_total += run_len;
|
||||
confirmed_by_genome.push(confirmed);
|
||||
}
|
||||
|
||||
debug!(
|
||||
n_dense_would_be,
|
||||
n_sparse_entries,
|
||||
n_runs = n_runs_total,
|
||||
avg_run_len = if n_runs_total > 0 { run_len_total as f64 / n_runs_total as f64 } else { 0.0 },
|
||||
z,
|
||||
"sparse Findere"
|
||||
);
|
||||
|
||||
// Actual bytes retained by the sparse hit structures (by_genome +
|
||||
// confirmed_by_genome, both alive simultaneously at this point).
|
||||
const HIT_ENTRY_BYTES: u64 = std::mem::size_of::<(u32, u32, u32)>() as u64;
|
||||
let by_genome_bytes: u64 = by_genome.iter().map(|v| v.capacity() as u64 * HIT_ENTRY_BYTES).sum();
|
||||
let confirmed_bytes: u64 = confirmed_by_genome.iter().map(|v| v.capacity() as u64 * HIT_ENTRY_BYTES).sum();
|
||||
let retained_bytes = by_genome_bytes + confirmed_bytes;
|
||||
|
||||
debug!(
|
||||
by_genome_bytes,
|
||||
confirmed_bytes,
|
||||
retained_bytes,
|
||||
chunk_bytes,
|
||||
empirical_multiplier = retained_bytes as f64 / chunk_bytes.max(1) as f64,
|
||||
"sparse memory retained"
|
||||
);
|
||||
|
||||
// ── Accumulate: genome totals (per genome, from confirmed hits) ──────────
|
||||
let mut accs: Vec<SeqAcc> = (0..n_seqs).map(|_| SeqAcc::new(n_genomes)).collect();
|
||||
let mut confirmed_any = vec![false; total_out];
|
||||
|
||||
for (g, hits) in confirmed_by_genome.iter().enumerate() {
|
||||
for &(seq_idx, pos_out, c) in hits {
|
||||
let abs_out = out_offsets[seq_idx as usize] + pos_out as usize;
|
||||
confirmed_any[abs_out] = true;
|
||||
accs[seq_idx as usize].genome_totals[g] += c;
|
||||
}
|
||||
}
|
||||
|
||||
// ── Accumulate: kmer_count / kmer_missing (per position, genome-independent) ─
|
||||
for seq_idx in 0..n_seqs {
|
||||
let out_n = n_kmers_out[seq_idx];
|
||||
let acc = &mut accs[seq_idx];
|
||||
for pos in 0..out_n {
|
||||
let abs_out = out_offsets[seq_idx] + pos;
|
||||
if confirmed_any[abs_out] {
|
||||
acc.kmer_count += 1;
|
||||
} else if !smer_index.is_in_index(seq_idx, pos) {
|
||||
acc.kmer_missing += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── Coverage (--detail): densify only when actually requested ────────────
|
||||
let mut cov: Vec<Vec<Vec<u32>>> = if detail {
|
||||
n_kmers_out.iter().map(|&n| vec![vec![0u32; n]; n_genomes]).collect()
|
||||
} else {
|
||||
Vec::new()
|
||||
};
|
||||
if detail {
|
||||
for (g, hits) in confirmed_by_genome.iter().enumerate() {
|
||||
for &(seq_idx, pos_out, c) in hits {
|
||||
cov[seq_idx as usize][g][pos_out as usize] += c;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Capacity estimate: actual sequence + ID bytes, plus JSON overhead per record.
|
||||
let seq_bytes: usize = batch.seqs.iter().map(|s| s.len()).sum();
|
||||
let id_bytes: usize = batch.ids.iter().map(|s| s.len()).sum();
|
||||
let cap = seq_bytes + id_bytes + n_seqs * (4 + 50 + n_genomes * 20) + 100;
|
||||
let mut buf = Vec::with_capacity(cap);
|
||||
emit_batch(
|
||||
&batch,
|
||||
&accs,
|
||||
genomes,
|
||||
count_missing,
|
||||
detail,
|
||||
&cov,
|
||||
&mut buf,
|
||||
);
|
||||
|
||||
debug!(
|
||||
chunk_bytes,
|
||||
n_seqs,
|
||||
n_smers = batch.n_kmers.iter().map(|&n| n as u64).sum::<u64>(),
|
||||
wall_ms = chunk_start.elapsed().as_millis() as u64,
|
||||
"process_chunk"
|
||||
);
|
||||
|
||||
buf
|
||||
}
|
||||
@@ -0,0 +1,75 @@
|
||||
use std::collections::VecDeque;
|
||||
|
||||
/// One confirmed z-window: genome `g`'s window ending at k-mer `pos` (the
|
||||
/// *leftmost* s-mer of the window, i.e. the k_user-mer's output position) is
|
||||
/// fully present and nonzero, with window-minimum `value`.
|
||||
pub(crate) type ConfirmedHit = (u32, u32, u32); // (seq_idx, pos_out, value)
|
||||
|
||||
/// Reduce one genome's raw sparse s-mer hits — `(seq_idx, pos_smer, raw_value)`,
|
||||
/// unsorted, exactly as delivered by `QueryHit::Value` — into confirmed
|
||||
/// z-windows, without ever visiting a position that had no hit at all.
|
||||
///
|
||||
/// A z-window is confirmed only when all z s-mers in it are present *and*
|
||||
/// nonzero for this genome (matching the dense sliding-window's semantics,
|
||||
/// where "not in index" or a zero value both contribute 0 to the window
|
||||
/// minimum) — which can only happen inside a maximal run of consecutive
|
||||
/// `pos_smer` values for the same sequence. `hits` is sorted in place by
|
||||
/// `(seq_idx, pos_smer)` to expose those runs; the monotone-deque
|
||||
/// window-minimum then runs per run, on run-relative indices, identical in
|
||||
/// spirit to the dense version's whole-sequence scan.
|
||||
///
|
||||
/// Returns the confirmed hits plus `(n_runs, total_run_len)` for logging —
|
||||
/// a low average run length relative to `z` means most hits fail to form a
|
||||
/// complete window.
|
||||
pub(crate) fn sparse_findere_for_genome(
|
||||
hits: &mut [(u32, u32, u32)],
|
||||
z: usize,
|
||||
presence: bool,
|
||||
threshold: u32,
|
||||
) -> (Vec<ConfirmedHit>, usize, usize) {
|
||||
hits.sort_unstable_by_key(|&(seq, pos, _)| (seq, pos));
|
||||
|
||||
let mut confirmed = Vec::new();
|
||||
let mut n_runs = 0usize;
|
||||
let mut total_run_len = 0usize;
|
||||
let mut dq: VecDeque<(usize, u32)> = VecDeque::new(); // (run-relative index, value)
|
||||
|
||||
let mut i = 0;
|
||||
while i < hits.len() {
|
||||
let seq = hits[i].0;
|
||||
let mut j = i + 1;
|
||||
while j < hits.len() && hits[j].0 == seq && hits[j].1 == hits[j - 1].1 + 1 {
|
||||
j += 1;
|
||||
}
|
||||
let run = &hits[i..j];
|
||||
n_runs += 1;
|
||||
total_run_len += run.len();
|
||||
|
||||
dq.clear();
|
||||
for (k, &(_, pos, val)) in run.iter().enumerate() {
|
||||
while dq.back().map_or(false, |&(_, v)| v >= val) {
|
||||
dq.pop_back();
|
||||
}
|
||||
dq.push_back((k, val));
|
||||
while dq.front().map_or(false, |&(fk, _)| fk + z <= k) {
|
||||
dq.pop_front();
|
||||
}
|
||||
if k + 1 >= z {
|
||||
let win_min = dq.front().unwrap().1;
|
||||
if win_min > 0 {
|
||||
let pos_out = pos + 1 - z as u32;
|
||||
let c = if presence { u32::from(win_min >= threshold) } else { win_min };
|
||||
confirmed.push((seq, pos_out, c));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
i = j;
|
||||
}
|
||||
|
||||
(confirmed, n_runs, total_run_len)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "tests/findere.rs"]
|
||||
mod tests;
|
||||
@@ -1,10 +1,16 @@
|
||||
//! 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.
|
||||
//! Query-side operations on an `obikidxcache::IndexCache`: matching query
|
||||
//! sequences' k-mers against an already-built `obikindex::KmerIndex` and
|
||||
//! formatting the result. 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;
|
||||
mod batch;
|
||||
mod chunk;
|
||||
mod findere;
|
||||
mod smer_index;
|
||||
mod output;
|
||||
|
||||
pub use query_layer::{KmerDesc, QueryHit, QueryStats};
|
||||
pub use query_layer::{KmerDesc, QueryHit, QueryPartition, QueryStats};
|
||||
pub use chunk::process_chunk;
|
||||
|
||||
@@ -0,0 +1,52 @@
|
||||
use std::io::Write;
|
||||
|
||||
use obikindex::GenomeInfo;
|
||||
|
||||
use crate::batch::QueryBatch;
|
||||
use crate::chunk::SeqAcc;
|
||||
|
||||
pub(crate) fn emit_batch(
|
||||
batch: &QueryBatch,
|
||||
accs: &[SeqAcc],
|
||||
genomes: &[GenomeInfo],
|
||||
count_missing: bool,
|
||||
detail: bool,
|
||||
cov: &[Vec<Vec<u32>>],
|
||||
out: &mut impl Write,
|
||||
) {
|
||||
for (seq_idx, (id, seq)) in batch.ids.iter().zip(batch.seqs.iter()).enumerate() {
|
||||
let acc = &accs[seq_idx];
|
||||
let mut ann = serde_json::Map::new();
|
||||
|
||||
ann.insert("kmer_count".into(), acc.kmer_count.into());
|
||||
if count_missing {
|
||||
ann.insert("kmer_missing".into(), acc.kmer_missing.into());
|
||||
}
|
||||
|
||||
let mut match_map = serde_json::Map::new();
|
||||
for (g, genome) in genomes.iter().enumerate() {
|
||||
if acc.genome_totals[g] != 0 {
|
||||
match_map.insert(genome.label.clone(), acc.genome_totals[g].into());
|
||||
}
|
||||
}
|
||||
ann.insert("kmer_strict_matches".into(), match_map.into());
|
||||
|
||||
if detail && !cov.is_empty() {
|
||||
let mut cov_map = serde_json::Map::new();
|
||||
for (g, genome) in genomes.iter().enumerate() {
|
||||
let v: Vec<serde_json::Value> = cov[seq_idx][g].iter().map(|&x| x.into()).collect();
|
||||
cov_map.insert(genome.label.clone(), v.into());
|
||||
}
|
||||
ann.insert("coverage".into(), cov_map.into());
|
||||
}
|
||||
|
||||
// OBITools4 FASTA format: >id {"key":value,...}
|
||||
let _ = out.write_all(b">");
|
||||
let _ = out.write_all(id.as_bytes());
|
||||
let _ = out.write_all(b" ");
|
||||
let _ = serde_json::to_writer(&mut *out, &ann);
|
||||
let _ = out.write_all(b"\n");
|
||||
let _ = out.write_all(seq);
|
||||
let _ = out.write_all(b"\n");
|
||||
}
|
||||
}
|
||||
@@ -1,78 +1,7 @@
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
|
||||
use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix};
|
||||
use obikidxcache::index_cache::IndexCache;
|
||||
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 ──────────────────
|
||||
|
||||
@@ -85,7 +14,7 @@ pub struct KmerDesc {
|
||||
}
|
||||
|
||||
/// Aggregate counters for one `query_partition_with` call — feeds the
|
||||
/// dedup-ratio and column-scan logging in `obikmer::cmd::query` (occurrences
|
||||
/// dedup-ratio and column-scan logging in `obikquery::chunk` (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).
|
||||
@@ -93,7 +22,7 @@ pub struct KmerDesc {
|
||||
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
|
||||
/// Total MPHF-membership checks 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,
|
||||
@@ -102,7 +31,7 @@ pub struct QueryStats {
|
||||
/// 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
|
||||
/// Total nonzero-cell fetches issued during the column-major fetch pass
|
||||
/// (`n_columns_scanned` × hits-per-layer, summed over layers).
|
||||
pub n_col_get_calls: usize,
|
||||
}
|
||||
@@ -119,7 +48,7 @@ impl std::ops::AddAssign for QueryStats {
|
||||
|
||||
// ── QueryHit — one event delivered to query_partition_with's callback ───────
|
||||
|
||||
/// One event from [`KmerPartition::query_partition_with`]'s two-stage query:
|
||||
/// One event from [`QueryPartition::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,
|
||||
@@ -131,9 +60,21 @@ pub enum QueryHit<'a> {
|
||||
Value(&'a [KmerDesc], usize, u32),
|
||||
}
|
||||
|
||||
// ── KmerPartition::query_partition_with ──────────────────────────────────────
|
||||
// ── QueryPartition — extension trait over obikidxcache::IndexCache ─────────
|
||||
|
||||
impl KmerIndex {
|
||||
/// Query-side operations on an already-opened [`IndexCache`] — `IndexCache`
|
||||
/// is a foreign type (`obikidxcache`), so this is an extension trait rather
|
||||
/// than an inherent `impl` (`obikindex`'s original `impl KmerIndex { .. }`
|
||||
/// shape doesn't compile outside `obikindex` itself — the orphan rule).
|
||||
///
|
||||
/// Built on `IndexCache` rather than re-opening each layer's MPHF/matrix
|
||||
/// files per call (the previous, per-call-`QueryLayer::open` design): every
|
||||
/// layer touched by a query is already open, for the whole lifetime of the
|
||||
/// cache, so a query pass costs zero disk I/O beyond the first chunk. This
|
||||
/// also drops the `with_counts` parameter the old design needed: whether a
|
||||
/// layer holds a count or a presence matrix is no longer guessed from an
|
||||
/// index-wide flag, it's read straight off `KmerLayer`'s own variant.
|
||||
pub trait QueryPartition {
|
||||
/// Query a single partition for a pre-deduplicated map of canonical
|
||||
/// k-mers → their occurrences (`seq_idx`, `pos`) in the query batch.
|
||||
///
|
||||
@@ -151,43 +92,54 @@ impl KmerIndex {
|
||||
/// 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>(
|
||||
///
|
||||
/// Infallible: every layer this can touch was already opened (and
|
||||
/// validated) when the cache was built — see `IndexCache::new`'s own
|
||||
/// panic-on-failure contract. A `part_idx` this cache doesn't hold, or
|
||||
/// an empty `kmers` map, both just return default (all-zero) stats.
|
||||
fn query_partition_with<F>(
|
||||
&self,
|
||||
part_idx: usize,
|
||||
kmers: &HashMap<CanonicalKmer, Vec<KmerDesc>>,
|
||||
n_genomes: usize,
|
||||
on_event: F,
|
||||
) -> QueryStats
|
||||
where
|
||||
F: FnMut(QueryHit);
|
||||
}
|
||||
|
||||
impl QueryPartition for IndexCache {
|
||||
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>
|
||||
) -> QueryStats
|
||||
where
|
||||
F: FnMut(QueryHit),
|
||||
{
|
||||
let mut stats = QueryStats::default();
|
||||
|
||||
if kmers.is_empty() {
|
||||
return Ok(stats);
|
||||
return 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<_>>()?;
|
||||
let n_layer = match self.n_layer(part_idx) {
|
||||
Some(n) if n > 0 => n,
|
||||
_ => return stats,
|
||||
};
|
||||
|
||||
// ── 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();
|
||||
(0..n_layer).map(|_| HashMap::new()).collect();
|
||||
|
||||
for (kmer, descs) in kmers {
|
||||
stats.n_unique_kmers += 1;
|
||||
for (layer_idx, layer) in layers.iter().enumerate() {
|
||||
for l in 0..n_layer {
|
||||
stats.n_mphf_calls += 1;
|
||||
if let Some(slot) = layer.find_slot(*kmer) {
|
||||
by_layer[layer_idx].insert(slot, descs);
|
||||
if let Some(slot) = self.find_in_layer(part_idx, l, *kmer) {
|
||||
by_layer[l].insert(slot, descs);
|
||||
on_event(QueryHit::Found(descs));
|
||||
stats.n_hits += 1;
|
||||
break;
|
||||
@@ -196,16 +148,13 @@ impl KmerIndex {
|
||||
}
|
||||
|
||||
// ── 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() {
|
||||
for (l, slots) in by_layer.iter().enumerate() {
|
||||
if slots.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let layer = &layers[layer_idx];
|
||||
let layer = self
|
||||
.get_layer(part_idx, l)
|
||||
.expect("layer within n_layer(part_idx)");
|
||||
let n_cols = layer.n_cols().min(n_genomes);
|
||||
stats.n_columns_scanned += n_cols;
|
||||
|
||||
@@ -221,7 +170,7 @@ impl KmerIndex {
|
||||
}
|
||||
}
|
||||
|
||||
Ok(stats)
|
||||
stats
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,41 @@
|
||||
/// Tracks, per (sequence, s-mer position), whether the k-mer was found in the
|
||||
/// index at all — independent of *which* genome(s) matched. Sized
|
||||
/// `total_smers` (one `bool` per s-mer occurrence in the chunk), **not**
|
||||
/// multiplied by `n_genomes`: this is the O(1)-per-position bookkeeping that
|
||||
/// `kmer_missing` needs (the leftmost-s-mer-of-window membership test), kept
|
||||
/// dense because it's already cheap — the `n_genomes`-scaled data lives in
|
||||
/// the sparse per-genome hit lists built alongside it (see `chunk::process_chunk`).
|
||||
pub(crate) struct SmerIndex {
|
||||
in_index: Vec<bool>, // total_smers
|
||||
offsets: Vec<usize>, // offsets[i]..offsets[i+1] = s-mer range for sequence i
|
||||
}
|
||||
|
||||
impl SmerIndex {
|
||||
pub(crate) fn new(n_kmers_per_seq: &[u32]) -> Self {
|
||||
let mut offsets = Vec::with_capacity(n_kmers_per_seq.len() + 1);
|
||||
let mut total = 0usize;
|
||||
offsets.push(0);
|
||||
for &n in n_kmers_per_seq {
|
||||
total += n as usize;
|
||||
offsets.push(total);
|
||||
}
|
||||
Self {
|
||||
in_index: vec![false; total],
|
||||
offsets,
|
||||
}
|
||||
}
|
||||
|
||||
/// Mark the k-mer at (seq, kmer) as found in the index — independent of
|
||||
/// any particular genome's value. Called once per hit k-mer (stage 1 of
|
||||
/// `query_partition_with`), regardless of how the column-major fetch
|
||||
/// (stage 2) later reports per-genome values.
|
||||
pub(crate) fn mark_found(&mut self, seq: usize, kmer: usize) {
|
||||
let abs = self.offsets[seq] + kmer;
|
||||
self.in_index[abs] = true;
|
||||
}
|
||||
|
||||
#[inline]
|
||||
pub(crate) fn is_in_index(&self, seq: usize, kmer: usize) -> bool {
|
||||
self.in_index[self.offsets[seq] + kmer]
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,127 @@
|
||||
use super::*;
|
||||
use obikrope::Rope;
|
||||
use obikseq::CanonicalKmer;
|
||||
use obiread::record::parse_chunk;
|
||||
|
||||
const K: usize = 11;
|
||||
const M: usize = 5;
|
||||
|
||||
/// Build a `QueryBatch` from raw FASTA text, going through the same
|
||||
/// `Rope` + `parse_chunk` path `process_chunk` uses — avoids hand-building a
|
||||
/// `normalized` `Rope`, which is an implementation detail of `obiread`.
|
||||
///
|
||||
/// `obikseq`'s global K/M params are thread-local under `test-utils` (see
|
||||
/// `obikseq::params`), so setting them here is per-test-thread and does not
|
||||
/// need coordination with other tests.
|
||||
fn batch_from_fasta(fasta: &str, k: usize, n_partitions: usize) -> QueryBatch {
|
||||
obikseq::set_k(k);
|
||||
obikseq::set_m(M);
|
||||
let mut rope = Rope::new(Some("text/fasta"));
|
||||
rope.push(fasta.as_bytes().to_vec());
|
||||
let records = parse_chunk(&rope, k);
|
||||
QueryBatch::from_records(records, k, 6, 0.7, n_partitions)
|
||||
}
|
||||
|
||||
fn total_occurrences(batch: &QueryBatch) -> u64 {
|
||||
batch.n_kmers.iter().map(|&n| n as u64).sum()
|
||||
}
|
||||
|
||||
fn total_unique_kmers(batch: &QueryBatch) -> u64 {
|
||||
batch.by_partition.iter().map(|m| m.len() as u64).sum()
|
||||
}
|
||||
|
||||
// A 60 bp sequence, arbitrary but fixed — no attempt is made to prove it is
|
||||
// free of internal k=11 repeats; the tests below only rely on inequalities
|
||||
// that hold regardless (see each test's comment).
|
||||
const SEQ: &str = "CATTAGCGTACCTGATCAGGTTACAGCTTAGGCATCCAGTTGACCATGACTGGACTTAGC";
|
||||
|
||||
#[test]
|
||||
fn single_sequence_yields_plausible_kmer_counts() {
|
||||
// A single record can still contain internal repeats (SEQ isn't
|
||||
// guaranteed repeat-free at k=11) — this only checks the batch is
|
||||
// internally consistent, not a specific dedup ratio. The cross-record
|
||||
// tests below make the actual, unconditional dedup claims.
|
||||
let fasta = format!(">r1\n{SEQ}\n");
|
||||
let batch = batch_from_fasta(&fasta, K, 1);
|
||||
|
||||
assert_eq!(batch.ids, vec!["r1".to_string()]);
|
||||
let occurrences = total_occurrences(&batch);
|
||||
let unique = total_unique_kmers(&batch);
|
||||
assert!(occurrences > 0, "sequence should yield at least one k-mer");
|
||||
assert!(unique > 0 && unique <= occurrences);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn duplicated_sequence_across_records_deduplicates() {
|
||||
// Two records with byte-identical sequences: every k-mer in record 1
|
||||
// exactly duplicates one in record 0, so unique kmers <= n_kmers[0],
|
||||
// strictly less than the summed occurrences (2 * n_kmers[0]) as long as
|
||||
// the sequence yields at least one k-mer. This holds regardless of
|
||||
// whether SEQ has internal repeats.
|
||||
let fasta = format!(">r1\n{SEQ}\n>r2\n{SEQ}\n");
|
||||
let batch = batch_from_fasta(&fasta, K, 1);
|
||||
|
||||
assert_eq!(batch.ids.len(), 2);
|
||||
let occurrences = total_occurrences(&batch);
|
||||
let unique = total_unique_kmers(&batch);
|
||||
|
||||
assert!(batch.n_kmers[0] > 0);
|
||||
assert_eq!(occurrences, batch.n_kmers[0] as u64 + batch.n_kmers[1] as u64);
|
||||
assert!(
|
||||
unique <= batch.n_kmers[0] as u64,
|
||||
"identical sequences must not produce more unique k-mers than one copy has"
|
||||
);
|
||||
assert!(
|
||||
unique < occurrences,
|
||||
"k-mer-level dedup must collapse at least the cross-record duplication"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn duplicated_sequence_broadcasts_to_both_seq_indices() {
|
||||
// Stronger than the ratio check above: pick any k-mer that hit in both
|
||||
// records and confirm its occurrence list actually references both
|
||||
// seq_idx 0 and seq_idx 1 — this is the specific new capability (dedup
|
||||
// reaching across records/superkmers), not just a smaller unique count.
|
||||
let fasta = format!(">r1\n{SEQ}\n>r2\n{SEQ}\n");
|
||||
let batch = batch_from_fasta(&fasta, K, 1);
|
||||
|
||||
let shared = batch.by_partition[0]
|
||||
.values()
|
||||
.find(|descs| descs.iter().any(|d| d.seq_idx == 0) && descs.iter().any(|d| d.seq_idx == 1));
|
||||
|
||||
assert!(
|
||||
shared.is_some(),
|
||||
"expected at least one k-mer shared between the two identical records"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_records_yield_empty_batch() {
|
||||
let batch = batch_from_fasta("", K, 1);
|
||||
assert!(batch.ids.is_empty());
|
||||
assert_eq!(total_occurrences(&batch), 0);
|
||||
assert_eq!(total_unique_kmers(&batch), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn partition_routing_is_a_pure_function_of_the_kmer() {
|
||||
// With n_partitions=4, every occurrence of a given k-mer must land in
|
||||
// the same partition bucket as every other occurrence of that k-mer
|
||||
// (partition routing is derived from the minimizer, shared by
|
||||
// definition among instances of the same k-mer's containing superkmer
|
||||
// in this test's single-sequence-pair setup).
|
||||
let fasta = format!(">r1\n{SEQ}\n>r2\n{SEQ}\n");
|
||||
let batch = batch_from_fasta(&fasta, K, 4);
|
||||
|
||||
let total_unique: u64 = batch.by_partition.iter().map(|m| m.len() as u64).sum();
|
||||
assert!(total_unique > 0);
|
||||
|
||||
// No k-mer key appears in more than one partition's map.
|
||||
let mut seen: std::collections::HashSet<CanonicalKmer> = std::collections::HashSet::new();
|
||||
for map in &batch.by_partition {
|
||||
for kmer in map.keys() {
|
||||
assert!(seen.insert(*kmer), "k-mer routed to more than one partition");
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,105 @@
|
||||
use super::*;
|
||||
|
||||
// ── sparse_findere_for_genome vs. a dense reference implementation ──────────
|
||||
//
|
||||
// No property-testing crate (proptest/quickcheck) is a workspace dependency
|
||||
// (checked before writing this — not adding one for a single test module,
|
||||
// per this project's dependency-approval rule). A tiny deterministic xorshift
|
||||
// PRNG, std-only, stands in for one.
|
||||
|
||||
/// Faithful reimplementation of the pre-phase-5 dense sliding-window scan —
|
||||
/// the algorithm `sparse_findere_for_genome` replaced — used here only as a
|
||||
/// correctness oracle, not in production code. Operates on one genome's
|
||||
/// hits across possibly many sequences, exactly like the sparse version.
|
||||
fn dense_reference_findere(
|
||||
hits: &[(u32, u32, u32)],
|
||||
seq_lens: &[usize],
|
||||
z: usize,
|
||||
presence: bool,
|
||||
threshold: u32,
|
||||
) -> Vec<(u32, u32, u32)> {
|
||||
let mut by_seq: Vec<Vec<u32>> = seq_lens.iter().map(|&n| vec![0u32; n]).collect();
|
||||
for &(seq, pos, val) in hits {
|
||||
by_seq[seq as usize][pos as usize] = val;
|
||||
}
|
||||
|
||||
let mut confirmed = Vec::new();
|
||||
for (seq_idx, values) in by_seq.iter().enumerate() {
|
||||
let n = values.len();
|
||||
let mut dq: std::collections::VecDeque<(usize, u32)> = std::collections::VecDeque::new();
|
||||
for i in 0..n {
|
||||
let v_i = values[i];
|
||||
while dq.front().map_or(false, |&(f, _)| f + z <= i) {
|
||||
dq.pop_front();
|
||||
}
|
||||
while dq.back().map_or(false, |&(_, v)| v >= v_i) {
|
||||
dq.pop_back();
|
||||
}
|
||||
dq.push_back((i, v_i));
|
||||
if i + 1 >= z {
|
||||
let win_min = dq.front().unwrap().1;
|
||||
if win_min > 0 {
|
||||
let pos_out = (i + 1 - z) as u32;
|
||||
let c = if presence { u32::from(win_min >= threshold) } else { win_min };
|
||||
confirmed.push((seq_idx as u32, pos_out, c));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
confirmed
|
||||
}
|
||||
|
||||
/// Minimal std-only xorshift64 PRNG — deterministic, seedable, no dependency.
|
||||
struct Xorshift64(u64);
|
||||
impl Xorshift64 {
|
||||
fn next(&mut self) -> u64 {
|
||||
self.0 ^= self.0 << 13;
|
||||
self.0 ^= self.0 >> 7;
|
||||
self.0 ^= self.0 << 17;
|
||||
self.0
|
||||
}
|
||||
fn range(&mut self, n: u32) -> u32 {
|
||||
(self.next() % n as u64) as u32
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sparse_findere_matches_dense_reference_on_random_inputs() {
|
||||
let mut rng = Xorshift64(0x5eed_5eed_5eed_5eedu64);
|
||||
|
||||
for case in 0..200 {
|
||||
let n_seqs = 1 + rng.range(4) as usize;
|
||||
let seq_lens: Vec<usize> = (0..n_seqs).map(|_| 1 + rng.range(30) as usize).collect();
|
||||
let z = 1 + rng.range(4) as usize;
|
||||
let presence = rng.range(2) == 0;
|
||||
let threshold = 1 + rng.range(3);
|
||||
|
||||
// Sparse density varies across cases, including edge cases (empty,
|
||||
// fully dense) — deliberately not uniform, to stress both few-hits
|
||||
// and many-overlapping-runs scenarios.
|
||||
let density = rng.range(101);
|
||||
let mut hits: Vec<(u32, u32, u32)> = Vec::new();
|
||||
for (seq_idx, &len) in seq_lens.iter().enumerate() {
|
||||
for pos in 0..len {
|
||||
if rng.range(100) < density {
|
||||
let val = 1 + rng.range(5); // never 0 — matches QueryHit::Value's invariant
|
||||
hits.push((seq_idx as u32, pos as u32, val));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let mut sparse_input = hits.clone();
|
||||
let (mut sparse_result, _, _) =
|
||||
sparse_findere_for_genome(&mut sparse_input, z, presence, threshold);
|
||||
let mut dense_result = dense_reference_findere(&hits, &seq_lens, z, presence, threshold);
|
||||
|
||||
sparse_result.sort_unstable();
|
||||
dense_result.sort_unstable();
|
||||
|
||||
assert_eq!(
|
||||
sparse_result, dense_result,
|
||||
"case {case}: n_seqs={n_seqs} seq_lens={seq_lens:?} z={z} presence={presence} \
|
||||
threshold={threshold} density={density} hits={hits:?}"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,5 +1,7 @@
|
||||
use super::*;
|
||||
use obikindex::IndexConfig;
|
||||
use obikidxcache::index_cache::IndexCache;
|
||||
use obikindex::{IndexConfig, KmerIndex};
|
||||
use std::sync::Arc;
|
||||
|
||||
// ── QueryStats::AddAssign ───────────────────────────────────────────────────
|
||||
|
||||
@@ -39,33 +41,39 @@ fn query_stats_default_is_zero() {
|
||||
// ── 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.
|
||||
/// `index/` subdirectory under any partition — `IndexCache::new` still opens
|
||||
/// fine on it (zero layers found), and `query_partition_with` must recognise
|
||||
/// this and return default (all-zero) stats rather than panicking, exactly
|
||||
/// like an empty `kmers` map.
|
||||
#[test]
|
||||
fn query_partition_with_missing_index_dir_returns_default_stats() {
|
||||
let tmp = tempfile::tempdir().expect("tempdir");
|
||||
// Same k/m pair as `batch`'s tests (11, 5) — `obikseq`'s global k/m
|
||||
// params are process-wide atomics in test builds (see
|
||||
// `obikseq::params`'s own doc comment), shared by every test in this
|
||||
// crate's single test binary regardless of thread; a different pair
|
||||
// here would race with `batch`'s tests running concurrently.
|
||||
let config = IndexConfig {
|
||||
kmer_size: 21,
|
||||
minimizer_size: 9,
|
||||
kmer_size: 11,
|
||||
minimizer_size: 5,
|
||||
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 cache = IndexCache::new(Arc::new(index), Some(vec![0]));
|
||||
|
||||
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.
|
||||
// before ever attempting an MPHF lookup, since the partition has zero
|
||||
// layers.
|
||||
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");
|
||||
let stats = cache.query_partition_with(0, &kmers, 1, |_event| {
|
||||
panic!("on_event must not be called: no index was built");
|
||||
});
|
||||
|
||||
assert_eq!(stats, QueryStats::default());
|
||||
}
|
||||
@@ -73,22 +81,23 @@ 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");
|
||||
// Same k/m pair as `batch`'s tests — see the comment in the previous
|
||||
// test for why this must match across every test in this crate.
|
||||
let config = IndexConfig {
|
||||
kmer_size: 21,
|
||||
minimizer_size: 9,
|
||||
kmer_size: 11,
|
||||
minimizer_size: 5,
|
||||
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 cache = IndexCache::new(Arc::new(index), Some(vec![0]));
|
||||
|
||||
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");
|
||||
let stats = cache.query_partition_with(0, &kmers, 1, |_event| {
|
||||
panic!("on_event must not be called on an empty kmer map");
|
||||
});
|
||||
|
||||
assert_eq!(stats, QueryStats::default());
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user