refactor(query): deduplicate k-mers upfront and split MPHF lookup
Refactor `QueryBatch` construction to perform canonical k-mer deduplication and partition routing during initialization, eliminating post-batch splitting. Split the `QueryLayer` MPHF lookup into separate `find_slot` and `fill_row` methods, updating callbacks to pass occurrence descriptors instead of indices. Introduce `QueryStats` for tracking MPHF calls and dereplication ratios, and add comprehensive unit tests for batch construction, stats arithmetic, and safe partition handling. Expose new query-layer types in the public API.
This commit is contained in:
Eric Coissac
@@ -7,8 +7,9 @@ use std::time::Instant;
|
||||
|
||||
use clap::Args;
|
||||
use obikindex::KmerIndex;
|
||||
use obikpartitionner::{KmerDesc, QueryStats};
|
||||
use obikrope::Rope;
|
||||
use obikseq::RoutableSuperKmer;
|
||||
use obikseq::CanonicalKmer;
|
||||
use obilayeredmap::IndexMode;
|
||||
use obipipeline::{Throttled, ThrottleGuard, throttle};
|
||||
use obiread::chunk::read_sequence_chunks_sized;
|
||||
@@ -94,19 +95,18 @@ impl QueryArgs {
|
||||
}
|
||||
}
|
||||
|
||||
// ── SKDesc — one occurrence of a superkmer in the batch ───────────────────────
|
||||
|
||||
/// Describes one occurrence of a superkmer in the query batch.
|
||||
pub struct SKDesc {
|
||||
/// Index of the source sequence within the batch.
|
||||
pub seq_idx: u32,
|
||||
/// Kmer offset of the first kmer of this superkmer within its sequence.
|
||||
pub kmer_offset: u32,
|
||||
}
|
||||
|
||||
// ── QueryBatch ────────────────────────────────────────────────────────────────
|
||||
|
||||
/// A batch of query sequences with their superkmers deduplicated.
|
||||
/// 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 struct QueryBatch {
|
||||
/// Sequence ids in batch order.
|
||||
pub ids: Vec<String>,
|
||||
@@ -114,30 +114,40 @@ pub struct QueryBatch {
|
||||
pub seqs: Vec<Vec<u8>>,
|
||||
/// Total kmer count per sequence (used for `--detail` coverage allocation).
|
||||
pub n_kmers: Vec<u32>,
|
||||
/// Deduplicated superkmer map.
|
||||
pub map: HashMap<RoutableSuperKmer, Vec<SKDesc>>,
|
||||
/// Deduplicated k-mer occurrences, one map per partition.
|
||||
pub by_partition: Vec<HashMap<CanonicalKmer, Vec<KmerDesc>>>,
|
||||
}
|
||||
|
||||
impl QueryBatch {
|
||||
/// Build a batch from a vec of parsed sequence records.
|
||||
pub fn from_records(records: Vec<SeqRecord>, k: usize, level_max: usize, theta: f64) -> Self {
|
||||
/// Build a batch from a vec of parsed sequence records, deduplicating
|
||||
/// k-mers and routing them to partitions in the same pass.
|
||||
pub 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());
|
||||
// Upper-bound estimate: at most one superkmer per k bases.
|
||||
// Avoids repeated rehash on large chunks.
|
||||
let cap = records.iter().map(|r| r.normalized.len()).sum::<usize>() / k.max(1);
|
||||
let mut map: HashMap<RoutableSuperKmer, Vec<SKDesc>> = HashMap::with_capacity(cap);
|
||||
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 n = (rsk.seql() - k + 1) as u32;
|
||||
map.entry(rsk).or_default().push(SKDesc {
|
||||
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,
|
||||
kmer_offset,
|
||||
pos: kmer_offset + j as u32,
|
||||
});
|
||||
}
|
||||
let n = (rsk.seql() - k + 1) as u32;
|
||||
kmer_offset += n;
|
||||
}
|
||||
|
||||
@@ -150,20 +160,9 @@ impl QueryBatch {
|
||||
ids,
|
||||
seqs,
|
||||
n_kmers,
|
||||
map,
|
||||
by_partition,
|
||||
}
|
||||
}
|
||||
|
||||
/// Split the superkmer map by partition index.
|
||||
pub fn split_by_partition(&self, n_partitions: usize) -> Vec<Vec<&RoutableSuperKmer>> {
|
||||
let mask = (n_partitions as u64) - 1;
|
||||
let mut by_part: Vec<Vec<&RoutableSuperKmer>> = vec![Vec::new(); n_partitions];
|
||||
for rsk in self.map.keys() {
|
||||
let part = (rsk.minimizer().seq_hash() & mask) as usize;
|
||||
by_part[part].push(rsk);
|
||||
}
|
||||
by_part
|
||||
}
|
||||
}
|
||||
|
||||
// ── KmerResults — allocation-free ragged result matrix ───────────────────────
|
||||
@@ -260,33 +259,36 @@ fn process_chunk(
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
let batch = QueryBatch::from_records(records, k, 6, 0.7);
|
||||
let batch = QueryBatch::from_records(records, k, 6, 0.7, n_partitions);
|
||||
let n_seqs = batch.ids.len();
|
||||
|
||||
// Flat result matrix — one allocation for the whole chunk.
|
||||
let mut results = KmerResults::new(&batch.n_kmers, n_genomes);
|
||||
|
||||
let by_part = batch.split_by_partition(n_partitions);
|
||||
// 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 (see query.md,
|
||||
// Future work point 5). 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, part_sks) in by_part.iter().enumerate() {
|
||||
if part_sks.is_empty() {
|
||||
for (part_idx, kmers) in batch.by_partition.iter().enumerate() {
|
||||
if kmers.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
idx.partition()
|
||||
let stats = idx.partition()
|
||||
.query_partition_with(
|
||||
part_idx,
|
||||
part_sks,
|
||||
k,
|
||||
kmers,
|
||||
n_genomes,
|
||||
with_counts,
|
||||
|sk_idx, kmer_idx, row| {
|
||||
let rsk = part_sks[sk_idx];
|
||||
let descs = batch.map.get(rsk).expect("rsk must be in map");
|
||||
|descs, row| {
|
||||
for desc in descs {
|
||||
results.set(
|
||||
desc.seq_idx as usize,
|
||||
desc.kmer_offset as usize + kmer_idx,
|
||||
desc.pos as usize,
|
||||
row,
|
||||
);
|
||||
}
|
||||
@@ -296,8 +298,17 @@ fn process_chunk(
|
||||
eprintln!("query error on partition {part_idx}: {e}");
|
||||
std::process::exit(1);
|
||||
});
|
||||
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,
|
||||
"k-mer dedup"
|
||||
);
|
||||
|
||||
// Sliding window minimum — one reusable buffer and one deque per batch.
|
||||
//
|
||||
// win_min[pos * n_genomes + g] = min count across the z-window [pos, pos+z)
|
||||
@@ -686,3 +697,7 @@ fn emit_batch(
|
||||
let _ = out.write_all(b"\n");
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "tests/query.rs"]
|
||||
mod tests;
|
||||
|
||||
@@ -0,0 +1,124 @@
|
||||
use super::*;
|
||||
|
||||
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");
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -14,4 +14,5 @@ mod select_layer;
|
||||
pub use filter::{GroupQuorumFilter, KmerFilter, passes_all};
|
||||
pub use merge_layer::MergeMode;
|
||||
pub use partition::{KmerPartition, KmerSpectrum, PARTITIONS_SUBDIR};
|
||||
pub use query_layer::{KmerDesc, QueryStats};
|
||||
pub use select_layer::{AggOp, OutputCol};
|
||||
|
||||
@@ -1,7 +1,8 @@
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
|
||||
use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix};
|
||||
use obikseq::{CanonicalKmer, RoutableSuperKmer};
|
||||
use obikseq::CanonicalKmer;
|
||||
use obiskio::{SKError, SKResult};
|
||||
use obilayeredmap::{IndexMode, MphfLayer, OLMError};
|
||||
use obilayeredmap::meta::PartitionMeta;
|
||||
@@ -44,53 +45,86 @@ impl QueryLayer {
|
||||
}
|
||||
}
|
||||
|
||||
/// Write per-genome values into `buf` if `kmer` is indexed; returns true on hit.
|
||||
fn find_into(&self, kmer: CanonicalKmer, n_genomes: usize, buf: &mut [u32]) -> bool {
|
||||
/// MPHF lookup only — no matrix access. `Some(slot)` on hit.
|
||||
fn find_slot(&self, kmer: CanonicalKmer) -> Option<usize> {
|
||||
match self {
|
||||
QueryLayer::Presence(mphf, mat) => {
|
||||
if let Some(slot) = mphf.find(kmer) {
|
||||
mat.fill_row(slot, &mut buf[..n_genomes]);
|
||||
true
|
||||
} else {
|
||||
false
|
||||
}
|
||||
}
|
||||
QueryLayer::Count(mphf, mat) => {
|
||||
if let Some(slot) = mphf.find(kmer) {
|
||||
mat.fill_row(slot, &mut buf[..n_genomes]);
|
||||
true
|
||||
} else {
|
||||
false
|
||||
QueryLayer::Presence(mphf, _) | QueryLayer::Count(mphf, _) => mphf.find(kmer),
|
||||
}
|
||||
}
|
||||
|
||||
/// Write per-genome values for `slot` into `buf`. `slot` must come from
|
||||
/// [`find_slot`] on this same layer.
|
||||
fn fill_row(&self, slot: usize, n_genomes: usize, buf: &mut [u32]) {
|
||||
match self {
|
||||
QueryLayer::Presence(_, mat) => mat.fill_row(slot, &mut buf[..n_genomes]),
|
||||
QueryLayer::Count(_, mat) => mat.fill_row(slot, &mut buf[..n_genomes]),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── KmerPartition::query_partition* ──────────────────────────────────────────
|
||||
// ── 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 logging in `obikmer::cmd::query` (occurrences vs. unique
|
||||
/// k-mers is the whole justification for k-mer-level dereplication).
|
||||
#[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,
|
||||
}
|
||||
|
||||
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;
|
||||
}
|
||||
}
|
||||
|
||||
// ── KmerPartition::query_partition_with ──────────────────────────────────────
|
||||
|
||||
impl KmerPartition {
|
||||
/// Query a single partition, calling `on_hit(sk_idx, kmer_idx, row)` for
|
||||
/// every found k-mer without allocating intermediate result vectors.
|
||||
/// Query a single partition for a pre-deduplicated map of canonical
|
||||
/// k-mers → their occurrences (`seq_idx`, `pos`) in the query batch.
|
||||
///
|
||||
/// Each unique k-mer triggers at most one MPHF lookup per layer (stopping
|
||||
/// at the first hit) and, on hit, one matrix row fetch — regardless of
|
||||
/// how many times that k-mer occurs across the batch. `on_hit(descs, row)`
|
||||
/// is called once per hit, with every occurrence to broadcast the row to.
|
||||
pub fn query_partition_with<F>(
|
||||
&self,
|
||||
part_idx: usize,
|
||||
superkmers: &[&RoutableSuperKmer],
|
||||
_k: usize,
|
||||
kmers: &HashMap<CanonicalKmer, Vec<KmerDesc>>,
|
||||
n_genomes: usize,
|
||||
with_counts: bool,
|
||||
mut on_hit: F,
|
||||
) -> SKResult<()>
|
||||
) -> SKResult<QueryStats>
|
||||
where
|
||||
F: FnMut(usize, usize, &[u32]),
|
||||
F: FnMut(&[KmerDesc], &[u32]),
|
||||
{
|
||||
if superkmers.is_empty() {
|
||||
return Ok(());
|
||||
let mut stats = QueryStats::default();
|
||||
|
||||
if kmers.is_empty() {
|
||||
return Ok(stats);
|
||||
}
|
||||
|
||||
let index_dir = self.part_dir(part_idx).join(INDEX_SUBDIR);
|
||||
if !index_dir.exists() {
|
||||
return Ok(());
|
||||
return Ok(stats);
|
||||
}
|
||||
|
||||
let meta = PartitionMeta::load(&index_dir).map_err(olm_to_sk)?;
|
||||
@@ -100,67 +134,24 @@ impl KmerPartition {
|
||||
|
||||
let mut buf = vec![0u32; n_genomes];
|
||||
|
||||
for (sk_idx, rsk) in superkmers.iter().enumerate() {
|
||||
for (kmer_idx, kmer) in rsk.superkmer().iter_canonical_kmers().enumerate() {
|
||||
for (kmer, descs) in kmers {
|
||||
stats.n_unique_kmers += 1;
|
||||
for layer in &layers {
|
||||
if layer.find_into(kmer, n_genomes, &mut buf) {
|
||||
on_hit(sk_idx, kmer_idx, &buf);
|
||||
stats.n_mphf_calls += 1;
|
||||
if let Some(slot) = layer.find_slot(*kmer) {
|
||||
layer.fill_row(slot, n_genomes, &mut buf);
|
||||
on_hit(descs, &buf);
|
||||
buf.fill(0);
|
||||
stats.n_hits += 1;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Query a single partition for a slice of super-kmers, returning per-kmer rows.
|
||||
/// Prefer [`query_partition_with`] to avoid per-kmer heap allocations.
|
||||
pub fn query_partition(
|
||||
&self,
|
||||
part_idx: usize,
|
||||
superkmers: &[&RoutableSuperKmer],
|
||||
_k: usize,
|
||||
n_genomes: usize,
|
||||
with_counts: bool,
|
||||
) -> SKResult<Vec<Vec<Option<Box<[u32]>>>>> {
|
||||
if superkmers.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let index_dir = self.part_dir(part_idx).join(INDEX_SUBDIR);
|
||||
|
||||
if !index_dir.exists() {
|
||||
return Ok(superkmers
|
||||
.iter()
|
||||
.map(|rsk| vec![None; rsk.seql()])
|
||||
.collect());
|
||||
}
|
||||
|
||||
let meta = PartitionMeta::load(&index_dir).map_err(olm_to_sk)?;
|
||||
let layers: Vec<QueryLayer> = (0..meta.n_layers)
|
||||
.map(|i| QueryLayer::open(&index_dir.join(format!("layer_{i}")), with_counts, &meta.mode))
|
||||
.collect::<SKResult<_>>()?;
|
||||
|
||||
let mut buf = vec![0u32; n_genomes];
|
||||
Ok(superkmers
|
||||
.iter()
|
||||
.map(|rsk| {
|
||||
rsk.superkmer()
|
||||
.iter_canonical_kmers()
|
||||
.map(|kmer| {
|
||||
for layer in &layers {
|
||||
if layer.find_into(kmer, n_genomes, &mut buf) {
|
||||
let row: Box<[u32]> = buf[..n_genomes].into();
|
||||
buf.fill(0);
|
||||
return Some(row);
|
||||
}
|
||||
}
|
||||
None
|
||||
})
|
||||
.collect()
|
||||
})
|
||||
.collect())
|
||||
Ok(stats)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "tests/query_layer.rs"]
|
||||
mod tests;
|
||||
|
||||
@@ -0,0 +1,65 @@
|
||||
use super::*;
|
||||
|
||||
// ── 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 };
|
||||
total += QueryStats { n_unique_kmers: 1, n_mphf_calls: 4, n_hits: 1 };
|
||||
assert_eq!(total.n_unique_kmers, 4);
|
||||
assert_eq!(total.n_mphf_calls, 9);
|
||||
assert_eq!(total.n_hits, 3);
|
||||
}
|
||||
|
||||
#[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);
|
||||
}
|
||||
|
||||
// ── 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 partition = KmerPartition::create(tmp.path().join("idx"), 2, 21, 9, false)
|
||||
.expect("create partition");
|
||||
|
||||
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 = partition
|
||||
.query_partition_with(0, &kmers, 1, false, |_descs, _row| {
|
||||
panic!("on_hit must not be called: no index was built");
|
||||
})
|
||||
.expect("query_partition_with should not error on a missing index dir");
|
||||
|
||||
assert_eq!(stats.n_unique_kmers, 0);
|
||||
assert_eq!(stats.n_mphf_calls, 0);
|
||||
assert_eq!(stats.n_hits, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn query_partition_with_empty_kmers_is_a_noop() {
|
||||
let tmp = tempfile::tempdir().expect("tempdir");
|
||||
let partition = KmerPartition::create(tmp.path().join("idx"), 2, 21, 9, false)
|
||||
.expect("create partition");
|
||||
|
||||
let kmers: HashMap<CanonicalKmer, Vec<KmerDesc>> = HashMap::new();
|
||||
let stats = partition
|
||||
.query_partition_with(0, &kmers, 1, false, |_, _| {
|
||||
panic!("on_hit 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