Push zpwxxpnpktps #67
@@ -1,851 +0,0 @@
|
|||||||
use std::collections::{HashMap, VecDeque};
|
|
||||||
use std::io::{self, BufWriter, Write};
|
|
||||||
use std::path::PathBuf;
|
|
||||||
use std::sync::Arc;
|
|
||||||
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
|
|
||||||
use std::time::Instant;
|
|
||||||
|
|
||||||
use clap::Args;
|
|
||||||
use obikindex::KmerIndex;
|
|
||||||
use obikpartitionner::{KmerDesc, QueryHit, QueryStats};
|
|
||||||
use obikrope::Rope;
|
|
||||||
use obikseq::CanonicalKmer;
|
|
||||||
use obilayeredmap::IndexMode;
|
|
||||||
use obipipeline::{Throttled, ThrottleGuard, throttle};
|
|
||||||
use obiread::chunk::read_sequence_chunks_sized;
|
|
||||||
use obiread::record::{SeqRecord, parse_chunk};
|
|
||||||
use obiskbuilder::SuperKmerIter;
|
|
||||||
use obisys::{Reporter, Stage, available_memory_bytes, spinner};
|
|
||||||
use tracing::{debug, info};
|
|
||||||
|
|
||||||
// ── Pipeline data ─────────────────────────────────────────────────────────────
|
|
||||||
|
|
||||||
enum QueryData {
|
|
||||||
Path(Throttled<PathBuf>),
|
|
||||||
Chunk(Rope),
|
|
||||||
Output(Vec<u8>),
|
|
||||||
}
|
|
||||||
|
|
||||||
// SAFETY: Rope contains Cell<u8> which is !Sync, but pipeline items are owned
|
|
||||||
// exclusively through channels — no item is ever shared across threads.
|
|
||||||
unsafe impl Send for QueryData {}
|
|
||||||
unsafe impl Sync for QueryData {}
|
|
||||||
|
|
||||||
// ── CLI ───────────────────────────────────────────────────────────────────────
|
|
||||||
|
|
||||||
#[derive(Args)]
|
|
||||||
pub struct QueryArgs {
|
|
||||||
/// Index directory
|
|
||||||
pub index: PathBuf,
|
|
||||||
|
|
||||||
/// Input sequences (FASTA/FASTQ, optionally gzip-compressed)
|
|
||||||
#[arg(num_args = 1..)]
|
|
||||||
pub inputs: Vec<String>,
|
|
||||||
|
|
||||||
/// Report per-position coverage vectors per genome (adds "coverage" to JSON)
|
|
||||||
#[arg(long)]
|
|
||||||
pub detail: bool,
|
|
||||||
|
|
||||||
/// Enable 1-mismatch approximate matching
|
|
||||||
#[arg(long)]
|
|
||||||
pub mismatch: bool,
|
|
||||||
|
|
||||||
/// Count k-mers absent from the index (adds kmer_missing annotation)
|
|
||||||
#[arg(long)]
|
|
||||||
pub count_missing: bool,
|
|
||||||
|
|
||||||
/// Report per-genome presence (0/1) instead of raw counts
|
|
||||||
#[arg(long)]
|
|
||||||
pub force_presence: bool,
|
|
||||||
|
|
||||||
/// Minimum accumulated match count to declare a genome present (implies --force-presence)
|
|
||||||
#[arg(long, default_value_t = 1)]
|
|
||||||
pub presence_threshold: u32,
|
|
||||||
|
|
||||||
/// Override the Findere z parameter from index metadata
|
|
||||||
#[arg(short = 'z', long)]
|
|
||||||
pub findere_z: Option<usize>,
|
|
||||||
|
|
||||||
/// Number of worker threads
|
|
||||||
#[arg(
|
|
||||||
short = 'T',
|
|
||||||
long,
|
|
||||||
default_value_t = obisys::effective_parallelism()
|
|
||||||
)]
|
|
||||||
pub threads: usize,
|
|
||||||
|
|
||||||
/// I/O chunk size in MiB (default: auto-sized from available RAM and thread count)
|
|
||||||
#[arg(long)]
|
|
||||||
pub chunk_size: Option<usize>,
|
|
||||||
|
|
||||||
/// Maximum number of input files open simultaneously.
|
|
||||||
/// Defaults to threads/4 (minimum 1). Keep below the number of workers
|
|
||||||
/// to ensure CPU workers are always available for the transform stage.
|
|
||||||
#[arg(long)]
|
|
||||||
pub max_open_files: Option<usize>,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl QueryArgs {
|
|
||||||
pub fn effective_max_open(&self) -> usize {
|
|
||||||
self.max_open_files
|
|
||||||
.unwrap_or_else(|| (self.threads / 4).max(1))
|
|
||||||
.max(1)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// ── QueryBatch ────────────────────────────────────────────────────────────────
|
|
||||||
|
|
||||||
/// 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>,
|
|
||||||
/// Raw sequence bytes (for output), in batch order.
|
|
||||||
pub seqs: Vec<Vec<u8>>,
|
|
||||||
/// Total kmer count per sequence (used for `--detail` coverage allocation).
|
|
||||||
pub n_kmers: Vec<u32>,
|
|
||||||
/// 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, 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());
|
|
||||||
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,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// ── SmerIndex — sparse "was this k-mer found at all" bookkeeping ─────────────
|
|
||||||
|
|
||||||
/// 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 `process_chunk`).
|
|
||||||
struct SmerIndex {
|
|
||||||
in_index: Vec<bool>, // total_smers
|
|
||||||
offsets: Vec<usize>, // offsets[i]..offsets[i+1] = s-mer range for sequence i
|
|
||||||
}
|
|
||||||
|
|
||||||
impl SmerIndex {
|
|
||||||
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.
|
|
||||||
fn mark_found(&mut self, seq: usize, kmer: usize) {
|
|
||||||
let abs = self.offsets[seq] + kmer;
|
|
||||||
self.in_index[abs] = true;
|
|
||||||
}
|
|
||||||
|
|
||||||
#[inline]
|
|
||||||
fn is_in_index(&self, seq: usize, kmer: usize) -> bool {
|
|
||||||
self.in_index[self.offsets[seq] + kmer]
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// ── Sparse Findere: per-genome run detection + sliding-window minimum ────────
|
|
||||||
|
|
||||||
/// 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`.
|
|
||||||
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.
|
|
||||||
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)
|
|
||||||
}
|
|
||||||
|
|
||||||
// ── Per-sequence accumulator ──────────────────────────────────────────────────
|
|
||||||
|
|
||||||
struct SeqAcc {
|
|
||||||
kmer_count: u32,
|
|
||||||
kmer_missing: u32,
|
|
||||||
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],
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// ── process_chunk ─────────────────────────────────────────────────────────────
|
|
||||||
|
|
||||||
fn process_chunk(
|
|
||||||
idx: &KmerIndex,
|
|
||||||
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,
|
|
||||||
) -> 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 (roadmap point 5) — 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 — unlike those, chunk_bytes's formula
|
|
||||||
// (run()) does not account for this at all today. 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<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::<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 (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, kmers) in batch.by_partition.iter().enumerate() {
|
|
||||||
if kmers.is_empty() {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
|
|
||||||
let stats = idx.partition()
|
|
||||||
.query_partition_with(
|
|
||||||
part_idx,
|
|
||||||
kmers,
|
|
||||||
n_genomes,
|
|
||||||
with_counts,
|
|
||||||
|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));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
},
|
|
||||||
)
|
|
||||||
.unwrap_or_else(|e| {
|
|
||||||
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,
|
|
||||||
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 — see the
|
|
||||||
// chunk-size formula's comment in `run()`), by allocated capacity rather
|
|
||||||
// than logical length so this reflects real memory pressure including
|
|
||||||
// Vec growth slack. `empirical_multiplier` is directly comparable to
|
|
||||||
// BYTES_PER_KMER_PER_GENOME (`run()`) — the ratio a cluster run's logs
|
|
||||||
// need to judge whether that constant is over- or under-conservative for
|
|
||||||
// real data, instead of guessing.
|
|
||||||
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.
|
|
||||||
// JSON per record ≈ 50 fixed chars + ~20 per genome (label + count value) + 100 (overhead).
|
|
||||||
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,
|
|
||||||
idx.meta(),
|
|
||||||
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
|
|
||||||
}
|
|
||||||
|
|
||||||
// ── GuardedChunkIter — keeps the throttle slot guard alive until the file is exhausted ──
|
|
||||||
|
|
||||||
/// Wraps a per-file `Rope` chunk iterator together with its `ThrottleGuard`,
|
|
||||||
/// so the guard (and the throttle slot it holds) is only released once the
|
|
||||||
/// file has been fully read — never earlier, never held past that point.
|
|
||||||
struct GuardedChunkIter {
|
|
||||||
inner: Box<dyn Iterator<Item = Rope> + Send>,
|
|
||||||
_guard: ThrottleGuard,
|
|
||||||
files_open: Arc<AtomicU32>,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl Iterator for GuardedChunkIter {
|
|
||||||
type Item = Rope;
|
|
||||||
fn next(&mut self) -> Option<Rope> {
|
|
||||||
self.inner.next()
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
impl Drop for GuardedChunkIter {
|
|
||||||
fn drop(&mut self) {
|
|
||||||
self.files_open.fetch_sub(1, Ordering::Relaxed);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// ── Entry point ───────────────────────────────────────────────────────────────
|
|
||||||
|
|
||||||
pub fn run(args: QueryArgs) {
|
|
||||||
let idx = Arc::new(KmerIndex::open(&args.index).unwrap_or_else(|e| {
|
|
||||||
eprintln!("error opening index: {e}");
|
|
||||||
std::process::exit(1);
|
|
||||||
}));
|
|
||||||
|
|
||||||
let k = idx.kmer_size();
|
|
||||||
let n_genomes = idx.meta().genomes.len();
|
|
||||||
let n_partitions = idx.n_partitions();
|
|
||||||
let with_counts = idx.meta().config.with_counts;
|
|
||||||
let n_workers = args.threads.max(1);
|
|
||||||
|
|
||||||
// Chunk size: each chunk stays in memory for its entire processing lifetime.
|
|
||||||
//
|
|
||||||
// Per-chunk memory is no longer a dense n_genomes-wide buffer (removed in
|
|
||||||
// the sparse Findere rework, see process_chunk) — it now scales with
|
|
||||||
// *actual hit count*, not with total_kmers_in_chunk × n_genomes
|
|
||||||
// unconditionally. BYTES_PER_KMER_PER_GENOME below is therefore a
|
|
||||||
// pathological-case bound, not a typical-case estimate: it protects
|
|
||||||
// against a fully-dense hit pattern (every k-mer of the query matching
|
|
||||||
// every genome — a degenerate case, e.g. low-complexity input theta-
|
|
||||||
// filtering should mostly reject, or an index of near-duplicate genomes),
|
|
||||||
// where by_genome and confirmed_by_genome (process_chunk) both end up
|
|
||||||
// holding one (seq_idx, pos, value) entry — 3 × u32 = 12 bytes, vs. 4
|
|
||||||
// bytes for the old dense encoding, where position was implicit in the
|
|
||||||
// array index — per (k-mer, genome) pair, and *coexist simultaneously*
|
|
||||||
// (by_genome isn't freed before confirmed_by_genome is built), for a
|
|
||||||
// worst case of ~24 bytes/pair before Vec growth slack. `cov` remains
|
|
||||||
// fully dense when --detail is set (unaffected by the sparse rework),
|
|
||||||
// still roughly doubling the n_genomes-scaled cost.
|
|
||||||
//
|
|
||||||
// For realistic, sparse hit patterns actual memory is far below this
|
|
||||||
// bound — see the "sparse memory retained" debug log in process_chunk,
|
|
||||||
// which reports the empirical bytes-per-raw-byte multiplier actually
|
|
||||||
// observed per chunk, directly comparable to BYTES_PER_KMER_PER_GENOME
|
|
||||||
// below. Tightening this constant for typical-case throughput (at the
|
|
||||||
// cost of pathological-case safety margin) is a deliberate tuning
|
|
||||||
// decision to make from that data, not something to guess at here.
|
|
||||||
//
|
|
||||||
// BASE_OVERHEAD approximates what scales with chunk_bytes alone,
|
|
||||||
// independent of n_genomes: the Rope itself, parsed SeqRecord sequence +
|
|
||||||
// normalised bytes, the superkmer dedup map, and the JSON output buffer.
|
|
||||||
// Like the n_genomes-scaled term, this is an estimate — validate against
|
|
||||||
// actual peak RSS (Stage::stop's `rss` in the summary table) on real
|
|
||||||
// workloads rather than trusting it blindly.
|
|
||||||
//
|
|
||||||
// We target ≤ 50 % of available RAM across all concurrent workers
|
|
||||||
// (SAFETY_FACTOR).
|
|
||||||
const BASE_OVERHEAD: u64 = 4;
|
|
||||||
const BYTES_PER_KMER_PER_GENOME: u64 = 8; // pathological-case bound — see comment above
|
|
||||||
const SAFETY_FACTOR: u64 = 2;
|
|
||||||
|
|
||||||
let detail_factor: u64 = if args.detail { 2 } else { 1 };
|
|
||||||
let overhead_multiplier =
|
|
||||||
BASE_OVERHEAD + n_genomes as u64 * BYTES_PER_KMER_PER_GENOME * detail_factor;
|
|
||||||
|
|
||||||
let chunk_bytes = args
|
|
||||||
.chunk_size
|
|
||||||
.map(|mb| mb * 1024 * 1024)
|
|
||||||
.unwrap_or_else(|| {
|
|
||||||
let avail = available_memory_bytes();
|
|
||||||
let computed = avail / (n_workers as u64 * overhead_multiplier * SAFETY_FACTOR);
|
|
||||||
computed.clamp(4 * 1024 * 1024, 256 * 1024 * 1024) as usize
|
|
||||||
});
|
|
||||||
|
|
||||||
debug!(
|
|
||||||
chunk_bytes,
|
|
||||||
n_genomes,
|
|
||||||
detail = args.detail,
|
|
||||||
overhead_multiplier,
|
|
||||||
estimated_peak_chunk_bytes = chunk_bytes as u64 * overhead_multiplier,
|
|
||||||
"chunk-size formula resolved"
|
|
||||||
);
|
|
||||||
|
|
||||||
let effective_z: usize = args
|
|
||||||
.findere_z
|
|
||||||
.unwrap_or_else(|| match idx.meta().config.evidence {
|
|
||||||
IndexMode::Approx { z, .. } | IndexMode::Hybrid { z, .. } => z as usize,
|
|
||||||
IndexMode::Exact => 1,
|
|
||||||
});
|
|
||||||
|
|
||||||
info!(
|
|
||||||
"query: k={k}, {} genome(s), with_counts={with_counts}, z={effective_z}, \
|
|
||||||
mismatch={}, detail={}",
|
|
||||||
n_genomes, args.mismatch, args.detail
|
|
||||||
);
|
|
||||||
|
|
||||||
if args.mismatch {
|
|
||||||
eprintln!("warning: --mismatch not yet implemented, ignored");
|
|
||||||
}
|
|
||||||
|
|
||||||
let detail = args.detail;
|
|
||||||
let count_missing = args.count_missing;
|
|
||||||
let force_presence = args.force_presence;
|
|
||||||
let presence_threshold = args.presence_threshold;
|
|
||||||
|
|
||||||
// Throttled iterator over input file paths: at most `effective_max_open()`
|
|
||||||
// files are open at once. Opening + decompressing + chunking each file is
|
|
||||||
// now a Flat pipeline stage, executed across the `n_workers` pool — not
|
|
||||||
// serialised in the pipe's dedicated source thread (see steps::scatter /
|
|
||||||
// cmd::superkmer for the same pattern applied to indexing).
|
|
||||||
info!("query: chunk_size={}MiB, max_open_files={}", chunk_bytes / (1024 * 1024), args.effective_max_open());
|
|
||||||
|
|
||||||
let paths: Vec<PathBuf> = args.inputs.iter().map(PathBuf::from).collect();
|
|
||||||
let path_source = throttle(paths.into_iter(), args.effective_max_open());
|
|
||||||
|
|
||||||
// Instrumentation: total bytes processed (for the EMA throughput readout),
|
|
||||||
// number of files currently open/being chunked, and number of chunks
|
|
||||||
// currently being processed by a worker — all read from the spinner loop
|
|
||||||
// below, updated from inside the pipe closures.
|
|
||||||
let total_bytes = Arc::new(AtomicU64::new(0));
|
|
||||||
let files_open = Arc::new(AtomicU32::new(0));
|
|
||||||
let chunks_active = Arc::new(AtomicU32::new(0));
|
|
||||||
|
|
||||||
let pipe = obipipeline::make_pipe! {
|
|
||||||
QueryData : Throttled<PathBuf> => Vec<u8>,
|
|
||||||
|| {
|
|
||||||
let files_open = Arc::clone(&files_open);
|
|
||||||
move |pw: Throttled<PathBuf>| -> GuardedChunkIter {
|
|
||||||
let path = pw.item;
|
|
||||||
let guard = pw.guard;
|
|
||||||
let path_str = path.to_str().unwrap_or("").to_owned();
|
|
||||||
files_open.fetch_add(1, Ordering::Relaxed);
|
|
||||||
let open_start = Instant::now();
|
|
||||||
// Hard-exit on file-open failure (mirrors the previous behaviour):
|
|
||||||
// propagating this as a pipeline Err would hit a known scheduler
|
|
||||||
// hang on early stage errors (obipipeline::scheduler::WorkerPool::run
|
|
||||||
// breaks its main loop without unblocking the still-running source
|
|
||||||
// thread, so the final `h.join()` never returns) — worth fixing in
|
|
||||||
// obipipeline itself, but out of scope here; sidestepping it like the
|
|
||||||
// original code already did is the safe choice for this change.
|
|
||||||
let iter = read_sequence_chunks_sized(&path_str, chunk_bytes).unwrap_or_else(|e| {
|
|
||||||
eprintln!("error opening {path_str}: {e}");
|
|
||||||
std::process::exit(1);
|
|
||||||
});
|
|
||||||
debug!(
|
|
||||||
path = %path_str,
|
|
||||||
open_ms = open_start.elapsed().as_millis() as u64,
|
|
||||||
"opened query input file"
|
|
||||||
);
|
|
||||||
let err_path = path_str.clone();
|
|
||||||
GuardedChunkIter {
|
|
||||||
inner: Box::new(iter.filter_map(move |r| match r {
|
|
||||||
Ok(rope) => Some(rope),
|
|
||||||
Err(e) => {
|
|
||||||
eprintln!("read error: {err_path}: {e}");
|
|
||||||
None
|
|
||||||
}
|
|
||||||
})),
|
|
||||||
_guard: guard,
|
|
||||||
files_open: Arc::clone(&files_open),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
} : Path => Chunk,
|
|
||||||
| {
|
|
||||||
let idx = Arc::clone(&idx);
|
|
||||||
let total_bytes = Arc::clone(&total_bytes);
|
|
||||||
let chunks_active = Arc::clone(&chunks_active);
|
|
||||||
move |rope: Rope| {
|
|
||||||
chunks_active.fetch_add(1, Ordering::Relaxed);
|
|
||||||
let bytes = rope.len() as u64;
|
|
||||||
let out = process_chunk(
|
|
||||||
&idx, rope, k, n_genomes, n_partitions, with_counts,
|
|
||||||
effective_z, detail, count_missing, force_presence, presence_threshold,
|
|
||||||
);
|
|
||||||
total_bytes.fetch_add(bytes, Ordering::Relaxed);
|
|
||||||
chunks_active.fetch_sub(1, Ordering::Relaxed);
|
|
||||||
out
|
|
||||||
}
|
|
||||||
} : Chunk => Output,
|
|
||||||
};
|
|
||||||
|
|
||||||
let t = Stage::start("query");
|
|
||||||
let pb = spinner("query");
|
|
||||||
|
|
||||||
let mut ema_rate: f64 = 0.0;
|
|
||||||
let mut last_t = Instant::now();
|
|
||||||
let mut last_bytes: u64 = 0;
|
|
||||||
const ALPHA: f64 = 0.15;
|
|
||||||
|
|
||||||
let mut out = BufWriter::new(io::stdout());
|
|
||||||
for block in pipe.apply(path_source, n_workers, 2) {
|
|
||||||
if !block.is_empty() {
|
|
||||||
out.write_all(&block).expect("write error");
|
|
||||||
}
|
|
||||||
|
|
||||||
let now = Instant::now();
|
|
||||||
let dt = now.duration_since(last_t).as_secs_f64();
|
|
||||||
if dt > 0.1 {
|
|
||||||
let total = total_bytes.load(Ordering::Relaxed);
|
|
||||||
let instant = (total - last_bytes) as f64 / dt;
|
|
||||||
ema_rate = ALPHA * instant + (1.0 - ALPHA) * ema_rate;
|
|
||||||
last_t = now;
|
|
||||||
last_bytes = total;
|
|
||||||
let bp = total as f64;
|
|
||||||
let (count_str, rate_str) = if bp >= 1e9 {
|
|
||||||
(format!("{:.2} GB", bp / 1e9), format!("{:.0} MB/s", ema_rate / 1e6))
|
|
||||||
} else {
|
|
||||||
(format!("{:.0} MB", bp / 1e6), format!("{:.0} MB/s", ema_rate / 1e6))
|
|
||||||
};
|
|
||||||
let active = chunks_active.load(Ordering::Relaxed);
|
|
||||||
let open = files_open.load(Ordering::Relaxed);
|
|
||||||
pb.set_message(format!("{count_str} {rate_str} [files open: {open}, chunks in flight: {active}]"));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
out.flush().expect("flush error");
|
|
||||||
|
|
||||||
pb.finish_and_clear();
|
|
||||||
|
|
||||||
let mut rep = Reporter::new();
|
|
||||||
rep.push(t.stop());
|
|
||||||
rep.print();
|
|
||||||
}
|
|
||||||
|
|
||||||
// ── Output ────────────────────────────────────────────────────────────────────
|
|
||||||
|
|
||||||
fn emit_batch(
|
|
||||||
batch: &QueryBatch,
|
|
||||||
accs: &[SeqAcc],
|
|
||||||
meta: &obikindex::meta::IndexMeta,
|
|
||||||
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 meta.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 meta.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");
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[cfg(test)]
|
|
||||||
#[path = "tests/query.rs"]
|
|
||||||
mod tests;
|
|
||||||
@@ -0,0 +1,74 @@
|
|||||||
|
use std::collections::HashMap;
|
||||||
|
|
||||||
|
use obikpartitionner::KmerDesc;
|
||||||
|
use obikseq::CanonicalKmer;
|
||||||
|
use obiread::record::SeqRecord;
|
||||||
|
use obiskbuilder::SuperKmerIter;
|
||||||
|
|
||||||
|
/// 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>,
|
||||||
|
/// Raw sequence bytes (for output), in batch order.
|
||||||
|
pub seqs: Vec<Vec<u8>>,
|
||||||
|
/// Total kmer count per sequence (used for `--detail` coverage allocation).
|
||||||
|
pub n_kmers: Vec<u32>,
|
||||||
|
/// 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, 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());
|
||||||
|
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,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,288 @@
|
|||||||
|
use std::time::Instant;
|
||||||
|
|
||||||
|
use obikindex::KmerIndex;
|
||||||
|
use obikpartitionner::{KmerDesc, QueryHit, QueryStats};
|
||||||
|
use obikrope::Rope;
|
||||||
|
use obikseq::CanonicalKmer;
|
||||||
|
use obiread::record::parse_chunk;
|
||||||
|
use tracing::debug;
|
||||||
|
|
||||||
|
use super::batch::QueryBatch;
|
||||||
|
use super::findere::{ConfirmedHit, sparse_findere_for_genome};
|
||||||
|
use super::output::emit_batch;
|
||||||
|
use super::smer_index::SmerIndex;
|
||||||
|
|
||||||
|
pub(super) struct SeqAcc {
|
||||||
|
pub(super) kmer_count: u32,
|
||||||
|
pub(super) kmer_missing: u32,
|
||||||
|
pub(super) 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],
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn process_chunk(
|
||||||
|
idx: &KmerIndex,
|
||||||
|
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,
|
||||||
|
) -> 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 (roadmap point 5) — 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 — unlike those, chunk_bytes's formula
|
||||||
|
// (run()) does not account for this at all today. 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<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::<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 (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, kmers) in batch.by_partition.iter().enumerate() {
|
||||||
|
if kmers.is_empty() {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
let stats = idx.partition()
|
||||||
|
.query_partition_with(
|
||||||
|
part_idx,
|
||||||
|
kmers,
|
||||||
|
n_genomes,
|
||||||
|
with_counts,
|
||||||
|
|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));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
},
|
||||||
|
)
|
||||||
|
.unwrap_or_else(|e| {
|
||||||
|
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,
|
||||||
|
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 — see the
|
||||||
|
// chunk-size formula's comment in `run()`), by allocated capacity rather
|
||||||
|
// than logical length so this reflects real memory pressure including
|
||||||
|
// Vec growth slack. `empirical_multiplier` is directly comparable to
|
||||||
|
// BYTES_PER_KMER_PER_GENOME (`run()`) — the ratio a cluster run's logs
|
||||||
|
// need to judge whether that constant is over- or under-conservative for
|
||||||
|
// real data, instead of guessing.
|
||||||
|
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.
|
||||||
|
// JSON per record ≈ 50 fixed chars + ~20 per genome (label + count value) + 100 (overhead).
|
||||||
|
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,
|
||||||
|
idx.meta(),
|
||||||
|
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,71 @@
|
|||||||
|
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(super) 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(super) 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)
|
||||||
|
}
|
||||||
@@ -0,0 +1,340 @@
|
|||||||
|
mod batch;
|
||||||
|
mod chunk;
|
||||||
|
mod findere;
|
||||||
|
mod output;
|
||||||
|
mod smer_index;
|
||||||
|
|
||||||
|
use std::io::{self, BufWriter, Write};
|
||||||
|
use std::path::PathBuf;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
|
||||||
|
use std::time::Instant;
|
||||||
|
|
||||||
|
use clap::Args;
|
||||||
|
use obikindex::KmerIndex;
|
||||||
|
use obikrope::Rope;
|
||||||
|
use obilayeredmap::IndexMode;
|
||||||
|
use obipipeline::{Throttled, ThrottleGuard, throttle};
|
||||||
|
use obiread::chunk::read_sequence_chunks_sized;
|
||||||
|
use obisys::{Reporter, Stage, available_memory_bytes, spinner};
|
||||||
|
use tracing::{debug, info};
|
||||||
|
|
||||||
|
use chunk::process_chunk;
|
||||||
|
|
||||||
|
// ── Pipeline data ─────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
enum QueryData {
|
||||||
|
Path(Throttled<PathBuf>),
|
||||||
|
Chunk(Rope),
|
||||||
|
Output(Vec<u8>),
|
||||||
|
}
|
||||||
|
|
||||||
|
// SAFETY: Rope contains Cell<u8> which is !Sync, but pipeline items are owned
|
||||||
|
// exclusively through channels — no item is ever shared across threads.
|
||||||
|
unsafe impl Send for QueryData {}
|
||||||
|
unsafe impl Sync for QueryData {}
|
||||||
|
|
||||||
|
// ── CLI ───────────────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
#[derive(Args)]
|
||||||
|
pub struct QueryArgs {
|
||||||
|
/// Index directory
|
||||||
|
pub index: PathBuf,
|
||||||
|
|
||||||
|
/// Input sequences (FASTA/FASTQ, optionally gzip-compressed)
|
||||||
|
#[arg(num_args = 1..)]
|
||||||
|
pub inputs: Vec<String>,
|
||||||
|
|
||||||
|
/// Report per-position coverage vectors per genome (adds "coverage" to JSON)
|
||||||
|
#[arg(long)]
|
||||||
|
pub detail: bool,
|
||||||
|
|
||||||
|
/// Enable 1-mismatch approximate matching
|
||||||
|
#[arg(long)]
|
||||||
|
pub mismatch: bool,
|
||||||
|
|
||||||
|
/// Count k-mers absent from the index (adds kmer_missing annotation)
|
||||||
|
#[arg(long)]
|
||||||
|
pub count_missing: bool,
|
||||||
|
|
||||||
|
/// Report per-genome presence (0/1) instead of raw counts
|
||||||
|
#[arg(long)]
|
||||||
|
pub force_presence: bool,
|
||||||
|
|
||||||
|
/// Minimum accumulated match count to declare a genome present (implies --force-presence)
|
||||||
|
#[arg(long, default_value_t = 1)]
|
||||||
|
pub presence_threshold: u32,
|
||||||
|
|
||||||
|
/// Override the Findere z parameter from index metadata
|
||||||
|
#[arg(short = 'z', long)]
|
||||||
|
pub findere_z: Option<usize>,
|
||||||
|
|
||||||
|
/// Number of worker threads
|
||||||
|
#[arg(
|
||||||
|
short = 'T',
|
||||||
|
long,
|
||||||
|
default_value_t = obisys::effective_parallelism()
|
||||||
|
)]
|
||||||
|
pub threads: usize,
|
||||||
|
|
||||||
|
/// I/O chunk size in MiB (default: auto-sized from available RAM and thread count)
|
||||||
|
#[arg(long)]
|
||||||
|
pub chunk_size: Option<usize>,
|
||||||
|
|
||||||
|
/// Maximum number of input files open simultaneously.
|
||||||
|
/// Defaults to threads/4 (minimum 1). Keep below the number of workers
|
||||||
|
/// to ensure CPU workers are always available for the transform stage.
|
||||||
|
#[arg(long)]
|
||||||
|
pub max_open_files: Option<usize>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl QueryArgs {
|
||||||
|
pub fn effective_max_open(&self) -> usize {
|
||||||
|
self.max_open_files
|
||||||
|
.unwrap_or_else(|| (self.threads / 4).max(1))
|
||||||
|
.max(1)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── GuardedChunkIter — keeps the throttle slot guard alive until the file is exhausted ──
|
||||||
|
|
||||||
|
/// Wraps a per-file `Rope` chunk iterator together with its `ThrottleGuard`,
|
||||||
|
/// so the guard (and the throttle slot it holds) is only released once the
|
||||||
|
/// file has been fully read — never earlier, never held past that point.
|
||||||
|
struct GuardedChunkIter {
|
||||||
|
inner: Box<dyn Iterator<Item = Rope> + Send>,
|
||||||
|
_guard: ThrottleGuard,
|
||||||
|
files_open: Arc<AtomicU32>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Iterator for GuardedChunkIter {
|
||||||
|
type Item = Rope;
|
||||||
|
fn next(&mut self) -> Option<Rope> {
|
||||||
|
self.inner.next()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Drop for GuardedChunkIter {
|
||||||
|
fn drop(&mut self) {
|
||||||
|
self.files_open.fetch_sub(1, Ordering::Relaxed);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Entry point ───────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
pub fn run(args: QueryArgs) {
|
||||||
|
let idx = Arc::new(KmerIndex::open(&args.index).unwrap_or_else(|e| {
|
||||||
|
eprintln!("error opening index: {e}");
|
||||||
|
std::process::exit(1);
|
||||||
|
}));
|
||||||
|
|
||||||
|
let k = idx.kmer_size();
|
||||||
|
let n_genomes = idx.meta().genomes.len();
|
||||||
|
let n_partitions = idx.n_partitions();
|
||||||
|
let with_counts = idx.meta().config.with_counts;
|
||||||
|
let n_workers = args.threads.max(1);
|
||||||
|
|
||||||
|
// Chunk size: each chunk stays in memory for its entire processing lifetime.
|
||||||
|
//
|
||||||
|
// Per-chunk memory is no longer a dense n_genomes-wide buffer (removed in
|
||||||
|
// the sparse Findere rework, see process_chunk) — it now scales with
|
||||||
|
// *actual hit count*, not with total_kmers_in_chunk × n_genomes
|
||||||
|
// unconditionally. BYTES_PER_KMER_PER_GENOME below is therefore a
|
||||||
|
// pathological-case bound, not a typical-case estimate: it protects
|
||||||
|
// against a fully-dense hit pattern (every k-mer of the query matching
|
||||||
|
// every genome — a degenerate case, e.g. low-complexity input theta-
|
||||||
|
// filtering should mostly reject, or an index of near-duplicate genomes),
|
||||||
|
// where by_genome and confirmed_by_genome (process_chunk) both end up
|
||||||
|
// holding one (seq_idx, pos, value) entry — 3 × u32 = 12 bytes, vs. 4
|
||||||
|
// bytes for the old dense encoding, where position was implicit in the
|
||||||
|
// array index — per (k-mer, genome) pair, and *coexist simultaneously*
|
||||||
|
// (by_genome isn't freed before confirmed_by_genome is built), for a
|
||||||
|
// worst case of ~24 bytes/pair before Vec growth slack. `cov` remains
|
||||||
|
// fully dense when --detail is set (unaffected by the sparse rework),
|
||||||
|
// still roughly doubling the n_genomes-scaled cost.
|
||||||
|
//
|
||||||
|
// For realistic, sparse hit patterns actual memory is far below this
|
||||||
|
// bound — see the "sparse memory retained" debug log in process_chunk,
|
||||||
|
// which reports the empirical bytes-per-raw-byte multiplier actually
|
||||||
|
// observed per chunk, directly comparable to BYTES_PER_KMER_PER_GENOME
|
||||||
|
// below. Tightening this constant for typical-case throughput (at the
|
||||||
|
// cost of pathological-case safety margin) is a deliberate tuning
|
||||||
|
// decision to make from that data, not something to guess at here.
|
||||||
|
//
|
||||||
|
// BASE_OVERHEAD approximates what scales with chunk_bytes alone,
|
||||||
|
// independent of n_genomes: the Rope itself, parsed SeqRecord sequence +
|
||||||
|
// normalised bytes, the superkmer dedup map, and the JSON output buffer.
|
||||||
|
// Like the n_genomes-scaled term, this is an estimate — validate against
|
||||||
|
// actual peak RSS (Stage::stop's `rss` in the summary table) on real
|
||||||
|
// workloads rather than trusting it blindly.
|
||||||
|
//
|
||||||
|
// We target ≤ 50 % of available RAM across all concurrent workers
|
||||||
|
// (SAFETY_FACTOR).
|
||||||
|
const BASE_OVERHEAD: u64 = 4;
|
||||||
|
const BYTES_PER_KMER_PER_GENOME: u64 = 8; // pathological-case bound — see comment above
|
||||||
|
const SAFETY_FACTOR: u64 = 2;
|
||||||
|
|
||||||
|
let detail_factor: u64 = if args.detail { 2 } else { 1 };
|
||||||
|
let overhead_multiplier =
|
||||||
|
BASE_OVERHEAD + n_genomes as u64 * BYTES_PER_KMER_PER_GENOME * detail_factor;
|
||||||
|
|
||||||
|
let chunk_bytes = args
|
||||||
|
.chunk_size
|
||||||
|
.map(|mb| mb * 1024 * 1024)
|
||||||
|
.unwrap_or_else(|| {
|
||||||
|
let avail = available_memory_bytes();
|
||||||
|
let computed = avail / (n_workers as u64 * overhead_multiplier * SAFETY_FACTOR);
|
||||||
|
computed.clamp(4 * 1024 * 1024, 256 * 1024 * 1024) as usize
|
||||||
|
});
|
||||||
|
|
||||||
|
debug!(
|
||||||
|
chunk_bytes,
|
||||||
|
n_genomes,
|
||||||
|
detail = args.detail,
|
||||||
|
overhead_multiplier,
|
||||||
|
estimated_peak_chunk_bytes = chunk_bytes as u64 * overhead_multiplier,
|
||||||
|
"chunk-size formula resolved"
|
||||||
|
);
|
||||||
|
|
||||||
|
let effective_z: usize = args
|
||||||
|
.findere_z
|
||||||
|
.unwrap_or_else(|| match idx.meta().config.evidence {
|
||||||
|
IndexMode::Approx { z, .. } | IndexMode::Hybrid { z, .. } => z as usize,
|
||||||
|
IndexMode::Exact => 1,
|
||||||
|
});
|
||||||
|
|
||||||
|
info!(
|
||||||
|
"query: k={k}, {} genome(s), with_counts={with_counts}, z={effective_z}, \
|
||||||
|
mismatch={}, detail={}",
|
||||||
|
n_genomes, args.mismatch, args.detail
|
||||||
|
);
|
||||||
|
|
||||||
|
if args.mismatch {
|
||||||
|
eprintln!("warning: --mismatch not yet implemented, ignored");
|
||||||
|
}
|
||||||
|
|
||||||
|
let detail = args.detail;
|
||||||
|
let count_missing = args.count_missing;
|
||||||
|
let force_presence = args.force_presence;
|
||||||
|
let presence_threshold = args.presence_threshold;
|
||||||
|
|
||||||
|
// Throttled iterator over input file paths: at most `effective_max_open()`
|
||||||
|
// files are open at once. Opening + decompressing + chunking each file is
|
||||||
|
// now a Flat pipeline stage, executed across the `n_workers` pool — not
|
||||||
|
// serialised in the pipe's dedicated source thread (see steps::scatter /
|
||||||
|
// cmd::superkmer for the same pattern applied to indexing).
|
||||||
|
info!("query: chunk_size={}MiB, max_open_files={}", chunk_bytes / (1024 * 1024), args.effective_max_open());
|
||||||
|
|
||||||
|
let paths: Vec<PathBuf> = args.inputs.iter().map(PathBuf::from).collect();
|
||||||
|
let path_source = throttle(paths.into_iter(), args.effective_max_open());
|
||||||
|
|
||||||
|
// Instrumentation: total bytes processed (for the EMA throughput readout),
|
||||||
|
// number of files currently open/being chunked, and number of chunks
|
||||||
|
// currently being processed by a worker — all read from the spinner loop
|
||||||
|
// below, updated from inside the pipe closures.
|
||||||
|
let total_bytes = Arc::new(AtomicU64::new(0));
|
||||||
|
let files_open = Arc::new(AtomicU32::new(0));
|
||||||
|
let chunks_active = Arc::new(AtomicU32::new(0));
|
||||||
|
|
||||||
|
let pipe = obipipeline::make_pipe! {
|
||||||
|
QueryData : Throttled<PathBuf> => Vec<u8>,
|
||||||
|
|| {
|
||||||
|
let files_open = Arc::clone(&files_open);
|
||||||
|
move |pw: Throttled<PathBuf>| -> GuardedChunkIter {
|
||||||
|
let path = pw.item;
|
||||||
|
let guard = pw.guard;
|
||||||
|
let path_str = path.to_str().unwrap_or("").to_owned();
|
||||||
|
files_open.fetch_add(1, Ordering::Relaxed);
|
||||||
|
let open_start = Instant::now();
|
||||||
|
// Hard-exit on file-open failure (mirrors the previous behaviour):
|
||||||
|
// propagating this as a pipeline Err would hit a known scheduler
|
||||||
|
// hang on early stage errors (obipipeline::scheduler::WorkerPool::run
|
||||||
|
// breaks its main loop without unblocking the still-running source
|
||||||
|
// thread, so the final `h.join()` never returns) — worth fixing in
|
||||||
|
// obipipeline itself, but out of scope here; sidestepping it like the
|
||||||
|
// original code already did is the safe choice for this change.
|
||||||
|
let iter = read_sequence_chunks_sized(&path_str, chunk_bytes).unwrap_or_else(|e| {
|
||||||
|
eprintln!("error opening {path_str}: {e}");
|
||||||
|
std::process::exit(1);
|
||||||
|
});
|
||||||
|
debug!(
|
||||||
|
path = %path_str,
|
||||||
|
open_ms = open_start.elapsed().as_millis() as u64,
|
||||||
|
"opened query input file"
|
||||||
|
);
|
||||||
|
let err_path = path_str.clone();
|
||||||
|
GuardedChunkIter {
|
||||||
|
inner: Box::new(iter.filter_map(move |r| match r {
|
||||||
|
Ok(rope) => Some(rope),
|
||||||
|
Err(e) => {
|
||||||
|
eprintln!("read error: {err_path}: {e}");
|
||||||
|
None
|
||||||
|
}
|
||||||
|
})),
|
||||||
|
_guard: guard,
|
||||||
|
files_open: Arc::clone(&files_open),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} : Path => Chunk,
|
||||||
|
| {
|
||||||
|
let idx = Arc::clone(&idx);
|
||||||
|
let total_bytes = Arc::clone(&total_bytes);
|
||||||
|
let chunks_active = Arc::clone(&chunks_active);
|
||||||
|
move |rope: Rope| {
|
||||||
|
chunks_active.fetch_add(1, Ordering::Relaxed);
|
||||||
|
let bytes = rope.len() as u64;
|
||||||
|
let out = process_chunk(
|
||||||
|
&idx, rope, k, n_genomes, n_partitions, with_counts,
|
||||||
|
effective_z, detail, count_missing, force_presence, presence_threshold,
|
||||||
|
);
|
||||||
|
total_bytes.fetch_add(bytes, Ordering::Relaxed);
|
||||||
|
chunks_active.fetch_sub(1, Ordering::Relaxed);
|
||||||
|
out
|
||||||
|
}
|
||||||
|
} : Chunk => Output,
|
||||||
|
};
|
||||||
|
|
||||||
|
let t = Stage::start("query");
|
||||||
|
let pb = spinner("query");
|
||||||
|
|
||||||
|
let mut ema_rate: f64 = 0.0;
|
||||||
|
let mut last_t = Instant::now();
|
||||||
|
let mut last_bytes: u64 = 0;
|
||||||
|
const ALPHA: f64 = 0.15;
|
||||||
|
|
||||||
|
let mut out = BufWriter::new(io::stdout());
|
||||||
|
for block in pipe.apply(path_source, n_workers, 2) {
|
||||||
|
if !block.is_empty() {
|
||||||
|
out.write_all(&block).expect("write error");
|
||||||
|
}
|
||||||
|
|
||||||
|
let now = Instant::now();
|
||||||
|
let dt = now.duration_since(last_t).as_secs_f64();
|
||||||
|
if dt > 0.1 {
|
||||||
|
let total = total_bytes.load(Ordering::Relaxed);
|
||||||
|
let instant = (total - last_bytes) as f64 / dt;
|
||||||
|
ema_rate = ALPHA * instant + (1.0 - ALPHA) * ema_rate;
|
||||||
|
last_t = now;
|
||||||
|
last_bytes = total;
|
||||||
|
let bp = total as f64;
|
||||||
|
let (count_str, rate_str) = if bp >= 1e9 {
|
||||||
|
(format!("{:.2} GB", bp / 1e9), format!("{:.0} MB/s", ema_rate / 1e6))
|
||||||
|
} else {
|
||||||
|
(format!("{:.0} MB", bp / 1e6), format!("{:.0} MB/s", ema_rate / 1e6))
|
||||||
|
};
|
||||||
|
let active = chunks_active.load(Ordering::Relaxed);
|
||||||
|
let open = files_open.load(Ordering::Relaxed);
|
||||||
|
pb.set_message(format!("{count_str} {rate_str} [files open: {open}, chunks in flight: {active}]"));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
out.flush().expect("flush error");
|
||||||
|
|
||||||
|
pb.finish_and_clear();
|
||||||
|
|
||||||
|
let mut rep = Reporter::new();
|
||||||
|
rep.push(t.stop());
|
||||||
|
rep.print();
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests;
|
||||||
@@ -0,0 +1,50 @@
|
|||||||
|
use std::io::Write;
|
||||||
|
|
||||||
|
use super::batch::QueryBatch;
|
||||||
|
use super::chunk::SeqAcc;
|
||||||
|
|
||||||
|
pub(super) fn emit_batch(
|
||||||
|
batch: &QueryBatch,
|
||||||
|
accs: &[SeqAcc],
|
||||||
|
meta: &obikindex::meta::IndexMeta,
|
||||||
|
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 meta.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 meta.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");
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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 `process_chunk`).
|
||||||
|
pub(super) 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(super) 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(super) fn mark_found(&mut self, seq: usize, kmer: usize) {
|
||||||
|
let abs = self.offsets[seq] + kmer;
|
||||||
|
self.in_index[abs] = true;
|
||||||
|
}
|
||||||
|
|
||||||
|
#[inline]
|
||||||
|
pub(super) fn is_in_index(&self, seq: usize, kmer: usize) -> bool {
|
||||||
|
self.in_index[self.offsets[seq] + kmer]
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -1,4 +1,9 @@
|
|||||||
use super::*;
|
use obikrope::Rope;
|
||||||
|
use obikseq::CanonicalKmer;
|
||||||
|
use obiread::record::parse_chunk;
|
||||||
|
|
||||||
|
use super::batch::QueryBatch;
|
||||||
|
use super::findere::sparse_findere_for_genome;
|
||||||
|
|
||||||
const K: usize = 11;
|
const K: usize = 11;
|
||||||
const M: usize = 5;
|
const M: usize = 5;
|
||||||
@@ -0,0 +1,112 @@
|
|||||||
|
use std::path::PathBuf;
|
||||||
|
|
||||||
|
use obikindex::{validate_label, IndexBitsPerKmer, KmerIndex};
|
||||||
|
use tracing::info;
|
||||||
|
|
||||||
|
pub(super) fn run_stats(index_path: &PathBuf) {
|
||||||
|
let idx = KmerIndex::open(index_path).unwrap_or_else(|e| {
|
||||||
|
eprintln!("error opening index: {e}");
|
||||||
|
std::process::exit(1);
|
||||||
|
});
|
||||||
|
let (total, per_genome) = idx.genome_kmer_counts().unwrap_or_else(|e| {
|
||||||
|
eprintln!("error computing stats: {e}");
|
||||||
|
std::process::exit(1);
|
||||||
|
});
|
||||||
|
println!("genome,n_kmers");
|
||||||
|
for (g, &n) in idx.meta().genomes.iter().zip(per_genome.iter()) {
|
||||||
|
println!("{},{}", g.label, n);
|
||||||
|
}
|
||||||
|
println!("total,{total}");
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn run_bits_per_kmer(index_path: &PathBuf) {
|
||||||
|
let idx = KmerIndex::open(index_path).unwrap_or_else(|e| {
|
||||||
|
eprintln!("error opening index: {e}");
|
||||||
|
std::process::exit(1);
|
||||||
|
});
|
||||||
|
let stats: IndexBitsPerKmer = idx.bits_per_kmer().unwrap_or_else(|e| {
|
||||||
|
eprintln!("error computing bits/kmer: {e}");
|
||||||
|
std::process::exit(1);
|
||||||
|
});
|
||||||
|
println!("k-mers : {}", stats.n_kmers);
|
||||||
|
println!("genomes : {}", stats.n_genomes);
|
||||||
|
println!("mphf : {:6.2} bits/kmer", stats.mphf);
|
||||||
|
println!("evidence : {:6.2} bits/kmer", stats.evidence);
|
||||||
|
println!(
|
||||||
|
"matrix : {:6.2} bits/kmer ({:.2} bits/kmer/genome)",
|
||||||
|
stats.matrix, stats.matrix_per_genome
|
||||||
|
);
|
||||||
|
println!("total : {:6.2} bits/kmer", stats.total);
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn run_upgrade_index(index_path: &PathBuf) {
|
||||||
|
let idx = KmerIndex::open(index_path).unwrap_or_else(|e| {
|
||||||
|
eprintln!("error opening index: {e}");
|
||||||
|
std::process::exit(1);
|
||||||
|
});
|
||||||
|
idx.upgrade_layer_meta().unwrap_or_else(|e| {
|
||||||
|
eprintln!("upgrade error: {e}");
|
||||||
|
std::process::exit(1);
|
||||||
|
});
|
||||||
|
info!("upgrade-index: layer_meta.json written to all layers that were missing it");
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn run_rename(index_path: &PathBuf, spec: &str) {
|
||||||
|
let (old_label, new_label) = parse_rename_spec(spec);
|
||||||
|
|
||||||
|
let mut idx = KmerIndex::open(index_path).unwrap_or_else(|e| {
|
||||||
|
eprintln!("error opening index: {e}");
|
||||||
|
std::process::exit(1);
|
||||||
|
});
|
||||||
|
|
||||||
|
let pos = idx
|
||||||
|
.meta()
|
||||||
|
.genomes
|
||||||
|
.iter()
|
||||||
|
.position(|g| g.label == old_label)
|
||||||
|
.unwrap_or_else(|| {
|
||||||
|
eprintln!("error: genome '{old_label}' not found in index");
|
||||||
|
std::process::exit(1);
|
||||||
|
});
|
||||||
|
|
||||||
|
validate_label(&new_label).unwrap_or_else(|e| {
|
||||||
|
eprintln!("error: --new-label: {e}");
|
||||||
|
std::process::exit(1);
|
||||||
|
});
|
||||||
|
|
||||||
|
if idx.meta().genomes.iter().any(|g| g.label == new_label) {
|
||||||
|
eprintln!("error: label '{new_label}' already exists in index");
|
||||||
|
std::process::exit(1);
|
||||||
|
}
|
||||||
|
|
||||||
|
idx.meta_mut().genomes[pos].label = new_label.clone();
|
||||||
|
idx.meta_mut().write(index_path).unwrap_or_else(|e| {
|
||||||
|
eprintln!("error writing index metadata: {e}");
|
||||||
|
std::process::exit(1);
|
||||||
|
});
|
||||||
|
|
||||||
|
let spectrums_dir = index_path.join("spectrums");
|
||||||
|
let old_spectrum = spectrums_dir.join(format!("{old_label}.json"));
|
||||||
|
let new_spectrum = spectrums_dir.join(format!("{new_label}.json"));
|
||||||
|
if old_spectrum.exists() {
|
||||||
|
std::fs::rename(&old_spectrum, &new_spectrum).unwrap_or_else(|e| {
|
||||||
|
eprintln!("warning: could not rename spectrum file: {e}");
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
info!("renamed genome '{old_label}' → '{new_label}'");
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_rename_spec(spec: &str) -> (String, String) {
|
||||||
|
let eq = spec.find('=').unwrap_or_else(|| {
|
||||||
|
eprintln!("error: --new-label expects NEW_LABEL=OLD_LABEL, got '{spec}'");
|
||||||
|
std::process::exit(1);
|
||||||
|
});
|
||||||
|
let new = spec[..eq].trim().to_string();
|
||||||
|
let old = spec[eq + 1..].trim().to_string();
|
||||||
|
if old.is_empty() || new.is_empty() {
|
||||||
|
eprintln!("error: --new-label: both old and new labels must be non-empty");
|
||||||
|
std::process::exit(1);
|
||||||
|
}
|
||||||
|
(old, new)
|
||||||
|
}
|
||||||
@@ -0,0 +1,85 @@
|
|||||||
|
mod maintenance;
|
||||||
|
mod partition_stats;
|
||||||
|
|
||||||
|
use std::path::PathBuf;
|
||||||
|
|
||||||
|
use clap::Args;
|
||||||
|
|
||||||
|
use maintenance::{run_bits_per_kmer, run_stats, run_upgrade_index, run_rename};
|
||||||
|
use partition_stats::run_partition_stats;
|
||||||
|
|
||||||
|
#[derive(Args)]
|
||||||
|
pub struct UtilsArgs {
|
||||||
|
/// Index directories to operate on (one or more)
|
||||||
|
#[arg(required = true, num_args = 1..)]
|
||||||
|
pub indexes: Vec<PathBuf>,
|
||||||
|
|
||||||
|
/// Set a new genome label: NEW_LABEL=OLD_LABEL (single-index only)
|
||||||
|
#[arg(long, value_name = "NEW=OLD")]
|
||||||
|
pub new_label: Option<String>,
|
||||||
|
|
||||||
|
/// Add missing layer_meta.json files to each layer (single-index only)
|
||||||
|
#[arg(long)]
|
||||||
|
pub upgrade_index: bool,
|
||||||
|
|
||||||
|
/// Print bits-per-kmer statistics (single-index only)
|
||||||
|
#[arg(long)]
|
||||||
|
pub bits_per_kmer: bool,
|
||||||
|
|
||||||
|
/// Print per-genome k-mer counts as CSV (single-index only)
|
||||||
|
#[arg(long)]
|
||||||
|
pub stats: bool,
|
||||||
|
|
||||||
|
/// Print partition size distribution report (accepts multiple indexes)
|
||||||
|
#[arg(long)]
|
||||||
|
pub partition_stats: bool,
|
||||||
|
|
||||||
|
/// Write per-(partition, source) raw data as CSV to FILE (used with --partition-stats)
|
||||||
|
#[arg(long, value_name = "FILE")]
|
||||||
|
pub csv: Option<PathBuf>,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn run(args: UtilsArgs) {
|
||||||
|
let mut any = false;
|
||||||
|
|
||||||
|
if let Some(spec) = &args.new_label {
|
||||||
|
any = true;
|
||||||
|
run_rename(single_index(&args), spec);
|
||||||
|
}
|
||||||
|
|
||||||
|
if args.upgrade_index {
|
||||||
|
any = true;
|
||||||
|
run_upgrade_index(single_index(&args));
|
||||||
|
}
|
||||||
|
|
||||||
|
if args.bits_per_kmer {
|
||||||
|
any = true;
|
||||||
|
run_bits_per_kmer(single_index(&args));
|
||||||
|
}
|
||||||
|
|
||||||
|
if args.stats {
|
||||||
|
any = true;
|
||||||
|
run_stats(single_index(&args));
|
||||||
|
}
|
||||||
|
|
||||||
|
if args.partition_stats {
|
||||||
|
any = true;
|
||||||
|
run_partition_stats(&args.indexes, args.csv.as_deref());
|
||||||
|
}
|
||||||
|
|
||||||
|
if !any {
|
||||||
|
eprintln!(
|
||||||
|
"utils: no operation specified. \
|
||||||
|
Available: --new-label, --upgrade-index, --bits-per-kmer, --stats, --partition-stats"
|
||||||
|
);
|
||||||
|
std::process::exit(1);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn single_index(args: &UtilsArgs) -> &PathBuf {
|
||||||
|
if args.indexes.len() > 1 {
|
||||||
|
eprintln!("utils: this option requires exactly one index (got {})", args.indexes.len());
|
||||||
|
std::process::exit(1);
|
||||||
|
}
|
||||||
|
&args.indexes[0]
|
||||||
|
}
|
||||||
@@ -1,89 +1,7 @@
|
|||||||
use std::io::{self, Write};
|
use std::io::{self, Write};
|
||||||
use std::path::PathBuf;
|
use std::path::PathBuf;
|
||||||
|
|
||||||
use clap::Args;
|
use obikindex::KmerIndex;
|
||||||
use obikindex::{validate_label, IndexBitsPerKmer, KmerIndex};
|
|
||||||
use tracing::info;
|
|
||||||
|
|
||||||
#[derive(Args)]
|
|
||||||
pub struct UtilsArgs {
|
|
||||||
/// Index directories to operate on (one or more)
|
|
||||||
#[arg(required = true, num_args = 1..)]
|
|
||||||
pub indexes: Vec<PathBuf>,
|
|
||||||
|
|
||||||
/// Set a new genome label: NEW_LABEL=OLD_LABEL (single-index only)
|
|
||||||
#[arg(long, value_name = "NEW=OLD")]
|
|
||||||
pub new_label: Option<String>,
|
|
||||||
|
|
||||||
/// Add missing layer_meta.json files to each layer (single-index only)
|
|
||||||
#[arg(long)]
|
|
||||||
pub upgrade_index: bool,
|
|
||||||
|
|
||||||
/// Print bits-per-kmer statistics (single-index only)
|
|
||||||
#[arg(long)]
|
|
||||||
pub bits_per_kmer: bool,
|
|
||||||
|
|
||||||
/// Print per-genome k-mer counts as CSV (single-index only)
|
|
||||||
#[arg(long)]
|
|
||||||
pub stats: bool,
|
|
||||||
|
|
||||||
/// Print partition size distribution report (accepts multiple indexes)
|
|
||||||
#[arg(long)]
|
|
||||||
pub partition_stats: bool,
|
|
||||||
|
|
||||||
/// Write per-(partition, source) raw data as CSV to FILE (used with --partition-stats)
|
|
||||||
#[arg(long, value_name = "FILE")]
|
|
||||||
pub csv: Option<PathBuf>,
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn run(args: UtilsArgs) {
|
|
||||||
let mut any = false;
|
|
||||||
|
|
||||||
if let Some(spec) = &args.new_label {
|
|
||||||
any = true;
|
|
||||||
run_rename(single_index(&args), spec);
|
|
||||||
}
|
|
||||||
|
|
||||||
if args.upgrade_index {
|
|
||||||
any = true;
|
|
||||||
run_upgrade_index(single_index(&args));
|
|
||||||
}
|
|
||||||
|
|
||||||
if args.bits_per_kmer {
|
|
||||||
any = true;
|
|
||||||
run_bits_per_kmer(single_index(&args));
|
|
||||||
}
|
|
||||||
|
|
||||||
if args.stats {
|
|
||||||
any = true;
|
|
||||||
run_stats(single_index(&args));
|
|
||||||
}
|
|
||||||
|
|
||||||
if args.partition_stats {
|
|
||||||
any = true;
|
|
||||||
run_partition_stats(&args.indexes, args.csv.as_deref());
|
|
||||||
}
|
|
||||||
|
|
||||||
if !any {
|
|
||||||
eprintln!(
|
|
||||||
"utils: no operation specified. \
|
|
||||||
Available: --new-label, --upgrade-index, --bits-per-kmer, --stats, --partition-stats"
|
|
||||||
);
|
|
||||||
std::process::exit(1);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// ── helpers ───────────────────────────────────────────────────────────────────
|
|
||||||
|
|
||||||
fn single_index(args: &UtilsArgs) -> &PathBuf {
|
|
||||||
if args.indexes.len() > 1 {
|
|
||||||
eprintln!("utils: this option requires exactly one index (got {})", args.indexes.len());
|
|
||||||
std::process::exit(1);
|
|
||||||
}
|
|
||||||
&args.indexes[0]
|
|
||||||
}
|
|
||||||
|
|
||||||
// ── --partition-stats ─────────────────────────────────────────────────────────
|
|
||||||
|
|
||||||
/// Per-partition, per-source byte count of all unitigs.bin files summed across layers.
|
/// Per-partition, per-source byte count of all unitigs.bin files summed across layers.
|
||||||
struct PartRow {
|
struct PartRow {
|
||||||
@@ -198,7 +116,7 @@ fn ascii_histogram(totals: &[u64], n_buckets: usize, bar_width: usize) -> String
|
|||||||
out
|
out
|
||||||
}
|
}
|
||||||
|
|
||||||
fn run_partition_stats(indexes: &[PathBuf], csv_path: Option<&std::path::Path>) {
|
pub(super) fn run_partition_stats(indexes: &[PathBuf], csv_path: Option<&std::path::Path>) {
|
||||||
let rows = collect_rows(indexes);
|
let rows = collect_rows(indexes);
|
||||||
if rows.is_empty() {
|
if rows.is_empty() {
|
||||||
eprintln!("partition-stats: no data found");
|
eprintln!("partition-stats: no data found");
|
||||||
@@ -282,113 +200,3 @@ fn run_partition_stats(indexes: &[PathBuf], csv_path: Option<&std::path::Path>)
|
|||||||
eprintln!("CSV written to {}", csv_out.display());
|
eprintln!("CSV written to {}", csv_out.display());
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// ── existing single-index operations ─────────────────────────────────────────
|
|
||||||
|
|
||||||
fn run_stats(index_path: &PathBuf) {
|
|
||||||
let idx = KmerIndex::open(index_path).unwrap_or_else(|e| {
|
|
||||||
eprintln!("error opening index: {e}");
|
|
||||||
std::process::exit(1);
|
|
||||||
});
|
|
||||||
let (total, per_genome) = idx.genome_kmer_counts().unwrap_or_else(|e| {
|
|
||||||
eprintln!("error computing stats: {e}");
|
|
||||||
std::process::exit(1);
|
|
||||||
});
|
|
||||||
println!("genome,n_kmers");
|
|
||||||
for (g, &n) in idx.meta().genomes.iter().zip(per_genome.iter()) {
|
|
||||||
println!("{},{}", g.label, n);
|
|
||||||
}
|
|
||||||
println!("total,{total}");
|
|
||||||
}
|
|
||||||
|
|
||||||
fn run_bits_per_kmer(index_path: &PathBuf) {
|
|
||||||
let idx = KmerIndex::open(index_path).unwrap_or_else(|e| {
|
|
||||||
eprintln!("error opening index: {e}");
|
|
||||||
std::process::exit(1);
|
|
||||||
});
|
|
||||||
let stats: IndexBitsPerKmer = idx.bits_per_kmer().unwrap_or_else(|e| {
|
|
||||||
eprintln!("error computing bits/kmer: {e}");
|
|
||||||
std::process::exit(1);
|
|
||||||
});
|
|
||||||
println!("k-mers : {}", stats.n_kmers);
|
|
||||||
println!("genomes : {}", stats.n_genomes);
|
|
||||||
println!("mphf : {:6.2} bits/kmer", stats.mphf);
|
|
||||||
println!("evidence : {:6.2} bits/kmer", stats.evidence);
|
|
||||||
println!(
|
|
||||||
"matrix : {:6.2} bits/kmer ({:.2} bits/kmer/genome)",
|
|
||||||
stats.matrix, stats.matrix_per_genome
|
|
||||||
);
|
|
||||||
println!("total : {:6.2} bits/kmer", stats.total);
|
|
||||||
}
|
|
||||||
|
|
||||||
fn run_upgrade_index(index_path: &PathBuf) {
|
|
||||||
let idx = KmerIndex::open(index_path).unwrap_or_else(|e| {
|
|
||||||
eprintln!("error opening index: {e}");
|
|
||||||
std::process::exit(1);
|
|
||||||
});
|
|
||||||
idx.upgrade_layer_meta().unwrap_or_else(|e| {
|
|
||||||
eprintln!("upgrade error: {e}");
|
|
||||||
std::process::exit(1);
|
|
||||||
});
|
|
||||||
info!("upgrade-index: layer_meta.json written to all layers that were missing it");
|
|
||||||
}
|
|
||||||
|
|
||||||
fn run_rename(index_path: &PathBuf, spec: &str) {
|
|
||||||
let (old_label, new_label) = parse_rename_spec(spec);
|
|
||||||
|
|
||||||
let mut idx = KmerIndex::open(index_path).unwrap_or_else(|e| {
|
|
||||||
eprintln!("error opening index: {e}");
|
|
||||||
std::process::exit(1);
|
|
||||||
});
|
|
||||||
|
|
||||||
let pos = idx
|
|
||||||
.meta()
|
|
||||||
.genomes
|
|
||||||
.iter()
|
|
||||||
.position(|g| g.label == old_label)
|
|
||||||
.unwrap_or_else(|| {
|
|
||||||
eprintln!("error: genome '{old_label}' not found in index");
|
|
||||||
std::process::exit(1);
|
|
||||||
});
|
|
||||||
|
|
||||||
validate_label(&new_label).unwrap_or_else(|e| {
|
|
||||||
eprintln!("error: --new-label: {e}");
|
|
||||||
std::process::exit(1);
|
|
||||||
});
|
|
||||||
|
|
||||||
if idx.meta().genomes.iter().any(|g| g.label == new_label) {
|
|
||||||
eprintln!("error: label '{new_label}' already exists in index");
|
|
||||||
std::process::exit(1);
|
|
||||||
}
|
|
||||||
|
|
||||||
idx.meta_mut().genomes[pos].label = new_label.clone();
|
|
||||||
idx.meta_mut().write(index_path).unwrap_or_else(|e| {
|
|
||||||
eprintln!("error writing index metadata: {e}");
|
|
||||||
std::process::exit(1);
|
|
||||||
});
|
|
||||||
|
|
||||||
let spectrums_dir = index_path.join("spectrums");
|
|
||||||
let old_spectrum = spectrums_dir.join(format!("{old_label}.json"));
|
|
||||||
let new_spectrum = spectrums_dir.join(format!("{new_label}.json"));
|
|
||||||
if old_spectrum.exists() {
|
|
||||||
std::fs::rename(&old_spectrum, &new_spectrum).unwrap_or_else(|e| {
|
|
||||||
eprintln!("warning: could not rename spectrum file: {e}");
|
|
||||||
});
|
|
||||||
}
|
|
||||||
|
|
||||||
info!("renamed genome '{old_label}' → '{new_label}'");
|
|
||||||
}
|
|
||||||
|
|
||||||
fn parse_rename_spec(spec: &str) -> (String, String) {
|
|
||||||
let eq = spec.find('=').unwrap_or_else(|| {
|
|
||||||
eprintln!("error: --new-label expects NEW_LABEL=OLD_LABEL, got '{spec}'");
|
|
||||||
std::process::exit(1);
|
|
||||||
});
|
|
||||||
let new = spec[..eq].trim().to_string();
|
|
||||||
let old = spec[eq + 1..].trim().to_string();
|
|
||||||
if old.is_empty() || new.is_empty() {
|
|
||||||
eprintln!("error: --new-label: both old and new labels must be non-empty");
|
|
||||||
std::process::exit(1);
|
|
||||||
}
|
|
||||||
(old, new)
|
|
||||||
}
|
|
||||||
Reference in New Issue
Block a user