Push smluomvxpptv #59
@@ -7,8 +7,9 @@ use std::time::Instant;
|
|||||||
|
|
||||||
use clap::Args;
|
use clap::Args;
|
||||||
use obikindex::KmerIndex;
|
use obikindex::KmerIndex;
|
||||||
|
use obikpartitionner::{KmerDesc, QueryStats};
|
||||||
use obikrope::Rope;
|
use obikrope::Rope;
|
||||||
use obikseq::RoutableSuperKmer;
|
use obikseq::CanonicalKmer;
|
||||||
use obilayeredmap::IndexMode;
|
use obilayeredmap::IndexMode;
|
||||||
use obipipeline::{Throttled, ThrottleGuard, throttle};
|
use obipipeline::{Throttled, ThrottleGuard, throttle};
|
||||||
use obiread::chunk::read_sequence_chunks_sized;
|
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 ────────────────────────────────────────────────────────────────
|
// ── 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 {
|
pub struct QueryBatch {
|
||||||
/// Sequence ids in batch order.
|
/// Sequence ids in batch order.
|
||||||
pub ids: Vec<String>,
|
pub ids: Vec<String>,
|
||||||
@@ -114,30 +114,40 @@ pub struct QueryBatch {
|
|||||||
pub seqs: Vec<Vec<u8>>,
|
pub seqs: Vec<Vec<u8>>,
|
||||||
/// Total kmer count per sequence (used for `--detail` coverage allocation).
|
/// Total kmer count per sequence (used for `--detail` coverage allocation).
|
||||||
pub n_kmers: Vec<u32>,
|
pub n_kmers: Vec<u32>,
|
||||||
/// Deduplicated superkmer map.
|
/// Deduplicated k-mer occurrences, one map per partition.
|
||||||
pub map: HashMap<RoutableSuperKmer, Vec<SKDesc>>,
|
pub by_partition: Vec<HashMap<CanonicalKmer, Vec<KmerDesc>>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl QueryBatch {
|
impl QueryBatch {
|
||||||
/// Build a batch from a vec of parsed sequence records.
|
/// Build a batch from a vec of parsed sequence records, deduplicating
|
||||||
pub fn from_records(records: Vec<SeqRecord>, k: usize, level_max: usize, theta: f64) -> Self {
|
/// 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 ids = Vec::with_capacity(records.len());
|
||||||
let mut seqs = Vec::with_capacity(records.len());
|
let mut seqs = Vec::with_capacity(records.len());
|
||||||
let mut n_kmers = Vec::with_capacity(records.len());
|
let mut n_kmers = Vec::with_capacity(records.len());
|
||||||
// Upper-bound estimate: at most one superkmer per k bases.
|
let mask = (n_partitions as u64) - 1;
|
||||||
// Avoids repeated rehash on large chunks.
|
let mut by_partition: Vec<HashMap<CanonicalKmer, Vec<KmerDesc>>> =
|
||||||
let cap = records.iter().map(|r| r.normalized.len()).sum::<usize>() / k.max(1);
|
(0..n_partitions).map(|_| HashMap::new()).collect();
|
||||||
let mut map: HashMap<RoutableSuperKmer, Vec<SKDesc>> = HashMap::with_capacity(cap);
|
|
||||||
|
|
||||||
for (seq_idx, record) in records.into_iter().enumerate() {
|
for (seq_idx, record) in records.into_iter().enumerate() {
|
||||||
let mut kmer_offset = 0u32;
|
let mut kmer_offset = 0u32;
|
||||||
|
|
||||||
for rsk in SuperKmerIter::new(&record.normalized, k, level_max, theta) {
|
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;
|
let n = (rsk.seql() - k + 1) as u32;
|
||||||
map.entry(rsk).or_default().push(SKDesc {
|
|
||||||
seq_idx: seq_idx as u32,
|
|
||||||
kmer_offset,
|
|
||||||
});
|
|
||||||
kmer_offset += n;
|
kmer_offset += n;
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -150,20 +160,9 @@ impl QueryBatch {
|
|||||||
ids,
|
ids,
|
||||||
seqs,
|
seqs,
|
||||||
n_kmers,
|
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 ───────────────────────
|
// ── KmerResults — allocation-free ragged result matrix ───────────────────────
|
||||||
@@ -260,33 +259,36 @@ fn process_chunk(
|
|||||||
return Vec::new();
|
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();
|
let n_seqs = batch.ids.len();
|
||||||
|
|
||||||
// Flat result matrix — one allocation for the whole chunk.
|
// Flat result matrix — one allocation for the whole chunk.
|
||||||
let mut results = KmerResults::new(&batch.n_kmers, n_genomes);
|
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() {
|
for (part_idx, kmers) in batch.by_partition.iter().enumerate() {
|
||||||
if part_sks.is_empty() {
|
if kmers.is_empty() {
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
|
|
||||||
idx.partition()
|
let stats = idx.partition()
|
||||||
.query_partition_with(
|
.query_partition_with(
|
||||||
part_idx,
|
part_idx,
|
||||||
part_sks,
|
kmers,
|
||||||
k,
|
|
||||||
n_genomes,
|
n_genomes,
|
||||||
with_counts,
|
with_counts,
|
||||||
|sk_idx, kmer_idx, row| {
|
|descs, row| {
|
||||||
let rsk = part_sks[sk_idx];
|
|
||||||
let descs = batch.map.get(rsk).expect("rsk must be in map");
|
|
||||||
for desc in descs {
|
for desc in descs {
|
||||||
results.set(
|
results.set(
|
||||||
desc.seq_idx as usize,
|
desc.seq_idx as usize,
|
||||||
desc.kmer_offset as usize + kmer_idx,
|
desc.pos as usize,
|
||||||
row,
|
row,
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
@@ -296,8 +298,17 @@ fn process_chunk(
|
|||||||
eprintln!("query error on partition {part_idx}: {e}");
|
eprintln!("query error on partition {part_idx}: {e}");
|
||||||
std::process::exit(1);
|
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.
|
// 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)
|
// 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");
|
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 filter::{GroupQuorumFilter, KmerFilter, passes_all};
|
||||||
pub use merge_layer::MergeMode;
|
pub use merge_layer::MergeMode;
|
||||||
pub use partition::{KmerPartition, KmerSpectrum, PARTITIONS_SUBDIR};
|
pub use partition::{KmerPartition, KmerSpectrum, PARTITIONS_SUBDIR};
|
||||||
|
pub use query_layer::{KmerDesc, QueryStats};
|
||||||
pub use select_layer::{AggOp, OutputCol};
|
pub use select_layer::{AggOp, OutputCol};
|
||||||
|
|||||||
@@ -1,7 +1,8 @@
|
|||||||
|
use std::collections::HashMap;
|
||||||
use std::path::Path;
|
use std::path::Path;
|
||||||
|
|
||||||
use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix};
|
use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix};
|
||||||
use obikseq::{CanonicalKmer, RoutableSuperKmer};
|
use obikseq::CanonicalKmer;
|
||||||
use obiskio::{SKError, SKResult};
|
use obiskio::{SKError, SKResult};
|
||||||
use obilayeredmap::{IndexMode, MphfLayer, OLMError};
|
use obilayeredmap::{IndexMode, MphfLayer, OLMError};
|
||||||
use obilayeredmap::meta::PartitionMeta;
|
use obilayeredmap::meta::PartitionMeta;
|
||||||
@@ -44,53 +45,86 @@ impl QueryLayer {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Write per-genome values into `buf` if `kmer` is indexed; returns true on hit.
|
/// MPHF lookup only — no matrix access. `Some(slot)` on hit.
|
||||||
fn find_into(&self, kmer: CanonicalKmer, n_genomes: usize, buf: &mut [u32]) -> bool {
|
fn find_slot(&self, kmer: CanonicalKmer) -> Option<usize> {
|
||||||
match self {
|
match self {
|
||||||
QueryLayer::Presence(mphf, mat) => {
|
QueryLayer::Presence(mphf, _) | QueryLayer::Count(mphf, _) => mphf.find(kmer),
|
||||||
if let Some(slot) = mphf.find(kmer) {
|
}
|
||||||
mat.fill_row(slot, &mut buf[..n_genomes]);
|
}
|
||||||
true
|
|
||||||
} else {
|
/// Write per-genome values for `slot` into `buf`. `slot` must come from
|
||||||
false
|
/// [`find_slot`] on this same layer.
|
||||||
}
|
fn fill_row(&self, slot: usize, n_genomes: usize, buf: &mut [u32]) {
|
||||||
}
|
match self {
|
||||||
QueryLayer::Count(mphf, mat) => {
|
QueryLayer::Presence(_, mat) => mat.fill_row(slot, &mut buf[..n_genomes]),
|
||||||
if let Some(slot) = mphf.find(kmer) {
|
QueryLayer::Count(_, mat) => mat.fill_row(slot, &mut buf[..n_genomes]),
|
||||||
mat.fill_row(slot, &mut buf[..n_genomes]);
|
|
||||||
true
|
|
||||||
} else {
|
|
||||||
false
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// ── 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 {
|
impl KmerPartition {
|
||||||
/// Query a single partition, calling `on_hit(sk_idx, kmer_idx, row)` for
|
/// Query a single partition for a pre-deduplicated map of canonical
|
||||||
/// every found k-mer without allocating intermediate result vectors.
|
/// 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>(
|
pub fn query_partition_with<F>(
|
||||||
&self,
|
&self,
|
||||||
part_idx: usize,
|
part_idx: usize,
|
||||||
superkmers: &[&RoutableSuperKmer],
|
kmers: &HashMap<CanonicalKmer, Vec<KmerDesc>>,
|
||||||
_k: usize,
|
|
||||||
n_genomes: usize,
|
n_genomes: usize,
|
||||||
with_counts: bool,
|
with_counts: bool,
|
||||||
mut on_hit: F,
|
mut on_hit: F,
|
||||||
) -> SKResult<()>
|
) -> SKResult<QueryStats>
|
||||||
where
|
where
|
||||||
F: FnMut(usize, usize, &[u32]),
|
F: FnMut(&[KmerDesc], &[u32]),
|
||||||
{
|
{
|
||||||
if superkmers.is_empty() {
|
let mut stats = QueryStats::default();
|
||||||
return Ok(());
|
|
||||||
|
if kmers.is_empty() {
|
||||||
|
return Ok(stats);
|
||||||
}
|
}
|
||||||
|
|
||||||
let index_dir = self.part_dir(part_idx).join(INDEX_SUBDIR);
|
let index_dir = self.part_dir(part_idx).join(INDEX_SUBDIR);
|
||||||
if !index_dir.exists() {
|
if !index_dir.exists() {
|
||||||
return Ok(());
|
return Ok(stats);
|
||||||
}
|
}
|
||||||
|
|
||||||
let meta = PartitionMeta::load(&index_dir).map_err(olm_to_sk)?;
|
let meta = PartitionMeta::load(&index_dir).map_err(olm_to_sk)?;
|
||||||
@@ -100,67 +134,24 @@ impl KmerPartition {
|
|||||||
|
|
||||||
let mut buf = vec![0u32; n_genomes];
|
let mut buf = vec![0u32; n_genomes];
|
||||||
|
|
||||||
for (sk_idx, rsk) in superkmers.iter().enumerate() {
|
for (kmer, descs) in kmers {
|
||||||
for (kmer_idx, kmer) in rsk.superkmer().iter_canonical_kmers().enumerate() {
|
stats.n_unique_kmers += 1;
|
||||||
for layer in &layers {
|
for layer in &layers {
|
||||||
if layer.find_into(kmer, n_genomes, &mut buf) {
|
stats.n_mphf_calls += 1;
|
||||||
on_hit(sk_idx, kmer_idx, &buf);
|
if let Some(slot) = layer.find_slot(*kmer) {
|
||||||
buf.fill(0);
|
layer.fill_row(slot, n_genomes, &mut buf);
|
||||||
break;
|
on_hit(descs, &buf);
|
||||||
}
|
buf.fill(0);
|
||||||
|
stats.n_hits += 1;
|
||||||
|
break;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
Ok(())
|
Ok(stats)
|
||||||
}
|
|
||||||
|
|
||||||
/// 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())
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[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