refactor: Shift KmerIndex ownership to Arc for thread-safe sharing
Replaces lifetime-bound references with runtime reference counting across multiple crates. This enables safe concurrent access across parallel workers without explicit cloning or manual lifetime management. Introduces the `query` and `utils` CLI commands in obikmer2, along with supporting modules for batch processing, sparse indexing, sliding-window findere logic, and output formatting. Updates dependency manifests and aligns test suites with the new ownership model.
This commit is contained in:
Generated
+12
@@ -1673,12 +1673,16 @@ dependencies = [
|
|||||||
"obikalgorithm",
|
"obikalgorithm",
|
||||||
"obikdump",
|
"obikdump",
|
||||||
"obikfilter",
|
"obikfilter",
|
||||||
|
"obikidxcache",
|
||||||
"obikindex",
|
"obikindex",
|
||||||
"obikindexer",
|
"obikindexer",
|
||||||
"obikmerge",
|
"obikmerge",
|
||||||
|
"obikquery",
|
||||||
"obikrebuild",
|
"obikrebuild",
|
||||||
|
"obikrope",
|
||||||
"obikselect",
|
"obikselect",
|
||||||
"obikseq",
|
"obikseq",
|
||||||
|
"obikstats",
|
||||||
"obipipeline",
|
"obipipeline",
|
||||||
"obiread",
|
"obiread",
|
||||||
"obiskbuilder",
|
"obiskbuilder",
|
||||||
@@ -1732,7 +1736,15 @@ dependencies = [
|
|||||||
name = "obikquery"
|
name = "obikquery"
|
||||||
version = "0.1.0"
|
version = "0.1.0"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
|
"obikidxcache",
|
||||||
"obikindex",
|
"obikindex",
|
||||||
|
"obikrope",
|
||||||
|
"obikseq",
|
||||||
|
"obiread",
|
||||||
|
"obiskbuilder",
|
||||||
|
"serde_json",
|
||||||
|
"tempfile",
|
||||||
|
"tracing",
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
|
|||||||
@@ -1,4 +1,5 @@
|
|||||||
use std::io::Write;
|
use std::io::Write;
|
||||||
|
use std::sync::Arc;
|
||||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||||
|
|
||||||
use rayon::prelude::*;
|
use rayon::prelude::*;
|
||||||
@@ -39,7 +40,7 @@ pub trait IndexDump {
|
|||||||
) -> OKIResult<()>;
|
) -> OKIResult<()>;
|
||||||
}
|
}
|
||||||
|
|
||||||
impl IndexDump for KmerIndex {
|
impl IndexDump for Arc<KmerIndex> {
|
||||||
fn dump<W: Write, F: Fn() + Send + Sync>(
|
fn dump<W: Write, F: Fn() + Send + Sync>(
|
||||||
&self,
|
&self,
|
||||||
out: &mut W,
|
out: &mut W,
|
||||||
@@ -87,7 +88,7 @@ impl IndexDump for KmerIndex {
|
|||||||
Ok(_) => { write_row(buf, row, prefix); true }
|
Ok(_) => { write_row(buf, row, prefix); true }
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
let cache = IndexCache::new(self, Some(vec![i]));
|
let cache = IndexCache::new(Arc::clone(self), Some(vec![i]));
|
||||||
if debug {
|
if debug {
|
||||||
cache.iter_partition_kmers_located(i, use_counts, n_genomes, filters, |part, layer, kmer, row| {
|
cache.iter_partition_kmers_located(i, use_counts, n_genomes, filters, |part, layer, kmer, row| {
|
||||||
let seq = String::from_utf8(kmer.to_ascii()).unwrap_or_else(|_| "?".repeat(kmer_size));
|
let seq = String::from_utf8(kmer.to_ascii()).unwrap_or_else(|_| "?".repeat(kmer_size));
|
||||||
@@ -106,7 +107,7 @@ impl IndexDump for KmerIndex {
|
|||||||
// ── Unbounded: no atomic, no contention ───────────────────────────
|
// ── Unbounded: no atomic, no contention ───────────────────────────
|
||||||
(0..n).into_par_iter().map(|i| {
|
(0..n).into_par_iter().map(|i| {
|
||||||
let mut buf = Vec::<u8>::new();
|
let mut buf = Vec::<u8>::new();
|
||||||
let cache = IndexCache::new(self, Some(vec![i]));
|
let cache = IndexCache::new(Arc::clone(self), Some(vec![i]));
|
||||||
if debug {
|
if debug {
|
||||||
cache.iter_partition_kmers_located(i, use_counts, n_genomes, filters, |part, layer, kmer, row| {
|
cache.iter_partition_kmers_located(i, use_counts, n_genomes, filters, |part, layer, kmer, row| {
|
||||||
let seq = String::from_utf8(kmer.to_ascii()).unwrap_or_else(|_| "?".repeat(kmer_size));
|
let seq = String::from_utf8(kmer.to_ascii()).unwrap_or_else(|_| "?".repeat(kmer_size));
|
||||||
|
|||||||
@@ -5,6 +5,7 @@
|
|||||||
//! foreign type, so this is an extension trait rather than an inherent `impl`.
|
//! foreign type, so this is an extension trait rather than an inherent `impl`.
|
||||||
|
|
||||||
use std::io::Write;
|
use std::io::Write;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
use obidebruinj::GraphDeBruijn;
|
use obidebruinj::GraphDeBruijn;
|
||||||
use obifastwrite::write_unitig;
|
use obifastwrite::write_unitig;
|
||||||
@@ -25,7 +26,7 @@ pub trait IndexUnitigs {
|
|||||||
) -> OKIResult<usize>;
|
) -> OKIResult<usize>;
|
||||||
}
|
}
|
||||||
|
|
||||||
impl IndexUnitigs for KmerIndex {
|
impl IndexUnitigs for Arc<KmerIndex> {
|
||||||
fn write_unitigs<W: Write + Send, F: Fn() + Send + Sync>(
|
fn write_unitigs<W: Write + Send, F: Fn() + Send + Sync>(
|
||||||
&self,
|
&self,
|
||||||
out: &mut W,
|
out: &mut W,
|
||||||
@@ -40,7 +41,7 @@ impl IndexUnitigs for KmerIndex {
|
|||||||
let g = (0..n)
|
let g = (0..n)
|
||||||
.into_par_iter()
|
.into_par_iter()
|
||||||
.try_fold(GraphDeBruijn::new, |mut local_g, i| -> OKIResult<GraphDeBruijn> {
|
.try_fold(GraphDeBruijn::new, |mut local_g, i| -> OKIResult<GraphDeBruijn> {
|
||||||
let cache = IndexCache::new(self, Some(vec![i]));
|
let cache = IndexCache::new(Arc::clone(self), Some(vec![i]));
|
||||||
cache.iter_partition_kmers(i, use_counts, n_genomes, filters, |kmer, _row| {
|
cache.iter_partition_kmers(i, use_counts, n_genomes, filters, |kmer, _row| {
|
||||||
local_g.push(kmer);
|
local_g.push(kmer);
|
||||||
true
|
true
|
||||||
|
|||||||
@@ -9,6 +9,7 @@
|
|||||||
|
|
||||||
use std::io;
|
use std::io;
|
||||||
use std::path::PathBuf;
|
use std::path::PathBuf;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
use obikalgorithm::Algorithm;
|
use obikalgorithm::Algorithm;
|
||||||
use obikindex::{IndexBuilder, IndexState, KmerIndex};
|
use obikindex::{IndexBuilder, IndexState, KmerIndex};
|
||||||
@@ -19,7 +20,7 @@ use crate::filter::KmerFilter;
|
|||||||
use crate::filter_partition::filter_partition;
|
use crate::filter_partition::filter_partition;
|
||||||
|
|
||||||
pub struct Filter<'a> {
|
pub struct Filter<'a> {
|
||||||
src: &'a KmerIndex,
|
src: Arc<KmerIndex>,
|
||||||
output: PathBuf,
|
output: PathBuf,
|
||||||
filters: &'a [Box<dyn KmerFilter>],
|
filters: &'a [Box<dyn KmerFilter>],
|
||||||
presence: bool,
|
presence: bool,
|
||||||
@@ -30,7 +31,7 @@ pub struct Filter<'a> {
|
|||||||
}
|
}
|
||||||
|
|
||||||
impl<'a> Filter<'a> {
|
impl<'a> Filter<'a> {
|
||||||
pub fn new(src: &'a KmerIndex, output: impl Into<PathBuf>, filters: &'a [Box<dyn KmerFilter>]) -> Self {
|
pub fn new(src: Arc<KmerIndex>, output: impl Into<PathBuf>, filters: &'a [Box<dyn KmerFilter>]) -> Self {
|
||||||
Self {
|
Self {
|
||||||
src,
|
src,
|
||||||
output: output.into(),
|
output: output.into(),
|
||||||
@@ -80,7 +81,7 @@ impl Algorithm for Filter<'_> {
|
|||||||
type Output = KmerIndex;
|
type Output = KmerIndex;
|
||||||
|
|
||||||
fn run(&mut self) -> obikalgorithm::Result<KmerIndex> {
|
fn run(&mut self) -> obikalgorithm::Result<KmerIndex> {
|
||||||
let src = self.src;
|
let src = &self.src;
|
||||||
let output = self.output.clone();
|
let output = self.output.clone();
|
||||||
|
|
||||||
if src.state()? != IndexState::Indexed {
|
if src.state()? != IndexState::Indexed {
|
||||||
|
|||||||
@@ -5,6 +5,7 @@
|
|||||||
//! copied — removing kmers changes MPHF slot assignment layer-wide.
|
//! copied — removing kmers changes MPHF slot assignment layer-wide.
|
||||||
|
|
||||||
use std::collections::HashMap;
|
use std::collections::HashMap;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
use obidebruinj::GraphDeBruijn;
|
use obidebruinj::GraphDeBruijn;
|
||||||
use obikidxcache::index_cache::IndexCache;
|
use obikidxcache::index_cache::IndexCache;
|
||||||
@@ -25,7 +26,7 @@ use crate::partition_iter::FilteredPartitionIter;
|
|||||||
/// output from a count source.
|
/// output from a count source.
|
||||||
pub(crate) fn filter_partition(
|
pub(crate) fn filter_partition(
|
||||||
dst: &KmerIndex,
|
dst: &KmerIndex,
|
||||||
src: &KmerIndex,
|
src: &Arc<KmerIndex>,
|
||||||
i: usize,
|
i: usize,
|
||||||
filters: &[Box<dyn KmerFilter>],
|
filters: &[Box<dyn KmerFilter>],
|
||||||
use_counts: bool,
|
use_counts: bool,
|
||||||
@@ -35,7 +36,7 @@ pub(crate) fn filter_partition(
|
|||||||
) -> OKIResult<()> {
|
) -> OKIResult<()> {
|
||||||
let n_genomes = src.meta().genomes().map_err(OKIError::Io)?.len();
|
let n_genomes = src.meta().genomes().map_err(OKIError::Io)?.len();
|
||||||
|
|
||||||
let cache = IndexCache::new(src, Some(vec![i]));
|
let cache = IndexCache::new(Arc::clone(src), Some(vec![i]));
|
||||||
let dst_partition = dst.partition(i)?;
|
let dst_partition = dst.partition(i)?;
|
||||||
let n_layers = src.partition(i)?.n_layers();
|
let n_layers = src.partition(i)?.n_layers();
|
||||||
|
|
||||||
|
|||||||
@@ -52,7 +52,7 @@ pub trait FilteredPartitionIter {
|
|||||||
) -> OKIResult<bool>;
|
) -> OKIResult<bool>;
|
||||||
}
|
}
|
||||||
|
|
||||||
impl FilteredPartitionIter for IndexCache<'_> {
|
impl FilteredPartitionIter for IndexCache {
|
||||||
fn iter_partition_kmers(
|
fn iter_partition_kmers(
|
||||||
&self,
|
&self,
|
||||||
part: usize,
|
part: usize,
|
||||||
|
|||||||
@@ -1,18 +1,23 @@
|
|||||||
use std::collections::HashMap;
|
use std::collections::HashMap;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
use obikindex::{KmerIndex, layer::KmerLayer};
|
use obikindex::{KmerIndex, layer::KmerLayer};
|
||||||
use obikseq::CanonicalKmer;
|
use obikseq::CanonicalKmer;
|
||||||
|
|
||||||
use crate::meta_cache::MetaCache;
|
use crate::meta_cache::MetaCache;
|
||||||
|
|
||||||
pub struct IndexCache<'a> {
|
pub struct IndexCache {
|
||||||
/// Ties this cache's lifetime to the `KmerIndex` it was opened from —
|
/// Owns a strong reference to the `KmerIndex` it was opened from — the
|
||||||
/// borrow-checker-enforced, not just documentation: an `IndexCache`
|
/// cache cannot outlive its index (the `Arc` keeps it alive at least as
|
||||||
/// cannot outlive its index, so it can never be read against an index
|
/// long as the cache itself, enforced at runtime by refcounting rather
|
||||||
/// that has since been mutated or dropped. `index()` is a genuine
|
/// than by a borrow-checked lifetime). `Arc`, not `&'a KmerIndex`,
|
||||||
/// accessor, but the field's real job is this lifetime bound, kept even
|
/// because `IndexCache` must itself be `'static`-capable to be shared
|
||||||
/// on paths that never call `index()`.
|
/// across `obipipeline`'s `thread::spawn`-based workers (see
|
||||||
raw_index: &'a KmerIndex,
|
/// `obikquery::query_layer`) — a plain borrow can never satisfy that,
|
||||||
|
/// regardless of how it's threaded through. `index()` is a genuine
|
||||||
|
/// accessor, but the field's real job is this ownership guarantee, kept
|
||||||
|
/// even on paths that never call `index()`.
|
||||||
|
raw_index: Arc<KmerIndex>,
|
||||||
meta: MetaCache,
|
meta: MetaCache,
|
||||||
|
|
||||||
/// Keyed by each partition's own index number, not renumbered —
|
/// Keyed by each partition's own index number, not renumbered —
|
||||||
@@ -24,7 +29,7 @@ pub struct IndexCache<'a> {
|
|||||||
layer_cache: HashMap<usize, Vec<KmerLayer>>,
|
layer_cache: HashMap<usize, Vec<KmerLayer>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl<'a> IndexCache<'a> {
|
impl IndexCache {
|
||||||
/// Opens every layer of the given `partitions` (every partition of the
|
/// Opens every layer of the given `partitions` (every partition of the
|
||||||
/// index, if `None`) once, up front, and keeps them alive for the
|
/// index, if `None`) once, up front, and keeps them alive for the
|
||||||
/// cache's lifetime — each partition's own layer count is read
|
/// cache's lifetime — each partition's own layer count is read
|
||||||
@@ -41,7 +46,7 @@ impl<'a> IndexCache<'a> {
|
|||||||
/// cache can only be built on partitions that actually exist and are
|
/// cache can only be built on partitions that actually exist and are
|
||||||
/// complete, so any failure here is a caller error, not a case to
|
/// complete, so any failure here is a caller error, not a case to
|
||||||
/// design around with a `Result`.
|
/// design around with a `Result`.
|
||||||
pub fn new(index: &'a KmerIndex, partitions: Option<Vec<usize>>) -> Self {
|
pub fn new(index: Arc<KmerIndex>, partitions: Option<Vec<usize>>) -> Self {
|
||||||
let selected = partitions.unwrap_or_else(|| (0..index.n_partitions()).collect());
|
let selected = partitions.unwrap_or_else(|| (0..index.n_partitions()).collect());
|
||||||
|
|
||||||
let mut layer_cache = HashMap::with_capacity(selected.len());
|
let mut layer_cache = HashMap::with_capacity(selected.len());
|
||||||
@@ -62,17 +67,18 @@ impl<'a> IndexCache<'a> {
|
|||||||
layer_cache.insert(p, layers);
|
layer_cache.insert(p, layers);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
let meta = MetaCache::from_meta(&index.meta());
|
||||||
IndexCache {
|
IndexCache {
|
||||||
raw_index: index,
|
raw_index: index,
|
||||||
meta: MetaCache::from_meta(&index.meta()),
|
meta,
|
||||||
layer_cache,
|
layer_cache,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// The index this cache was opened from.
|
/// The index this cache was opened from.
|
||||||
#[inline]
|
#[inline]
|
||||||
pub fn index(&self) -> &'a KmerIndex {
|
pub fn index(&self) -> &KmerIndex {
|
||||||
self.raw_index
|
&self.raw_index
|
||||||
}
|
}
|
||||||
|
|
||||||
/// The index's metadata, snapshotted once at construction — see
|
/// The index's metadata, snapshotted once at construction — see
|
||||||
|
|||||||
@@ -20,6 +20,10 @@ obikfilter = { path = "../obikfilter" }
|
|||||||
obikselect = { path = "../obikselect" }
|
obikselect = { path = "../obikselect" }
|
||||||
obikdump = { path = "../obikdump" }
|
obikdump = { path = "../obikdump" }
|
||||||
obikrebuild = { path = "../obikrebuild" }
|
obikrebuild = { path = "../obikrebuild" }
|
||||||
|
obikstats = { path = "../obikstats" }
|
||||||
|
obikquery = { path = "../obikquery" }
|
||||||
|
obikidxcache = { path = "../obikidxcache" }
|
||||||
|
obikrope = { path = "../obikrope" }
|
||||||
obifastwrite = { path = "../obifastwrite" }
|
obifastwrite = { path = "../obifastwrite" }
|
||||||
obiskbuilder = { path = "../obiskbuilder" }
|
obiskbuilder = { path = "../obiskbuilder" }
|
||||||
clap = { version = "4", features = ["derive"] }
|
clap = { version = "4", features = ["derive"] }
|
||||||
|
|||||||
@@ -1,5 +1,6 @@
|
|||||||
use std::io::{self, BufWriter};
|
use std::io::{self, BufWriter};
|
||||||
use std::path::PathBuf;
|
use std::path::PathBuf;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
use clap::Args;
|
use clap::Args;
|
||||||
use obikdump::IndexDump;
|
use obikdump::IndexDump;
|
||||||
@@ -32,10 +33,10 @@ pub struct DumpArgs {
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub fn run(args: DumpArgs) {
|
pub fn run(args: DumpArgs) {
|
||||||
let idx = KmerIndex::open(&args.index).unwrap_or_else(|e| {
|
let idx = Arc::new(KmerIndex::open(&args.index).unwrap_or_else(|e| {
|
||||||
eprintln!("error opening index: {e}");
|
eprintln!("error opening index: {e}");
|
||||||
std::process::exit(1);
|
std::process::exit(1);
|
||||||
});
|
}));
|
||||||
|
|
||||||
let n_genomes = idx.meta().genomes().unwrap_or_else(|e| {
|
let n_genomes = idx.meta().genomes().unwrap_or_else(|e| {
|
||||||
eprintln!("error reading index metadata: {e}");
|
eprintln!("error reading index metadata: {e}");
|
||||||
|
|||||||
@@ -1,4 +1,5 @@
|
|||||||
use std::path::PathBuf;
|
use std::path::PathBuf;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
use clap::Args;
|
use clap::Args;
|
||||||
use obikalgorithm::Algorithm;
|
use obikalgorithm::Algorithm;
|
||||||
@@ -51,10 +52,10 @@ pub struct FilterArgs {
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub fn run(args: FilterArgs) {
|
pub fn run(args: FilterArgs) {
|
||||||
let src = KmerIndex::open(&args.source).unwrap_or_else(|e| {
|
let src = Arc::new(KmerIndex::open(&args.source).unwrap_or_else(|e| {
|
||||||
eprintln!("error opening source index: {e}");
|
eprintln!("error opening source index: {e}");
|
||||||
std::process::exit(1);
|
std::process::exit(1);
|
||||||
});
|
}));
|
||||||
|
|
||||||
let mut filters: Vec<Box<dyn KmerFilter>> =
|
let mut filters: Vec<Box<dyn KmerFilter>> =
|
||||||
vec![Box::new(args.group_filter.build_filter(&src.meta()))];
|
vec![Box::new(args.group_filter.build_filter(&src.meta()))];
|
||||||
@@ -80,7 +81,7 @@ pub fn run(args: FilterArgs) {
|
|||||||
let mut rep = Reporter::new();
|
let mut rep = Reporter::new();
|
||||||
let t = Stage::start("filter");
|
let t = Stage::start("filter");
|
||||||
let pb = progress_bar("filter", src.n_partitions() as u64, "partitions");
|
let pb = progress_bar("filter", src.n_partitions() as u64, "partitions");
|
||||||
let mut alg = Filter::new(&src, &args.output, &filters)
|
let mut alg = Filter::new(Arc::clone(&src), &args.output, &filters)
|
||||||
.presence(args.presence)
|
.presence(args.presence)
|
||||||
.force(args.force)
|
.force(args.force)
|
||||||
.sparse(!args.dense)
|
.sparse(!args.dense)
|
||||||
|
|||||||
@@ -6,6 +6,8 @@ pub mod index;
|
|||||||
pub mod merge;
|
pub mod merge;
|
||||||
pub mod pack;
|
pub mod pack;
|
||||||
mod predicate;
|
mod predicate;
|
||||||
|
pub mod query;
|
||||||
pub mod select;
|
pub mod select;
|
||||||
pub mod superkmer;
|
pub mod superkmer;
|
||||||
pub mod unitig;
|
pub mod unitig;
|
||||||
|
pub mod utils;
|
||||||
|
|||||||
@@ -1,4 +1,5 @@
|
|||||||
use std::path::PathBuf;
|
use std::path::PathBuf;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
use clap::Args;
|
use clap::Args;
|
||||||
use obikindex::KmerIndex;
|
use obikindex::KmerIndex;
|
||||||
@@ -32,10 +33,10 @@ pub fn run(args: PackArgs) {
|
|||||||
std::process::exit(1);
|
std::process::exit(1);
|
||||||
});
|
});
|
||||||
|
|
||||||
let idx = KmerIndex::open(&args.index).unwrap_or_else(|e| {
|
let idx = Arc::new(KmerIndex::open(&args.index).unwrap_or_else(|e| {
|
||||||
eprintln!("error opening index: {e}");
|
eprintln!("error opening index: {e}");
|
||||||
std::process::exit(1);
|
std::process::exit(1);
|
||||||
});
|
}));
|
||||||
|
|
||||||
let n_genomes = idx.meta().genomes().unwrap_or_else(|e| {
|
let n_genomes = idx.meta().genomes().unwrap_or_else(|e| {
|
||||||
eprintln!("error reading index metadata: {e}");
|
eprintln!("error reading index metadata: {e}");
|
||||||
|
|||||||
@@ -0,0 +1,343 @@
|
|||||||
|
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 obikidxcache::index_cache::IndexCache;
|
||||||
|
use obikindex::KmerIndex;
|
||||||
|
use obikindex::layer::IndexMode;
|
||||||
|
use obikquery::process_chunk;
|
||||||
|
use obikrope::Rope;
|
||||||
|
use obipipeline::{Throttled, ThrottleGuard, throttle};
|
||||||
|
use obiread::chunk::read_sequence_chunks_sized;
|
||||||
|
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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── 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 genomes = idx.meta().genomes().unwrap_or_else(|e| {
|
||||||
|
eprintln!("error reading index metadata: {e}");
|
||||||
|
std::process::exit(1);
|
||||||
|
});
|
||||||
|
let n_genomes = genomes.len();
|
||||||
|
let genomes = Arc::new(genomes);
|
||||||
|
let n_partitions = idx.n_partitions();
|
||||||
|
let with_counts = idx.meta().config.with_counts;
|
||||||
|
let n_workers = args.threads.max(1);
|
||||||
|
|
||||||
|
// Every partition/layer the query might touch is opened once, up front,
|
||||||
|
// and shared (via Arc) across every `obipipeline` worker — a query pass
|
||||||
|
// is then pure in-memory lookups, never a per-chunk disk open (the
|
||||||
|
// previous `obikindex`-based design's cost). `IndexCache` owns its
|
||||||
|
// `Arc<KmerIndex>`, so it's itself `'static`-capable, satisfying
|
||||||
|
// obipipeline's `Send + Sync + 'static` requirement on pipeline data —
|
||||||
|
// see `obikquery::query_layer`'s doc comment for why that rules out a
|
||||||
|
// borrow-based cache here.
|
||||||
|
let cache = Arc::new(IndexCache::new(Arc::clone(&idx), None));
|
||||||
|
|
||||||
|
// Chunk size: each chunk stays in memory for its entire processing lifetime.
|
||||||
|
//
|
||||||
|
// Per-chunk memory is not a dense n_genomes-wide buffer — it 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 (obikquery::chunk::process_chunk)
|
||||||
|
// both end up holding one (seq_idx, pos, value) entry — 3 × u32 = 12
|
||||||
|
// bytes — 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, 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
|
||||||
|
// a Flat pipeline stage, executed across the `n_workers` pool.
|
||||||
|
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 cache = Arc::clone(&cache);
|
||||||
|
let genomes = Arc::clone(&genomes);
|
||||||
|
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(
|
||||||
|
&cache, rope, k, n_genomes, n_partitions, with_counts,
|
||||||
|
effective_z, detail, count_missing, force_presence, presence_threshold,
|
||||||
|
&genomes,
|
||||||
|
);
|
||||||
|
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();
|
||||||
|
}
|
||||||
@@ -1,5 +1,6 @@
|
|||||||
use std::io::{self, BufWriter};
|
use std::io::{self, BufWriter};
|
||||||
use std::path::PathBuf;
|
use std::path::PathBuf;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
use clap::Args;
|
use clap::Args;
|
||||||
use obikdump::IndexUnitigs;
|
use obikdump::IndexUnitigs;
|
||||||
@@ -20,10 +21,10 @@ pub struct UnitigArgs {
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub fn run(args: UnitigArgs) {
|
pub fn run(args: UnitigArgs) {
|
||||||
let idx = KmerIndex::open(&args.index).unwrap_or_else(|e| {
|
let idx = Arc::new(KmerIndex::open(&args.index).unwrap_or_else(|e| {
|
||||||
eprintln!("error opening index: {e}");
|
eprintln!("error opening index: {e}");
|
||||||
std::process::exit(1);
|
std::process::exit(1);
|
||||||
});
|
}));
|
||||||
|
|
||||||
info!(
|
info!(
|
||||||
"unitig: building de Bruijn graph from {} partition(s) (k={})",
|
"unitig: building de Bruijn graph from {} partition(s) (k={})",
|
||||||
|
|||||||
@@ -0,0 +1,113 @@
|
|||||||
|
use std::path::PathBuf;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use obikalgorithm::Algorithm;
|
||||||
|
use obikindex::{GenomeInfo, KmerIndex};
|
||||||
|
use obikstats::{BitsPerKmer, GenomeKmerCounts};
|
||||||
|
use tracing::info;
|
||||||
|
|
||||||
|
pub(super) fn run_stats(index_path: &PathBuf) {
|
||||||
|
let idx = Arc::new(KmerIndex::open(index_path).unwrap_or_else(|e| {
|
||||||
|
eprintln!("error opening index: {e}");
|
||||||
|
std::process::exit(1);
|
||||||
|
}));
|
||||||
|
let genomes = idx.meta().genomes().unwrap_or_else(|e| {
|
||||||
|
eprintln!("error reading index metadata: {e}");
|
||||||
|
std::process::exit(1);
|
||||||
|
});
|
||||||
|
|
||||||
|
let (total, per_genome) = GenomeKmerCounts::new(Arc::clone(&idx)).run().unwrap_or_else(|e| {
|
||||||
|
eprintln!("error computing stats: {e}");
|
||||||
|
std::process::exit(1);
|
||||||
|
});
|
||||||
|
|
||||||
|
println!("genome,n_kmers");
|
||||||
|
for (g, &n) in 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 = Arc::new(KmerIndex::open(index_path).unwrap_or_else(|e| {
|
||||||
|
eprintln!("error opening index: {e}");
|
||||||
|
std::process::exit(1);
|
||||||
|
}));
|
||||||
|
|
||||||
|
let stats = BitsPerKmer::new(idx).run().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_rename(index_path: &PathBuf, spec: &str) {
|
||||||
|
let (old_label, new_label) = parse_rename_spec(spec);
|
||||||
|
|
||||||
|
let idx = KmerIndex::open(index_path).unwrap_or_else(|e| {
|
||||||
|
eprintln!("error opening index: {e}");
|
||||||
|
std::process::exit(1);
|
||||||
|
});
|
||||||
|
|
||||||
|
let genomes = idx.meta().genomes().unwrap_or_else(|e| {
|
||||||
|
eprintln!("error reading index metadata: {e}");
|
||||||
|
std::process::exit(1);
|
||||||
|
});
|
||||||
|
|
||||||
|
let pos = genomes
|
||||||
|
.iter()
|
||||||
|
.position(|g| g.label == old_label)
|
||||||
|
.unwrap_or_else(|| {
|
||||||
|
eprintln!("error: genome '{old_label}' not found in index");
|
||||||
|
std::process::exit(1);
|
||||||
|
});
|
||||||
|
|
||||||
|
GenomeInfo::validate_label(&new_label).unwrap_or_else(|e| {
|
||||||
|
eprintln!("error: --new-label: {e}");
|
||||||
|
std::process::exit(1);
|
||||||
|
});
|
||||||
|
|
||||||
|
if genomes.iter().any(|g| g.label == new_label) {
|
||||||
|
eprintln!("error: label '{new_label}' already exists in index");
|
||||||
|
std::process::exit(1);
|
||||||
|
}
|
||||||
|
|
||||||
|
idx.meta().rename_genome(pos, new_label.clone()).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,76 @@
|
|||||||
|
mod maintenance;
|
||||||
|
mod partition_stats;
|
||||||
|
|
||||||
|
use std::path::PathBuf;
|
||||||
|
|
||||||
|
use clap::Args;
|
||||||
|
|
||||||
|
use maintenance::{run_bits_per_kmer, run_stats, 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>,
|
||||||
|
|
||||||
|
/// 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.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, --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]
|
||||||
|
}
|
||||||
@@ -0,0 +1,206 @@
|
|||||||
|
use std::io::{self, Write};
|
||||||
|
use std::path::PathBuf;
|
||||||
|
|
||||||
|
use obikindex::KmerIndex;
|
||||||
|
|
||||||
|
/// Per-partition, per-source byte count of all unitigs.bin files summed across layers.
|
||||||
|
struct PartRow {
|
||||||
|
partition: usize,
|
||||||
|
source: String,
|
||||||
|
bytes: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
fn collect_rows(indexes: &[PathBuf]) -> Vec<PartRow> {
|
||||||
|
let mut rows = Vec::new();
|
||||||
|
for path in indexes {
|
||||||
|
let idx = KmerIndex::open(path).unwrap_or_else(|e| {
|
||||||
|
eprintln!("error opening index {}: {e}", path.display());
|
||||||
|
std::process::exit(1);
|
||||||
|
});
|
||||||
|
let name = path
|
||||||
|
.file_name()
|
||||||
|
.map(|n| n.to_string_lossy().into_owned())
|
||||||
|
.unwrap_or_else(|| path.display().to_string());
|
||||||
|
let n_parts = idx.n_partitions();
|
||||||
|
for i in 0..n_parts {
|
||||||
|
let mut bytes = 0u64;
|
||||||
|
let n_layers = idx.n_layers(i).unwrap_or_else(|e| {
|
||||||
|
eprintln!("error reading partition {i} of {}: {e}", path.display());
|
||||||
|
std::process::exit(1);
|
||||||
|
});
|
||||||
|
for l in 0..n_layers {
|
||||||
|
let p = idx.layer_unitigs_path(i, l).unwrap_or_else(|e| {
|
||||||
|
eprintln!("error reading layer {l} of partition {i} of {}: {e}", path.display());
|
||||||
|
std::process::exit(1);
|
||||||
|
});
|
||||||
|
if let Ok(m) = std::fs::metadata(&p) {
|
||||||
|
bytes += m.len();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
rows.push(PartRow { partition: i, source: name.clone(), bytes });
|
||||||
|
}
|
||||||
|
}
|
||||||
|
rows
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Sum bytes per partition across all sources.
|
||||||
|
fn partition_totals(rows: &[PartRow], n_parts: usize) -> Vec<u64> {
|
||||||
|
let mut totals = vec![0u64; n_parts];
|
||||||
|
for r in rows {
|
||||||
|
totals[r.partition] += r.bytes;
|
||||||
|
}
|
||||||
|
totals
|
||||||
|
}
|
||||||
|
|
||||||
|
fn stats_summary(totals: &[u64]) -> (u64, u64, f64, f64, u64, u64, u64) {
|
||||||
|
let mut sorted = totals.to_vec();
|
||||||
|
sorted.sort_unstable();
|
||||||
|
let n = sorted.len();
|
||||||
|
let min = sorted[0];
|
||||||
|
let max = sorted[n - 1];
|
||||||
|
let mean = sorted.iter().sum::<u64>() as f64 / n as f64;
|
||||||
|
let median = if n % 2 == 0 {
|
||||||
|
(sorted[n / 2 - 1] + sorted[n / 2]) as f64 / 2.0
|
||||||
|
} else {
|
||||||
|
sorted[n / 2] as f64
|
||||||
|
};
|
||||||
|
let p95 = sorted[(n as f64 * 0.95) as usize];
|
||||||
|
let p99 = sorted[(n as f64 * 0.99) as usize];
|
||||||
|
let variance = sorted
|
||||||
|
.iter()
|
||||||
|
.map(|&v| (v as f64 - mean).powi(2))
|
||||||
|
.sum::<f64>()
|
||||||
|
/ n as f64;
|
||||||
|
let std_dev = variance.sqrt();
|
||||||
|
(min, max, mean, median, p95, p99, std_dev as u64)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn human_bytes(b: u64) -> String {
|
||||||
|
if b >= 1 << 30 {
|
||||||
|
format!("{:.1} GB", b as f64 / (1u64 << 30) as f64)
|
||||||
|
} else if b >= 1 << 20 {
|
||||||
|
format!("{:.1} MB", b as f64 / (1u64 << 20) as f64)
|
||||||
|
} else if b >= 1 << 10 {
|
||||||
|
format!("{:.1} KB", b as f64 / (1u64 << 10) as f64)
|
||||||
|
} else {
|
||||||
|
format!("{b} B")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn ascii_histogram(totals: &[u64], n_buckets: usize, bar_width: usize) -> String {
|
||||||
|
let min = *totals.iter().min().unwrap();
|
||||||
|
let max = *totals.iter().max().unwrap();
|
||||||
|
if min == max {
|
||||||
|
return format!(" (all partitions identical: {})\n", human_bytes(min));
|
||||||
|
}
|
||||||
|
|
||||||
|
let bucket_size = (max - min).max(1) as f64 / n_buckets as f64;
|
||||||
|
let mut counts = vec![0usize; n_buckets];
|
||||||
|
for &v in totals {
|
||||||
|
let b = (((v - min) as f64 / bucket_size) as usize).min(n_buckets - 1);
|
||||||
|
counts[b] += 1;
|
||||||
|
}
|
||||||
|
|
||||||
|
let max_count = *counts.iter().max().unwrap();
|
||||||
|
let mut out = String::new();
|
||||||
|
for (i, &c) in counts.iter().enumerate() {
|
||||||
|
let lo = min + (i as f64 * bucket_size) as u64;
|
||||||
|
let hi = min + ((i + 1) as f64 * bucket_size) as u64;
|
||||||
|
let bar_len = if max_count > 0 { c * bar_width / max_count } else { 0 };
|
||||||
|
let bar = "█".repeat(bar_len);
|
||||||
|
out.push_str(&format!(
|
||||||
|
" {:>8} – {:>8} │{:<width$} {}\n",
|
||||||
|
human_bytes(lo),
|
||||||
|
human_bytes(hi),
|
||||||
|
bar,
|
||||||
|
c,
|
||||||
|
width = bar_width
|
||||||
|
));
|
||||||
|
}
|
||||||
|
out
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn run_partition_stats(indexes: &[PathBuf], csv_path: Option<&std::path::Path>) {
|
||||||
|
let rows = collect_rows(indexes);
|
||||||
|
if rows.is_empty() {
|
||||||
|
eprintln!("partition-stats: no data found");
|
||||||
|
std::process::exit(1);
|
||||||
|
}
|
||||||
|
|
||||||
|
let n_parts = rows.iter().map(|r| r.partition).max().unwrap() + 1;
|
||||||
|
let totals = partition_totals(&rows, n_parts);
|
||||||
|
let (min, max, mean, median, p95, p99, std_dev) = stats_summary(&totals);
|
||||||
|
|
||||||
|
// outliers: > median + 1.5 × IQR (approximate via > 1.5 × median as fallback)
|
||||||
|
let mut sorted_t = totals.clone();
|
||||||
|
sorted_t.sort_unstable();
|
||||||
|
let q1 = sorted_t[n_parts / 4] as f64;
|
||||||
|
let q3 = sorted_t[3 * n_parts / 4] as f64;
|
||||||
|
let iqr = q3 - q1;
|
||||||
|
let outlier_threshold = q3 + 1.5 * iqr;
|
||||||
|
|
||||||
|
let mut out = String::new();
|
||||||
|
out.push_str("# Partition size report\n\n");
|
||||||
|
out.push_str(&format!(
|
||||||
|
"Sources: {} \nPartitions: {} \n\n",
|
||||||
|
indexes.len(),
|
||||||
|
n_parts
|
||||||
|
));
|
||||||
|
|
||||||
|
out.push_str("## Summary statistics (total unitigs.bin bytes per partition, sum across sources)\n\n");
|
||||||
|
out.push_str("| Stat | Value |\n|---|---|\n");
|
||||||
|
out.push_str(&format!("| min | {} |\n", human_bytes(min)));
|
||||||
|
out.push_str(&format!("| max | {} |\n", human_bytes(max)));
|
||||||
|
out.push_str(&format!("| mean | {} |\n", human_bytes(mean as u64)));
|
||||||
|
out.push_str(&format!("| median | {} |\n", human_bytes(median as u64)));
|
||||||
|
out.push_str(&format!("| p95 | {} |\n", human_bytes(p95)));
|
||||||
|
out.push_str(&format!("| p99 | {} |\n", human_bytes(p99)));
|
||||||
|
out.push_str(&format!("| std | {} |\n", human_bytes(std_dev)));
|
||||||
|
out.push_str(&format!("| max/median ratio | {:.2}× |\n\n", max as f64 / median));
|
||||||
|
|
||||||
|
out.push_str("## Histogram\n\n```\n");
|
||||||
|
out.push_str(&ascii_histogram(&totals, 30, 40));
|
||||||
|
out.push_str("```\n\n");
|
||||||
|
|
||||||
|
let outliers: Vec<(usize, u64)> = totals
|
||||||
|
.iter()
|
||||||
|
.enumerate()
|
||||||
|
.filter(|(_, v)| **v as f64 > outlier_threshold)
|
||||||
|
.map(|(i, v)| (i, *v))
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
if outliers.is_empty() {
|
||||||
|
out.push_str("## Outliers\n\nNone (threshold: Q3 + 1.5×IQR = ");
|
||||||
|
out.push_str(&human_bytes(outlier_threshold as u64));
|
||||||
|
out.push_str(").\n");
|
||||||
|
} else {
|
||||||
|
out.push_str(&format!(
|
||||||
|
"## Outliers (> Q3 + 1.5×IQR = {})\n\n| Partition | Total size | Ratio to median |\n|---|---|---|\n",
|
||||||
|
human_bytes(outlier_threshold as u64)
|
||||||
|
));
|
||||||
|
for (i, v) in &outliers {
|
||||||
|
out.push_str(&format!(
|
||||||
|
"| {} | {} | {:.2}× |\n",
|
||||||
|
i,
|
||||||
|
human_bytes(*v),
|
||||||
|
*v as f64 / median
|
||||||
|
));
|
||||||
|
}
|
||||||
|
out.push('\n');
|
||||||
|
}
|
||||||
|
|
||||||
|
print!("{out}");
|
||||||
|
|
||||||
|
if let Some(csv_out) = csv_path {
|
||||||
|
let file = std::fs::File::create(csv_out).unwrap_or_else(|e| {
|
||||||
|
eprintln!("error creating CSV file {}: {e}", csv_out.display());
|
||||||
|
std::process::exit(1);
|
||||||
|
});
|
||||||
|
let mut w = io::BufWriter::new(file);
|
||||||
|
writeln!(w, "partition,source,bytes").unwrap();
|
||||||
|
for r in &rows {
|
||||||
|
writeln!(w, "{},{},{}", r.partition, r.source, r.bytes).unwrap();
|
||||||
|
}
|
||||||
|
eprintln!("CSV written to {}", csv_out.display());
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -25,6 +25,8 @@ enum Commands {
|
|||||||
Select(cmd::select::SelectArgs),
|
Select(cmd::select::SelectArgs),
|
||||||
/// Dump an index's kmers as a CSV table
|
/// Dump an index's kmers as a CSV table
|
||||||
Dump(cmd::dump::DumpArgs),
|
Dump(cmd::dump::DumpArgs),
|
||||||
|
/// Query sequences against an index, annotating each with per-genome matches
|
||||||
|
Query(cmd::query::QueryArgs),
|
||||||
/// Assemble an index's kmers into unitigs and write them as FASTA
|
/// Assemble an index's kmers into unitigs and write them as FASTA
|
||||||
Unitig(cmd::unitig::UnitigArgs),
|
Unitig(cmd::unitig::UnitigArgs),
|
||||||
/// Pack an index's matrices into single-file, sparse format by default (--dense to opt out), in place
|
/// Pack an index's matrices into single-file, sparse format by default (--dense to opt out), in place
|
||||||
@@ -33,6 +35,8 @@ enum Commands {
|
|||||||
Estimate(cmd::estimate::EstimateArgs),
|
Estimate(cmd::estimate::EstimateArgs),
|
||||||
/// Read/write genome metadata (CSV) on an already-built index
|
/// Read/write genome metadata (CSV) on an already-built index
|
||||||
Annotate(cmd::annotate::AnnotateArgs),
|
Annotate(cmd::annotate::AnnotateArgs),
|
||||||
|
/// Maintenance/inspection operations on already-built indexes
|
||||||
|
Utils(cmd::utils::UtilsArgs),
|
||||||
}
|
}
|
||||||
|
|
||||||
fn main() {
|
fn main() {
|
||||||
@@ -51,9 +55,11 @@ fn main() {
|
|||||||
Commands::Filter(args) => cmd::filter::run(args),
|
Commands::Filter(args) => cmd::filter::run(args),
|
||||||
Commands::Select(args) => cmd::select::run(args),
|
Commands::Select(args) => cmd::select::run(args),
|
||||||
Commands::Dump(args) => cmd::dump::run(args),
|
Commands::Dump(args) => cmd::dump::run(args),
|
||||||
|
Commands::Query(args) => cmd::query::run(args),
|
||||||
Commands::Unitig(args) => cmd::unitig::run(args),
|
Commands::Unitig(args) => cmd::unitig::run(args),
|
||||||
Commands::Pack(args) => cmd::pack::run(args),
|
Commands::Pack(args) => cmd::pack::run(args),
|
||||||
Commands::Estimate(args) => cmd::estimate::run(args),
|
Commands::Estimate(args) => cmd::estimate::run(args),
|
||||||
Commands::Annotate(args) => cmd::annotate::run(args),
|
Commands::Annotate(args) => cmd::annotate::run(args),
|
||||||
|
Commands::Utils(args) => cmd::utils::run(args),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -5,3 +5,14 @@ edition = "2024"
|
|||||||
|
|
||||||
[dependencies]
|
[dependencies]
|
||||||
obikindex = { path = "../obikindex" }
|
obikindex = { path = "../obikindex" }
|
||||||
|
obikidxcache = { path = "../obikidxcache" }
|
||||||
|
obikseq = { path = "../obikseq" }
|
||||||
|
obikrope = { path = "../obikrope" }
|
||||||
|
obiread = { path = "../obiread" }
|
||||||
|
obiskbuilder = { path = "../obiskbuilder" }
|
||||||
|
serde_json = "1"
|
||||||
|
tracing = "0.1.44"
|
||||||
|
|
||||||
|
[dev-dependencies]
|
||||||
|
obikseq = { path = "../obikseq", features = ["test-utils"] }
|
||||||
|
tempfile = "3"
|
||||||
|
|||||||
@@ -0,0 +1,79 @@
|
|||||||
|
use std::collections::HashMap;
|
||||||
|
|
||||||
|
use obikseq::CanonicalKmer;
|
||||||
|
use obiread::record::SeqRecord;
|
||||||
|
use obiskbuilder::SuperKmerIter;
|
||||||
|
|
||||||
|
use crate::query_layer::KmerDesc;
|
||||||
|
|
||||||
|
/// A batch of query sequences, with k-mers deduplicated directly (not just at
|
||||||
|
/// the superkmer level) and pre-split by partition.
|
||||||
|
///
|
||||||
|
/// Superkmer *construction* (`SuperKmerIter`) is still required — it's the
|
||||||
|
/// mechanism that computes minimizers and partition routing — but the dedup
|
||||||
|
/// key is the canonical k-mer, not the superkmer: two different superkmers
|
||||||
|
/// that happen to share a k-mer (read overlaps, repeats, a SNP splitting an
|
||||||
|
/// otherwise-identical run) are deduplicated too, not just identical whole
|
||||||
|
/// superkmers. This also means each unique k-mer triggers at most one MPHF
|
||||||
|
/// lookup, not one per occurrence.
|
||||||
|
pub(crate) struct QueryBatch {
|
||||||
|
/// Sequence ids in batch order.
|
||||||
|
pub(crate) ids: Vec<String>,
|
||||||
|
/// Raw sequence bytes (for output), in batch order.
|
||||||
|
pub(crate) seqs: Vec<Vec<u8>>,
|
||||||
|
/// Total kmer count per sequence (used for `--detail` coverage allocation).
|
||||||
|
pub(crate) n_kmers: Vec<u32>,
|
||||||
|
/// Deduplicated k-mer occurrences, one map per partition.
|
||||||
|
pub(crate) by_partition: Vec<HashMap<CanonicalKmer, Vec<KmerDesc>>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl QueryBatch {
|
||||||
|
/// Build a batch from a vec of parsed sequence records, deduplicating
|
||||||
|
/// k-mers and routing them to partitions in the same pass.
|
||||||
|
pub(crate) fn from_records(
|
||||||
|
records: Vec<SeqRecord>,
|
||||||
|
k: usize,
|
||||||
|
level_max: usize,
|
||||||
|
theta: f64,
|
||||||
|
n_partitions: usize,
|
||||||
|
) -> Self {
|
||||||
|
let mut ids = Vec::with_capacity(records.len());
|
||||||
|
let mut seqs = Vec::with_capacity(records.len());
|
||||||
|
let mut n_kmers = Vec::with_capacity(records.len());
|
||||||
|
let mask = (n_partitions as u64) - 1;
|
||||||
|
let mut by_partition: Vec<HashMap<CanonicalKmer, Vec<KmerDesc>>> =
|
||||||
|
(0..n_partitions).map(|_| HashMap::new()).collect();
|
||||||
|
|
||||||
|
for (seq_idx, record) in records.into_iter().enumerate() {
|
||||||
|
let mut kmer_offset = 0u32;
|
||||||
|
|
||||||
|
for rsk in SuperKmerIter::new(&record.normalized, k, level_max, theta) {
|
||||||
|
let part_idx = (rsk.minimizer().seq_hash() & mask) as usize;
|
||||||
|
let map = &mut by_partition[part_idx];
|
||||||
|
for (j, kmer) in rsk.superkmer().iter_canonical_kmers().enumerate() {
|
||||||
|
map.entry(kmer).or_default().push(KmerDesc {
|
||||||
|
seq_idx: seq_idx as u32,
|
||||||
|
pos: kmer_offset + j as u32,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
let n = (rsk.seql() - k + 1) as u32;
|
||||||
|
kmer_offset += n;
|
||||||
|
}
|
||||||
|
|
||||||
|
ids.push(record.id);
|
||||||
|
seqs.push(record.sequence);
|
||||||
|
n_kmers.push(kmer_offset);
|
||||||
|
}
|
||||||
|
|
||||||
|
Self {
|
||||||
|
ids,
|
||||||
|
seqs,
|
||||||
|
n_kmers,
|
||||||
|
by_partition,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
#[path = "tests/batch.rs"]
|
||||||
|
mod tests;
|
||||||
@@ -0,0 +1,279 @@
|
|||||||
|
use std::time::Instant;
|
||||||
|
|
||||||
|
use obikidxcache::index_cache::IndexCache;
|
||||||
|
use obikindex::GenomeInfo;
|
||||||
|
use obikrope::Rope;
|
||||||
|
use obikseq::CanonicalKmer;
|
||||||
|
use obiread::record::parse_chunk;
|
||||||
|
use tracing::debug;
|
||||||
|
|
||||||
|
use crate::batch::QueryBatch;
|
||||||
|
use crate::findere::{ConfirmedHit, sparse_findere_for_genome};
|
||||||
|
use crate::output::emit_batch;
|
||||||
|
use crate::query_layer::{QueryHit, QueryPartition, QueryStats};
|
||||||
|
use crate::smer_index::SmerIndex;
|
||||||
|
|
||||||
|
pub(crate) struct SeqAcc {
|
||||||
|
pub(crate) kmer_count: u32,
|
||||||
|
pub(crate) kmer_missing: u32,
|
||||||
|
pub(crate) genome_totals: Vec<u32>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl SeqAcc {
|
||||||
|
fn new(n_genomes: usize) -> Self {
|
||||||
|
Self {
|
||||||
|
kmer_count: 0,
|
||||||
|
kmer_missing: 0,
|
||||||
|
genome_totals: vec![0u32; n_genomes],
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Turn one chunk of raw sequence bytes into an obitools-style annotated
|
||||||
|
/// FASTA byte buffer, matching its k-mers against `cache`'s index.
|
||||||
|
///
|
||||||
|
/// `cache` is expected to already hold every partition/layer this chunk
|
||||||
|
/// might touch (built once, up front, for the whole run — see
|
||||||
|
/// `obikidxcache::IndexCache::new`), so every lookup here is a plain
|
||||||
|
/// in-memory operation, never a disk open.
|
||||||
|
#[allow(clippy::too_many_arguments)]
|
||||||
|
pub fn process_chunk(
|
||||||
|
cache: &IndexCache,
|
||||||
|
rope: Rope,
|
||||||
|
k: usize,
|
||||||
|
n_genomes: usize,
|
||||||
|
n_partitions: usize,
|
||||||
|
with_counts: bool,
|
||||||
|
effective_z: usize,
|
||||||
|
detail: bool,
|
||||||
|
count_missing: bool,
|
||||||
|
force_presence: bool,
|
||||||
|
presence_threshold: u32,
|
||||||
|
genomes: &[GenomeInfo],
|
||||||
|
) -> Vec<u8> {
|
||||||
|
let chunk_start = Instant::now();
|
||||||
|
let chunk_bytes = rope.len();
|
||||||
|
|
||||||
|
let records = parse_chunk(&rope, k);
|
||||||
|
if records.is_empty() {
|
||||||
|
return Vec::new();
|
||||||
|
}
|
||||||
|
|
||||||
|
let batch = QueryBatch::from_records(records, k, 6, 0.7, n_partitions);
|
||||||
|
let n_seqs = batch.ids.len();
|
||||||
|
|
||||||
|
// Estimate QueryBatch::by_partition's actual memory footprint: the
|
||||||
|
// k-mer-level dedup map — one HashMap<CanonicalKmer, Vec<KmerDesc>> per
|
||||||
|
// partition, sized by *unique* k-mers, not shrunk by dedup. On real
|
||||||
|
// workloads with a low intra-chunk duplication rate this can dwarf every
|
||||||
|
// other per-chunk structure, including the sparse Findere ones logged
|
||||||
|
// further down. Measured by allocated capacity, not logical length, to
|
||||||
|
// reflect real memory pressure (HashMap/Vec growth slack) — `by_partition`
|
||||||
|
// is alive for the entire process_chunk call (never drained, only
|
||||||
|
// iterated by reference), so this is its footprint for the whole chunk
|
||||||
|
// lifetime, not a transient.
|
||||||
|
let hashmap_slot_bytes = (std::mem::size_of::<CanonicalKmer>()
|
||||||
|
+ std::mem::size_of::<Vec<crate::query_layer::KmerDesc>>()
|
||||||
|
+ 1) as u64; // +1 ≈ hashbrown control byte per slot
|
||||||
|
let by_partition_map_bytes: u64 = batch
|
||||||
|
.by_partition
|
||||||
|
.iter()
|
||||||
|
.map(|m| m.capacity() as u64 * hashmap_slot_bytes)
|
||||||
|
.sum();
|
||||||
|
let by_partition_desc_bytes: u64 = batch
|
||||||
|
.by_partition
|
||||||
|
.iter()
|
||||||
|
.flat_map(|m| m.values())
|
||||||
|
.map(|v| v.capacity() as u64 * std::mem::size_of::<crate::query_layer::KmerDesc>() as u64)
|
||||||
|
.sum();
|
||||||
|
let by_partition_bytes = by_partition_map_bytes + by_partition_desc_bytes;
|
||||||
|
|
||||||
|
debug!(
|
||||||
|
n_unique_kmers_total = batch.by_partition.iter().map(|m| m.len() as u64).sum::<u64>(),
|
||||||
|
by_partition_map_bytes,
|
||||||
|
by_partition_desc_bytes,
|
||||||
|
by_partition_bytes,
|
||||||
|
chunk_bytes,
|
||||||
|
"by_partition memory retained"
|
||||||
|
);
|
||||||
|
|
||||||
|
// Sparse bookkeeping for the whole chunk:
|
||||||
|
// - smer_index: O(total_smers) — is this s-mer in the index at all.
|
||||||
|
// - by_genome[g]: raw (seq_idx, pos_smer, value) hits for genome g, only
|
||||||
|
// ever containing nonzero entries (query_partition_with never emits a
|
||||||
|
// QueryHit::Value for a zero value) — empty for every genome this chunk
|
||||||
|
// never matched, which is the common case for unrelated queries.
|
||||||
|
let mut smer_index = SmerIndex::new(&batch.n_kmers);
|
||||||
|
let mut by_genome: Vec<Vec<(u32, u32, u32)>> = (0..n_genomes).map(|_| Vec::new()).collect();
|
||||||
|
|
||||||
|
// Dedup-ratio bookkeeping: occurrences (from batch.n_kmers, computed
|
||||||
|
// before dedup) vs. unique k-mers actually queried (query_stats) — the
|
||||||
|
// entire justification for k-mer-level dereplication. If this ratio
|
||||||
|
// stays close to 1.0 on real data, dereplication isn't paying for
|
||||||
|
// itself and that should show up here.
|
||||||
|
let n_occurrences: u64 = batch.n_kmers.iter().map(|&n| n as u64).sum();
|
||||||
|
let mut query_stats = QueryStats::default();
|
||||||
|
|
||||||
|
for (part_idx, kmers) in batch.by_partition.iter().enumerate() {
|
||||||
|
if kmers.is_empty() {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
let stats = cache.query_partition_with(part_idx, kmers, n_genomes, |event| match event {
|
||||||
|
QueryHit::Found(descs) => {
|
||||||
|
for desc in descs {
|
||||||
|
smer_index.mark_found(desc.seq_idx as usize, desc.pos as usize);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
QueryHit::Value(descs, g, v) => {
|
||||||
|
for desc in descs {
|
||||||
|
by_genome[g].push((desc.seq_idx, desc.pos, v));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
query_stats += stats;
|
||||||
|
}
|
||||||
|
|
||||||
|
debug!(
|
||||||
|
n_occurrences,
|
||||||
|
n_unique_kmers = query_stats.n_unique_kmers,
|
||||||
|
n_mphf_calls = query_stats.n_mphf_calls,
|
||||||
|
n_hits = query_stats.n_hits,
|
||||||
|
n_columns_scanned = query_stats.n_columns_scanned,
|
||||||
|
n_col_get_calls = query_stats.n_col_get_calls,
|
||||||
|
"k-mer dedup + column-major fetch"
|
||||||
|
);
|
||||||
|
|
||||||
|
// ── Sparse Findere: per-genome run detection + sliding-window minimum ────
|
||||||
|
//
|
||||||
|
// Confirmed z-windows, per genome, replace the dense win_min matrix:
|
||||||
|
// total retained memory is O(actual hits), not O(total_smers × n_genomes)
|
||||||
|
// — the whole point of this pass. See sparse_findere_for_genome's doc for
|
||||||
|
// why run detection is equivalent to the dense scan's semantics.
|
||||||
|
let presence = force_presence || !with_counts;
|
||||||
|
let threshold = presence_threshold;
|
||||||
|
let z = effective_z;
|
||||||
|
|
||||||
|
let n_kmers_out: Vec<usize> = batch
|
||||||
|
.n_kmers
|
||||||
|
.iter()
|
||||||
|
.map(|&n| {
|
||||||
|
let n = n as usize;
|
||||||
|
if n >= z { n - z + 1 } else { 0 }
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
let mut out_offsets = Vec::with_capacity(n_seqs + 1);
|
||||||
|
{
|
||||||
|
let mut total = 0usize;
|
||||||
|
out_offsets.push(0);
|
||||||
|
for &n in &n_kmers_out {
|
||||||
|
total += n;
|
||||||
|
out_offsets.push(total);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let total_out = *out_offsets.last().unwrap_or(&0);
|
||||||
|
|
||||||
|
let n_dense_would_be = n_occurrences as u64 * n_genomes as u64;
|
||||||
|
let mut n_sparse_entries = 0u64;
|
||||||
|
let mut n_runs_total = 0usize;
|
||||||
|
let mut run_len_total = 0usize;
|
||||||
|
|
||||||
|
let mut confirmed_by_genome: Vec<Vec<ConfirmedHit>> = Vec::with_capacity(n_genomes);
|
||||||
|
for hits in &mut by_genome {
|
||||||
|
n_sparse_entries += hits.len() as u64;
|
||||||
|
let (confirmed, n_runs, run_len) = sparse_findere_for_genome(hits, z, presence, threshold);
|
||||||
|
n_runs_total += n_runs;
|
||||||
|
run_len_total += run_len;
|
||||||
|
confirmed_by_genome.push(confirmed);
|
||||||
|
}
|
||||||
|
|
||||||
|
debug!(
|
||||||
|
n_dense_would_be,
|
||||||
|
n_sparse_entries,
|
||||||
|
n_runs = n_runs_total,
|
||||||
|
avg_run_len = if n_runs_total > 0 { run_len_total as f64 / n_runs_total as f64 } else { 0.0 },
|
||||||
|
z,
|
||||||
|
"sparse Findere"
|
||||||
|
);
|
||||||
|
|
||||||
|
// Actual bytes retained by the sparse hit structures (by_genome +
|
||||||
|
// confirmed_by_genome, both alive simultaneously at this point).
|
||||||
|
const HIT_ENTRY_BYTES: u64 = std::mem::size_of::<(u32, u32, u32)>() as u64;
|
||||||
|
let by_genome_bytes: u64 = by_genome.iter().map(|v| v.capacity() as u64 * HIT_ENTRY_BYTES).sum();
|
||||||
|
let confirmed_bytes: u64 = confirmed_by_genome.iter().map(|v| v.capacity() as u64 * HIT_ENTRY_BYTES).sum();
|
||||||
|
let retained_bytes = by_genome_bytes + confirmed_bytes;
|
||||||
|
|
||||||
|
debug!(
|
||||||
|
by_genome_bytes,
|
||||||
|
confirmed_bytes,
|
||||||
|
retained_bytes,
|
||||||
|
chunk_bytes,
|
||||||
|
empirical_multiplier = retained_bytes as f64 / chunk_bytes.max(1) as f64,
|
||||||
|
"sparse memory retained"
|
||||||
|
);
|
||||||
|
|
||||||
|
// ── Accumulate: genome totals (per genome, from confirmed hits) ──────────
|
||||||
|
let mut accs: Vec<SeqAcc> = (0..n_seqs).map(|_| SeqAcc::new(n_genomes)).collect();
|
||||||
|
let mut confirmed_any = vec![false; total_out];
|
||||||
|
|
||||||
|
for (g, hits) in confirmed_by_genome.iter().enumerate() {
|
||||||
|
for &(seq_idx, pos_out, c) in hits {
|
||||||
|
let abs_out = out_offsets[seq_idx as usize] + pos_out as usize;
|
||||||
|
confirmed_any[abs_out] = true;
|
||||||
|
accs[seq_idx as usize].genome_totals[g] += c;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Accumulate: kmer_count / kmer_missing (per position, genome-independent) ─
|
||||||
|
for seq_idx in 0..n_seqs {
|
||||||
|
let out_n = n_kmers_out[seq_idx];
|
||||||
|
let acc = &mut accs[seq_idx];
|
||||||
|
for pos in 0..out_n {
|
||||||
|
let abs_out = out_offsets[seq_idx] + pos;
|
||||||
|
if confirmed_any[abs_out] {
|
||||||
|
acc.kmer_count += 1;
|
||||||
|
} else if !smer_index.is_in_index(seq_idx, pos) {
|
||||||
|
acc.kmer_missing += 1;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Coverage (--detail): densify only when actually requested ────────────
|
||||||
|
let mut cov: Vec<Vec<Vec<u32>>> = if detail {
|
||||||
|
n_kmers_out.iter().map(|&n| vec![vec![0u32; n]; n_genomes]).collect()
|
||||||
|
} else {
|
||||||
|
Vec::new()
|
||||||
|
};
|
||||||
|
if detail {
|
||||||
|
for (g, hits) in confirmed_by_genome.iter().enumerate() {
|
||||||
|
for &(seq_idx, pos_out, c) in hits {
|
||||||
|
cov[seq_idx as usize][g][pos_out as usize] += c;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Capacity estimate: actual sequence + ID bytes, plus JSON overhead per record.
|
||||||
|
let seq_bytes: usize = batch.seqs.iter().map(|s| s.len()).sum();
|
||||||
|
let id_bytes: usize = batch.ids.iter().map(|s| s.len()).sum();
|
||||||
|
let cap = seq_bytes + id_bytes + n_seqs * (4 + 50 + n_genomes * 20) + 100;
|
||||||
|
let mut buf = Vec::with_capacity(cap);
|
||||||
|
emit_batch(
|
||||||
|
&batch,
|
||||||
|
&accs,
|
||||||
|
genomes,
|
||||||
|
count_missing,
|
||||||
|
detail,
|
||||||
|
&cov,
|
||||||
|
&mut buf,
|
||||||
|
);
|
||||||
|
|
||||||
|
debug!(
|
||||||
|
chunk_bytes,
|
||||||
|
n_seqs,
|
||||||
|
n_smers = batch.n_kmers.iter().map(|&n| n as u64).sum::<u64>(),
|
||||||
|
wall_ms = chunk_start.elapsed().as_millis() as u64,
|
||||||
|
"process_chunk"
|
||||||
|
);
|
||||||
|
|
||||||
|
buf
|
||||||
|
}
|
||||||
@@ -0,0 +1,75 @@
|
|||||||
|
use std::collections::VecDeque;
|
||||||
|
|
||||||
|
/// One confirmed z-window: genome `g`'s window ending at k-mer `pos` (the
|
||||||
|
/// *leftmost* s-mer of the window, i.e. the k_user-mer's output position) is
|
||||||
|
/// fully present and nonzero, with window-minimum `value`.
|
||||||
|
pub(crate) type ConfirmedHit = (u32, u32, u32); // (seq_idx, pos_out, value)
|
||||||
|
|
||||||
|
/// Reduce one genome's raw sparse s-mer hits — `(seq_idx, pos_smer, raw_value)`,
|
||||||
|
/// unsorted, exactly as delivered by `QueryHit::Value` — into confirmed
|
||||||
|
/// z-windows, without ever visiting a position that had no hit at all.
|
||||||
|
///
|
||||||
|
/// A z-window is confirmed only when all z s-mers in it are present *and*
|
||||||
|
/// nonzero for this genome (matching the dense sliding-window's semantics,
|
||||||
|
/// where "not in index" or a zero value both contribute 0 to the window
|
||||||
|
/// minimum) — which can only happen inside a maximal run of consecutive
|
||||||
|
/// `pos_smer` values for the same sequence. `hits` is sorted in place by
|
||||||
|
/// `(seq_idx, pos_smer)` to expose those runs; the monotone-deque
|
||||||
|
/// window-minimum then runs per run, on run-relative indices, identical in
|
||||||
|
/// spirit to the dense version's whole-sequence scan.
|
||||||
|
///
|
||||||
|
/// Returns the confirmed hits plus `(n_runs, total_run_len)` for logging —
|
||||||
|
/// a low average run length relative to `z` means most hits fail to form a
|
||||||
|
/// complete window.
|
||||||
|
pub(crate) fn sparse_findere_for_genome(
|
||||||
|
hits: &mut [(u32, u32, u32)],
|
||||||
|
z: usize,
|
||||||
|
presence: bool,
|
||||||
|
threshold: u32,
|
||||||
|
) -> (Vec<ConfirmedHit>, usize, usize) {
|
||||||
|
hits.sort_unstable_by_key(|&(seq, pos, _)| (seq, pos));
|
||||||
|
|
||||||
|
let mut confirmed = Vec::new();
|
||||||
|
let mut n_runs = 0usize;
|
||||||
|
let mut total_run_len = 0usize;
|
||||||
|
let mut dq: VecDeque<(usize, u32)> = VecDeque::new(); // (run-relative index, value)
|
||||||
|
|
||||||
|
let mut i = 0;
|
||||||
|
while i < hits.len() {
|
||||||
|
let seq = hits[i].0;
|
||||||
|
let mut j = i + 1;
|
||||||
|
while j < hits.len() && hits[j].0 == seq && hits[j].1 == hits[j - 1].1 + 1 {
|
||||||
|
j += 1;
|
||||||
|
}
|
||||||
|
let run = &hits[i..j];
|
||||||
|
n_runs += 1;
|
||||||
|
total_run_len += run.len();
|
||||||
|
|
||||||
|
dq.clear();
|
||||||
|
for (k, &(_, pos, val)) in run.iter().enumerate() {
|
||||||
|
while dq.back().map_or(false, |&(_, v)| v >= val) {
|
||||||
|
dq.pop_back();
|
||||||
|
}
|
||||||
|
dq.push_back((k, val));
|
||||||
|
while dq.front().map_or(false, |&(fk, _)| fk + z <= k) {
|
||||||
|
dq.pop_front();
|
||||||
|
}
|
||||||
|
if k + 1 >= z {
|
||||||
|
let win_min = dq.front().unwrap().1;
|
||||||
|
if win_min > 0 {
|
||||||
|
let pos_out = pos + 1 - z as u32;
|
||||||
|
let c = if presence { u32::from(win_min >= threshold) } else { win_min };
|
||||||
|
confirmed.push((seq, pos_out, c));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
i = j;
|
||||||
|
}
|
||||||
|
|
||||||
|
(confirmed, n_runs, total_run_len)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
#[path = "tests/findere.rs"]
|
||||||
|
mod tests;
|
||||||
@@ -1,10 +1,16 @@
|
|||||||
//! Query-side operations on an `obikindex::KmerIndex`: staging ground for
|
//! Query-side operations on an `obikidxcache::IndexCache`: matching query
|
||||||
//! code currently living in `obikindex::index`'s query path, to be migrated
|
//! sequences' k-mers against an already-built `obikindex::KmerIndex` and
|
||||||
//! here to lighten that crate — see `DevDocMD/` for the rationale. Kept as
|
//! formatting the result. Kept as a separate crate — not a module of
|
||||||
//! a separate crate — not a module of `obikindex` — so the dependency runs
|
//! `obikindex` — so the dependency runs one way only (query code depends on
|
||||||
//! one way only (query code depends on the data model, never the reverse),
|
//! the data model, never the reverse), same pattern as `obikindexer` for the
|
||||||
//! same pattern as `obikindexer` for the build side.
|
//! build side.
|
||||||
|
|
||||||
mod query_layer;
|
mod query_layer;
|
||||||
|
mod batch;
|
||||||
|
mod chunk;
|
||||||
|
mod findere;
|
||||||
|
mod smer_index;
|
||||||
|
mod output;
|
||||||
|
|
||||||
pub use query_layer::{KmerDesc, QueryHit, QueryStats};
|
pub use query_layer::{KmerDesc, QueryHit, QueryPartition, QueryStats};
|
||||||
|
pub use chunk::process_chunk;
|
||||||
|
|||||||
@@ -0,0 +1,52 @@
|
|||||||
|
use std::io::Write;
|
||||||
|
|
||||||
|
use obikindex::GenomeInfo;
|
||||||
|
|
||||||
|
use crate::batch::QueryBatch;
|
||||||
|
use crate::chunk::SeqAcc;
|
||||||
|
|
||||||
|
pub(crate) fn emit_batch(
|
||||||
|
batch: &QueryBatch,
|
||||||
|
accs: &[SeqAcc],
|
||||||
|
genomes: &[GenomeInfo],
|
||||||
|
count_missing: bool,
|
||||||
|
detail: bool,
|
||||||
|
cov: &[Vec<Vec<u32>>],
|
||||||
|
out: &mut impl Write,
|
||||||
|
) {
|
||||||
|
for (seq_idx, (id, seq)) in batch.ids.iter().zip(batch.seqs.iter()).enumerate() {
|
||||||
|
let acc = &accs[seq_idx];
|
||||||
|
let mut ann = serde_json::Map::new();
|
||||||
|
|
||||||
|
ann.insert("kmer_count".into(), acc.kmer_count.into());
|
||||||
|
if count_missing {
|
||||||
|
ann.insert("kmer_missing".into(), acc.kmer_missing.into());
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut match_map = serde_json::Map::new();
|
||||||
|
for (g, genome) in genomes.iter().enumerate() {
|
||||||
|
if acc.genome_totals[g] != 0 {
|
||||||
|
match_map.insert(genome.label.clone(), acc.genome_totals[g].into());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
ann.insert("kmer_strict_matches".into(), match_map.into());
|
||||||
|
|
||||||
|
if detail && !cov.is_empty() {
|
||||||
|
let mut cov_map = serde_json::Map::new();
|
||||||
|
for (g, genome) in genomes.iter().enumerate() {
|
||||||
|
let v: Vec<serde_json::Value> = cov[seq_idx][g].iter().map(|&x| x.into()).collect();
|
||||||
|
cov_map.insert(genome.label.clone(), v.into());
|
||||||
|
}
|
||||||
|
ann.insert("coverage".into(), cov_map.into());
|
||||||
|
}
|
||||||
|
|
||||||
|
// OBITools4 FASTA format: >id {"key":value,...}
|
||||||
|
let _ = out.write_all(b">");
|
||||||
|
let _ = out.write_all(id.as_bytes());
|
||||||
|
let _ = out.write_all(b" ");
|
||||||
|
let _ = serde_json::to_writer(&mut *out, &ann);
|
||||||
|
let _ = out.write_all(b"\n");
|
||||||
|
let _ = out.write_all(seq);
|
||||||
|
let _ = out.write_all(b"\n");
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -1,78 +1,7 @@
|
|||||||
use std::collections::HashMap;
|
use std::collections::HashMap;
|
||||||
use std::path::Path;
|
|
||||||
|
|
||||||
use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix};
|
use obikidxcache::index_cache::IndexCache;
|
||||||
use obikseq::CanonicalKmer;
|
use obikseq::CanonicalKmer;
|
||||||
use obikindex::layer::MphfLayer;
|
|
||||||
use obikindex::{OKIError, OKIResult};
|
|
||||||
|
|
||||||
use obikindex::KmerIndex;
|
|
||||||
|
|
||||||
// ── per-layer query handle ────────────────────────────────────────────────────
|
|
||||||
|
|
||||||
enum QueryLayer {
|
|
||||||
Presence(MphfLayer, PersistentBitMatrix),
|
|
||||||
Count(MphfLayer, PersistentCompactIntMatrix),
|
|
||||||
}
|
|
||||||
|
|
||||||
impl QueryLayer {
|
|
||||||
fn open(layer_dir: &Path, with_counts: bool) -> OKIResult<Self> {
|
|
||||||
let mphf = MphfLayer::open(layer_dir)?;
|
|
||||||
let counts_dir = layer_dir.join("counts");
|
|
||||||
let presence_dir = layer_dir.join("presence");
|
|
||||||
|
|
||||||
if with_counts && counts_dir.exists() {
|
|
||||||
let mat = PersistentCompactIntMatrix::open(layer_dir).map_err(OKIError::Io)?;
|
|
||||||
Ok(QueryLayer::Count(mphf, mat))
|
|
||||||
} else if presence_dir.exists() || !counts_dir.exists() {
|
|
||||||
// presence mode, or no matrix at all → Implicit handled inside open()
|
|
||||||
let mat = PersistentBitMatrix::open(layer_dir).map_err(OKIError::Io)?;
|
|
||||||
Ok(QueryLayer::Presence(mphf, mat))
|
|
||||||
} else {
|
|
||||||
// counts exist but not presence — count layer, no presence requested
|
|
||||||
let mat = PersistentCompactIntMatrix::open(layer_dir).map_err(OKIError::Io)?;
|
|
||||||
Ok(QueryLayer::Count(mphf, mat))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// MPHF lookup only — no matrix access. `Some(slot)` on hit.
|
|
||||||
fn find_slot(&self, kmer: CanonicalKmer) -> Option<usize> {
|
|
||||||
match self {
|
|
||||||
QueryLayer::Presence(mphf, _) | QueryLayer::Count(mphf, _) => mphf.find(kmer),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Number of genome columns this layer's matrix actually has. Bounds
|
|
||||||
/// column-major iteration — usually equal to the index's `n_genomes`, but
|
|
||||||
/// `PersistentBitMatrix::Implicit` (the documented mono-genome fast path)
|
|
||||||
/// always reports exactly `1`, regardless of the index's real genome
|
|
||||||
/// count, so callers must use this rather than assuming `n_genomes`.
|
|
||||||
fn n_cols(&self) -> usize {
|
|
||||||
match self {
|
|
||||||
QueryLayer::Presence(_, mat) => mat.n_cols(),
|
|
||||||
QueryLayer::Count(_, mat) => mat.n_cols(),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Every nonzero `(idx into slots, col, value)` triple among `slots`.
|
|
||||||
/// Format-agnostic: each matrix picks its own natural traversal
|
|
||||||
/// (`PersistentBitMatrix::nonzero_iter` dispatches to a genuinely
|
|
||||||
/// row-major decode on `Sparse`, not a column-major point-probe loop —
|
|
||||||
/// see `DevDocMD/architecture/siblings.md`, "`query` never benefits
|
|
||||||
/// from sparse row-major access"). Replaces the old per-`(genome,
|
|
||||||
/// slot)` `col_value` point lookup, which this layer's `Sparse`
|
|
||||||
/// presence matrices paid for badly: each such lookup rebuilt the
|
|
||||||
/// entire row just to return one cell.
|
|
||||||
fn nonzero_iter<'a>(
|
|
||||||
&'a self,
|
|
||||||
slots: &'a [usize],
|
|
||||||
) -> Box<dyn Iterator<Item = (usize, usize, u32)> + 'a> {
|
|
||||||
match self {
|
|
||||||
QueryLayer::Presence(_, mat) => mat.nonzero_iter(slots),
|
|
||||||
QueryLayer::Count(_, mat) => Box::new(mat.nonzero_iter(slots)),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// ── KmerDesc — one occurrence of a k-mer in the query batch ──────────────────
|
// ── KmerDesc — one occurrence of a k-mer in the query batch ──────────────────
|
||||||
|
|
||||||
@@ -85,7 +14,7 @@ pub struct KmerDesc {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Aggregate counters for one `query_partition_with` call — feeds the
|
/// Aggregate counters for one `query_partition_with` call — feeds the
|
||||||
/// dedup-ratio and column-scan logging in `obikmer::cmd::query` (occurrences
|
/// dedup-ratio and column-scan logging in `obikquery::chunk` (occurrences
|
||||||
/// vs. unique k-mers is the whole justification for k-mer-level
|
/// vs. unique k-mers is the whole justification for k-mer-level
|
||||||
/// dereplication; columns scanned / `get()` calls quantify the column-major
|
/// dereplication; columns scanned / `get()` calls quantify the column-major
|
||||||
/// fetch's locality claim).
|
/// fetch's locality claim).
|
||||||
@@ -93,7 +22,7 @@ pub struct KmerDesc {
|
|||||||
pub struct QueryStats {
|
pub struct QueryStats {
|
||||||
/// Distinct canonical k-mers queried in this partition.
|
/// Distinct canonical k-mers queried in this partition.
|
||||||
pub n_unique_kmers: usize,
|
pub n_unique_kmers: usize,
|
||||||
/// Total `MphfLayer::find` calls issued (a k-mer tried against more than
|
/// Total MPHF-membership checks issued (a k-mer tried against more than
|
||||||
/// one layer before a hit, or against all layers on a miss, counts once
|
/// one layer before a hit, or against all layers on a miss, counts once
|
||||||
/// per layer attempted).
|
/// per layer attempted).
|
||||||
pub n_mphf_calls: usize,
|
pub n_mphf_calls: usize,
|
||||||
@@ -102,7 +31,7 @@ pub struct QueryStats {
|
|||||||
/// Total genome columns scanned across all hit layers (sum of
|
/// Total genome columns scanned across all hit layers (sum of
|
||||||
/// `layer.n_cols()` over layers with at least one hit).
|
/// `layer.n_cols()` over layers with at least one hit).
|
||||||
pub n_columns_scanned: usize,
|
pub n_columns_scanned: usize,
|
||||||
/// Total `col_value` calls issued during the column-major fetch pass
|
/// Total nonzero-cell fetches issued during the column-major fetch pass
|
||||||
/// (`n_columns_scanned` × hits-per-layer, summed over layers).
|
/// (`n_columns_scanned` × hits-per-layer, summed over layers).
|
||||||
pub n_col_get_calls: usize,
|
pub n_col_get_calls: usize,
|
||||||
}
|
}
|
||||||
@@ -119,7 +48,7 @@ impl std::ops::AddAssign for QueryStats {
|
|||||||
|
|
||||||
// ── QueryHit — one event delivered to query_partition_with's callback ───────
|
// ── QueryHit — one event delivered to query_partition_with's callback ───────
|
||||||
|
|
||||||
/// One event from [`KmerPartition::query_partition_with`]'s two-stage query:
|
/// One event from [`QueryPartition::query_partition_with`]'s two-stage query:
|
||||||
/// a `Found` event once per hit k-mer (stage 1, MPHF-only — mark the k-mer as
|
/// a `Found` event once per hit k-mer (stage 1, MPHF-only — mark the k-mer as
|
||||||
/// indexed regardless of any genome's value), then a `Value` event per
|
/// indexed regardless of any genome's value), then a `Value` event per
|
||||||
/// `(hit k-mer, genome)` pair with a nonzero matrix value (stage 2,
|
/// `(hit k-mer, genome)` pair with a nonzero matrix value (stage 2,
|
||||||
@@ -131,9 +60,21 @@ pub enum QueryHit<'a> {
|
|||||||
Value(&'a [KmerDesc], usize, u32),
|
Value(&'a [KmerDesc], usize, u32),
|
||||||
}
|
}
|
||||||
|
|
||||||
// ── KmerPartition::query_partition_with ──────────────────────────────────────
|
// ── QueryPartition — extension trait over obikidxcache::IndexCache ─────────
|
||||||
|
|
||||||
impl KmerIndex {
|
/// Query-side operations on an already-opened [`IndexCache`] — `IndexCache`
|
||||||
|
/// is a foreign type (`obikidxcache`), so this is an extension trait rather
|
||||||
|
/// than an inherent `impl` (`obikindex`'s original `impl KmerIndex { .. }`
|
||||||
|
/// shape doesn't compile outside `obikindex` itself — the orphan rule).
|
||||||
|
///
|
||||||
|
/// Built on `IndexCache` rather than re-opening each layer's MPHF/matrix
|
||||||
|
/// files per call (the previous, per-call-`QueryLayer::open` design): every
|
||||||
|
/// layer touched by a query is already open, for the whole lifetime of the
|
||||||
|
/// cache, so a query pass costs zero disk I/O beyond the first chunk. This
|
||||||
|
/// also drops the `with_counts` parameter the old design needed: whether a
|
||||||
|
/// layer holds a count or a presence matrix is no longer guessed from an
|
||||||
|
/// index-wide flag, it's read straight off `KmerLayer`'s own variant.
|
||||||
|
pub trait QueryPartition {
|
||||||
/// Query a single partition for a pre-deduplicated map of canonical
|
/// Query a single partition for a pre-deduplicated map of canonical
|
||||||
/// k-mers → their occurrences (`seq_idx`, `pos`) in the query batch.
|
/// k-mers → their occurrences (`seq_idx`, `pos`) in the query batch.
|
||||||
///
|
///
|
||||||
@@ -151,43 +92,54 @@ impl KmerIndex {
|
|||||||
/// matrix formats are column-oriented on disk (one `mmap`'d region per
|
/// matrix formats are column-oriented on disk (one `mmap`'d region per
|
||||||
/// genome), so scanning one column at a time touches far fewer
|
/// genome), so scanning one column at a time touches far fewer
|
||||||
/// distinct mmap regions than fetching one full row per hit.
|
/// distinct mmap regions than fetching one full row per hit.
|
||||||
pub fn query_partition_with<F>(
|
///
|
||||||
|
/// Infallible: every layer this can touch was already opened (and
|
||||||
|
/// validated) when the cache was built — see `IndexCache::new`'s own
|
||||||
|
/// panic-on-failure contract. A `part_idx` this cache doesn't hold, or
|
||||||
|
/// an empty `kmers` map, both just return default (all-zero) stats.
|
||||||
|
fn query_partition_with<F>(
|
||||||
|
&self,
|
||||||
|
part_idx: usize,
|
||||||
|
kmers: &HashMap<CanonicalKmer, Vec<KmerDesc>>,
|
||||||
|
n_genomes: usize,
|
||||||
|
on_event: F,
|
||||||
|
) -> QueryStats
|
||||||
|
where
|
||||||
|
F: FnMut(QueryHit);
|
||||||
|
}
|
||||||
|
|
||||||
|
impl QueryPartition for IndexCache {
|
||||||
|
fn query_partition_with<F>(
|
||||||
&self,
|
&self,
|
||||||
part_idx: usize,
|
part_idx: usize,
|
||||||
kmers: &HashMap<CanonicalKmer, Vec<KmerDesc>>,
|
kmers: &HashMap<CanonicalKmer, Vec<KmerDesc>>,
|
||||||
n_genomes: usize,
|
n_genomes: usize,
|
||||||
with_counts: bool,
|
|
||||||
mut on_event: F,
|
mut on_event: F,
|
||||||
) -> OKIResult<QueryStats>
|
) -> QueryStats
|
||||||
where
|
where
|
||||||
F: FnMut(QueryHit),
|
F: FnMut(QueryHit),
|
||||||
{
|
{
|
||||||
let mut stats = QueryStats::default();
|
let mut stats = QueryStats::default();
|
||||||
|
|
||||||
if kmers.is_empty() {
|
if kmers.is_empty() {
|
||||||
return Ok(stats);
|
return stats;
|
||||||
}
|
}
|
||||||
|
|
||||||
let index_dir = self.index_dir(part_idx);
|
let n_layer = match self.n_layer(part_idx) {
|
||||||
if !index_dir.exists() {
|
Some(n) if n > 0 => n,
|
||||||
return Ok(stats);
|
_ => return stats,
|
||||||
}
|
};
|
||||||
|
|
||||||
let meta = self.partition_meta(part_idx)?;
|
|
||||||
let layers: Vec<QueryLayer> = (0..meta.n_layers)
|
|
||||||
.map(|i| QueryLayer::open(&self.layer_dir(part_idx, i), with_counts))
|
|
||||||
.collect::<OKIResult<_>>()?;
|
|
||||||
|
|
||||||
// ── Stage 1: MPHF-only pass, bucket hits by (layer_idx, slot) ────────
|
// ── Stage 1: MPHF-only pass, bucket hits by (layer_idx, slot) ────────
|
||||||
let mut by_layer: Vec<HashMap<usize, &Vec<KmerDesc>>> =
|
let mut by_layer: Vec<HashMap<usize, &Vec<KmerDesc>>> =
|
||||||
(0..layers.len()).map(|_| HashMap::new()).collect();
|
(0..n_layer).map(|_| HashMap::new()).collect();
|
||||||
|
|
||||||
for (kmer, descs) in kmers {
|
for (kmer, descs) in kmers {
|
||||||
stats.n_unique_kmers += 1;
|
stats.n_unique_kmers += 1;
|
||||||
for (layer_idx, layer) in layers.iter().enumerate() {
|
for l in 0..n_layer {
|
||||||
stats.n_mphf_calls += 1;
|
stats.n_mphf_calls += 1;
|
||||||
if let Some(slot) = layer.find_slot(*kmer) {
|
if let Some(slot) = self.find_in_layer(part_idx, l, *kmer) {
|
||||||
by_layer[layer_idx].insert(slot, descs);
|
by_layer[l].insert(slot, descs);
|
||||||
on_event(QueryHit::Found(descs));
|
on_event(QueryHit::Found(descs));
|
||||||
stats.n_hits += 1;
|
stats.n_hits += 1;
|
||||||
break;
|
break;
|
||||||
@@ -196,16 +148,13 @@ impl KmerIndex {
|
|||||||
}
|
}
|
||||||
|
|
||||||
// ── Stage 2: nonzero-cell fetch, per layer ────────────────────────────
|
// ── Stage 2: nonzero-cell fetch, per layer ────────────────────────────
|
||||||
// Format-agnostic — see `QueryLayer::nonzero_iter`. `n_cols` still
|
for (l, slots) in by_layer.iter().enumerate() {
|
||||||
// bounds accepted genome columns (Implicit reports fewer than
|
|
||||||
// `n_genomes`; see `n_cols`'s doc), cells beyond it are dropped
|
|
||||||
// rather than ever produced, since `nonzero_iter` only knows the
|
|
||||||
// matrix's own column count, not the caller's `n_genomes`.
|
|
||||||
for (layer_idx, slots) in by_layer.iter().enumerate() {
|
|
||||||
if slots.is_empty() {
|
if slots.is_empty() {
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
let layer = &layers[layer_idx];
|
let layer = self
|
||||||
|
.get_layer(part_idx, l)
|
||||||
|
.expect("layer within n_layer(part_idx)");
|
||||||
let n_cols = layer.n_cols().min(n_genomes);
|
let n_cols = layer.n_cols().min(n_genomes);
|
||||||
stats.n_columns_scanned += n_cols;
|
stats.n_columns_scanned += n_cols;
|
||||||
|
|
||||||
@@ -221,7 +170,7 @@ impl KmerIndex {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
Ok(stats)
|
stats
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,41 @@
|
|||||||
|
/// Tracks, per (sequence, s-mer position), whether the k-mer was found in the
|
||||||
|
/// index at all — independent of *which* genome(s) matched. Sized
|
||||||
|
/// `total_smers` (one `bool` per s-mer occurrence in the chunk), **not**
|
||||||
|
/// multiplied by `n_genomes`: this is the O(1)-per-position bookkeeping that
|
||||||
|
/// `kmer_missing` needs (the leftmost-s-mer-of-window membership test), kept
|
||||||
|
/// dense because it's already cheap — the `n_genomes`-scaled data lives in
|
||||||
|
/// the sparse per-genome hit lists built alongside it (see `chunk::process_chunk`).
|
||||||
|
pub(crate) struct SmerIndex {
|
||||||
|
in_index: Vec<bool>, // total_smers
|
||||||
|
offsets: Vec<usize>, // offsets[i]..offsets[i+1] = s-mer range for sequence i
|
||||||
|
}
|
||||||
|
|
||||||
|
impl SmerIndex {
|
||||||
|
pub(crate) fn new(n_kmers_per_seq: &[u32]) -> Self {
|
||||||
|
let mut offsets = Vec::with_capacity(n_kmers_per_seq.len() + 1);
|
||||||
|
let mut total = 0usize;
|
||||||
|
offsets.push(0);
|
||||||
|
for &n in n_kmers_per_seq {
|
||||||
|
total += n as usize;
|
||||||
|
offsets.push(total);
|
||||||
|
}
|
||||||
|
Self {
|
||||||
|
in_index: vec![false; total],
|
||||||
|
offsets,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Mark the k-mer at (seq, kmer) as found in the index — independent of
|
||||||
|
/// any particular genome's value. Called once per hit k-mer (stage 1 of
|
||||||
|
/// `query_partition_with`), regardless of how the column-major fetch
|
||||||
|
/// (stage 2) later reports per-genome values.
|
||||||
|
pub(crate) fn mark_found(&mut self, seq: usize, kmer: usize) {
|
||||||
|
let abs = self.offsets[seq] + kmer;
|
||||||
|
self.in_index[abs] = true;
|
||||||
|
}
|
||||||
|
|
||||||
|
#[inline]
|
||||||
|
pub(crate) fn is_in_index(&self, seq: usize, kmer: usize) -> bool {
|
||||||
|
self.in_index[self.offsets[seq] + kmer]
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,127 @@
|
|||||||
|
use super::*;
|
||||||
|
use obikrope::Rope;
|
||||||
|
use obikseq::CanonicalKmer;
|
||||||
|
use obiread::record::parse_chunk;
|
||||||
|
|
||||||
|
const K: usize = 11;
|
||||||
|
const M: usize = 5;
|
||||||
|
|
||||||
|
/// Build a `QueryBatch` from raw FASTA text, going through the same
|
||||||
|
/// `Rope` + `parse_chunk` path `process_chunk` uses — avoids hand-building a
|
||||||
|
/// `normalized` `Rope`, which is an implementation detail of `obiread`.
|
||||||
|
///
|
||||||
|
/// `obikseq`'s global K/M params are thread-local under `test-utils` (see
|
||||||
|
/// `obikseq::params`), so setting them here is per-test-thread and does not
|
||||||
|
/// need coordination with other tests.
|
||||||
|
fn batch_from_fasta(fasta: &str, k: usize, n_partitions: usize) -> QueryBatch {
|
||||||
|
obikseq::set_k(k);
|
||||||
|
obikseq::set_m(M);
|
||||||
|
let mut rope = Rope::new(Some("text/fasta"));
|
||||||
|
rope.push(fasta.as_bytes().to_vec());
|
||||||
|
let records = parse_chunk(&rope, k);
|
||||||
|
QueryBatch::from_records(records, k, 6, 0.7, n_partitions)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn total_occurrences(batch: &QueryBatch) -> u64 {
|
||||||
|
batch.n_kmers.iter().map(|&n| n as u64).sum()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn total_unique_kmers(batch: &QueryBatch) -> u64 {
|
||||||
|
batch.by_partition.iter().map(|m| m.len() as u64).sum()
|
||||||
|
}
|
||||||
|
|
||||||
|
// A 60 bp sequence, arbitrary but fixed — no attempt is made to prove it is
|
||||||
|
// free of internal k=11 repeats; the tests below only rely on inequalities
|
||||||
|
// that hold regardless (see each test's comment).
|
||||||
|
const SEQ: &str = "CATTAGCGTACCTGATCAGGTTACAGCTTAGGCATCCAGTTGACCATGACTGGACTTAGC";
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn single_sequence_yields_plausible_kmer_counts() {
|
||||||
|
// A single record can still contain internal repeats (SEQ isn't
|
||||||
|
// guaranteed repeat-free at k=11) — this only checks the batch is
|
||||||
|
// internally consistent, not a specific dedup ratio. The cross-record
|
||||||
|
// tests below make the actual, unconditional dedup claims.
|
||||||
|
let fasta = format!(">r1\n{SEQ}\n");
|
||||||
|
let batch = batch_from_fasta(&fasta, K, 1);
|
||||||
|
|
||||||
|
assert_eq!(batch.ids, vec!["r1".to_string()]);
|
||||||
|
let occurrences = total_occurrences(&batch);
|
||||||
|
let unique = total_unique_kmers(&batch);
|
||||||
|
assert!(occurrences > 0, "sequence should yield at least one k-mer");
|
||||||
|
assert!(unique > 0 && unique <= occurrences);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn duplicated_sequence_across_records_deduplicates() {
|
||||||
|
// Two records with byte-identical sequences: every k-mer in record 1
|
||||||
|
// exactly duplicates one in record 0, so unique kmers <= n_kmers[0],
|
||||||
|
// strictly less than the summed occurrences (2 * n_kmers[0]) as long as
|
||||||
|
// the sequence yields at least one k-mer. This holds regardless of
|
||||||
|
// whether SEQ has internal repeats.
|
||||||
|
let fasta = format!(">r1\n{SEQ}\n>r2\n{SEQ}\n");
|
||||||
|
let batch = batch_from_fasta(&fasta, K, 1);
|
||||||
|
|
||||||
|
assert_eq!(batch.ids.len(), 2);
|
||||||
|
let occurrences = total_occurrences(&batch);
|
||||||
|
let unique = total_unique_kmers(&batch);
|
||||||
|
|
||||||
|
assert!(batch.n_kmers[0] > 0);
|
||||||
|
assert_eq!(occurrences, batch.n_kmers[0] as u64 + batch.n_kmers[1] as u64);
|
||||||
|
assert!(
|
||||||
|
unique <= batch.n_kmers[0] as u64,
|
||||||
|
"identical sequences must not produce more unique k-mers than one copy has"
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
unique < occurrences,
|
||||||
|
"k-mer-level dedup must collapse at least the cross-record duplication"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn duplicated_sequence_broadcasts_to_both_seq_indices() {
|
||||||
|
// Stronger than the ratio check above: pick any k-mer that hit in both
|
||||||
|
// records and confirm its occurrence list actually references both
|
||||||
|
// seq_idx 0 and seq_idx 1 — this is the specific new capability (dedup
|
||||||
|
// reaching across records/superkmers), not just a smaller unique count.
|
||||||
|
let fasta = format!(">r1\n{SEQ}\n>r2\n{SEQ}\n");
|
||||||
|
let batch = batch_from_fasta(&fasta, K, 1);
|
||||||
|
|
||||||
|
let shared = batch.by_partition[0]
|
||||||
|
.values()
|
||||||
|
.find(|descs| descs.iter().any(|d| d.seq_idx == 0) && descs.iter().any(|d| d.seq_idx == 1));
|
||||||
|
|
||||||
|
assert!(
|
||||||
|
shared.is_some(),
|
||||||
|
"expected at least one k-mer shared between the two identical records"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn empty_records_yield_empty_batch() {
|
||||||
|
let batch = batch_from_fasta("", K, 1);
|
||||||
|
assert!(batch.ids.is_empty());
|
||||||
|
assert_eq!(total_occurrences(&batch), 0);
|
||||||
|
assert_eq!(total_unique_kmers(&batch), 0);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn partition_routing_is_a_pure_function_of_the_kmer() {
|
||||||
|
// With n_partitions=4, every occurrence of a given k-mer must land in
|
||||||
|
// the same partition bucket as every other occurrence of that k-mer
|
||||||
|
// (partition routing is derived from the minimizer, shared by
|
||||||
|
// definition among instances of the same k-mer's containing superkmer
|
||||||
|
// in this test's single-sequence-pair setup).
|
||||||
|
let fasta = format!(">r1\n{SEQ}\n>r2\n{SEQ}\n");
|
||||||
|
let batch = batch_from_fasta(&fasta, K, 4);
|
||||||
|
|
||||||
|
let total_unique: u64 = batch.by_partition.iter().map(|m| m.len() as u64).sum();
|
||||||
|
assert!(total_unique > 0);
|
||||||
|
|
||||||
|
// No k-mer key appears in more than one partition's map.
|
||||||
|
let mut seen: std::collections::HashSet<CanonicalKmer> = std::collections::HashSet::new();
|
||||||
|
for map in &batch.by_partition {
|
||||||
|
for kmer in map.keys() {
|
||||||
|
assert!(seen.insert(*kmer), "k-mer routed to more than one partition");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,105 @@
|
|||||||
|
use super::*;
|
||||||
|
|
||||||
|
// ── sparse_findere_for_genome vs. a dense reference implementation ──────────
|
||||||
|
//
|
||||||
|
// No property-testing crate (proptest/quickcheck) is a workspace dependency
|
||||||
|
// (checked before writing this — not adding one for a single test module,
|
||||||
|
// per this project's dependency-approval rule). A tiny deterministic xorshift
|
||||||
|
// PRNG, std-only, stands in for one.
|
||||||
|
|
||||||
|
/// Faithful reimplementation of the pre-phase-5 dense sliding-window scan —
|
||||||
|
/// the algorithm `sparse_findere_for_genome` replaced — used here only as a
|
||||||
|
/// correctness oracle, not in production code. Operates on one genome's
|
||||||
|
/// hits across possibly many sequences, exactly like the sparse version.
|
||||||
|
fn dense_reference_findere(
|
||||||
|
hits: &[(u32, u32, u32)],
|
||||||
|
seq_lens: &[usize],
|
||||||
|
z: usize,
|
||||||
|
presence: bool,
|
||||||
|
threshold: u32,
|
||||||
|
) -> Vec<(u32, u32, u32)> {
|
||||||
|
let mut by_seq: Vec<Vec<u32>> = seq_lens.iter().map(|&n| vec![0u32; n]).collect();
|
||||||
|
for &(seq, pos, val) in hits {
|
||||||
|
by_seq[seq as usize][pos as usize] = val;
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut confirmed = Vec::new();
|
||||||
|
for (seq_idx, values) in by_seq.iter().enumerate() {
|
||||||
|
let n = values.len();
|
||||||
|
let mut dq: std::collections::VecDeque<(usize, u32)> = std::collections::VecDeque::new();
|
||||||
|
for i in 0..n {
|
||||||
|
let v_i = values[i];
|
||||||
|
while dq.front().map_or(false, |&(f, _)| f + z <= i) {
|
||||||
|
dq.pop_front();
|
||||||
|
}
|
||||||
|
while dq.back().map_or(false, |&(_, v)| v >= v_i) {
|
||||||
|
dq.pop_back();
|
||||||
|
}
|
||||||
|
dq.push_back((i, v_i));
|
||||||
|
if i + 1 >= z {
|
||||||
|
let win_min = dq.front().unwrap().1;
|
||||||
|
if win_min > 0 {
|
||||||
|
let pos_out = (i + 1 - z) as u32;
|
||||||
|
let c = if presence { u32::from(win_min >= threshold) } else { win_min };
|
||||||
|
confirmed.push((seq_idx as u32, pos_out, c));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
confirmed
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Minimal std-only xorshift64 PRNG — deterministic, seedable, no dependency.
|
||||||
|
struct Xorshift64(u64);
|
||||||
|
impl Xorshift64 {
|
||||||
|
fn next(&mut self) -> u64 {
|
||||||
|
self.0 ^= self.0 << 13;
|
||||||
|
self.0 ^= self.0 >> 7;
|
||||||
|
self.0 ^= self.0 << 17;
|
||||||
|
self.0
|
||||||
|
}
|
||||||
|
fn range(&mut self, n: u32) -> u32 {
|
||||||
|
(self.next() % n as u64) as u32
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn sparse_findere_matches_dense_reference_on_random_inputs() {
|
||||||
|
let mut rng = Xorshift64(0x5eed_5eed_5eed_5eedu64);
|
||||||
|
|
||||||
|
for case in 0..200 {
|
||||||
|
let n_seqs = 1 + rng.range(4) as usize;
|
||||||
|
let seq_lens: Vec<usize> = (0..n_seqs).map(|_| 1 + rng.range(30) as usize).collect();
|
||||||
|
let z = 1 + rng.range(4) as usize;
|
||||||
|
let presence = rng.range(2) == 0;
|
||||||
|
let threshold = 1 + rng.range(3);
|
||||||
|
|
||||||
|
// Sparse density varies across cases, including edge cases (empty,
|
||||||
|
// fully dense) — deliberately not uniform, to stress both few-hits
|
||||||
|
// and many-overlapping-runs scenarios.
|
||||||
|
let density = rng.range(101);
|
||||||
|
let mut hits: Vec<(u32, u32, u32)> = Vec::new();
|
||||||
|
for (seq_idx, &len) in seq_lens.iter().enumerate() {
|
||||||
|
for pos in 0..len {
|
||||||
|
if rng.range(100) < density {
|
||||||
|
let val = 1 + rng.range(5); // never 0 — matches QueryHit::Value's invariant
|
||||||
|
hits.push((seq_idx as u32, pos as u32, val));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut sparse_input = hits.clone();
|
||||||
|
let (mut sparse_result, _, _) =
|
||||||
|
sparse_findere_for_genome(&mut sparse_input, z, presence, threshold);
|
||||||
|
let mut dense_result = dense_reference_findere(&hits, &seq_lens, z, presence, threshold);
|
||||||
|
|
||||||
|
sparse_result.sort_unstable();
|
||||||
|
dense_result.sort_unstable();
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
sparse_result, dense_result,
|
||||||
|
"case {case}: n_seqs={n_seqs} seq_lens={seq_lens:?} z={z} presence={presence} \
|
||||||
|
threshold={threshold} density={density} hits={hits:?}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -1,5 +1,7 @@
|
|||||||
use super::*;
|
use super::*;
|
||||||
use obikindex::IndexConfig;
|
use obikidxcache::index_cache::IndexCache;
|
||||||
|
use obikindex::{IndexConfig, KmerIndex};
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
// ── QueryStats::AddAssign ───────────────────────────────────────────────────
|
// ── QueryStats::AddAssign ───────────────────────────────────────────────────
|
||||||
|
|
||||||
@@ -39,33 +41,39 @@ fn query_stats_default_is_zero() {
|
|||||||
// ── query_partition_with on a not-yet-indexed partition ─────────────────────
|
// ── query_partition_with on a not-yet-indexed partition ─────────────────────
|
||||||
|
|
||||||
/// A `KmerPartition` created but never taken through `build_layers` has no
|
/// A `KmerPartition` created but never taken through `build_layers` has no
|
||||||
/// `index/` subdirectory under any partition — `query_partition_with` must
|
/// `index/` subdirectory under any partition — `IndexCache::new` still opens
|
||||||
/// recognise this and return default (all-zero) stats rather than erroring,
|
/// fine on it (zero layers found), and `query_partition_with` must recognise
|
||||||
/// exactly like an empty `kmers` map.
|
/// this and return default (all-zero) stats rather than panicking, exactly
|
||||||
|
/// like an empty `kmers` map.
|
||||||
#[test]
|
#[test]
|
||||||
fn query_partition_with_missing_index_dir_returns_default_stats() {
|
fn query_partition_with_missing_index_dir_returns_default_stats() {
|
||||||
let tmp = tempfile::tempdir().expect("tempdir");
|
let tmp = tempfile::tempdir().expect("tempdir");
|
||||||
|
// Same k/m pair as `batch`'s tests (11, 5) — `obikseq`'s global k/m
|
||||||
|
// params are process-wide atomics in test builds (see
|
||||||
|
// `obikseq::params`'s own doc comment), shared by every test in this
|
||||||
|
// crate's single test binary regardless of thread; a different pair
|
||||||
|
// here would race with `batch`'s tests running concurrently.
|
||||||
let config = IndexConfig {
|
let config = IndexConfig {
|
||||||
kmer_size: 21,
|
kmer_size: 11,
|
||||||
minimizer_size: 9,
|
minimizer_size: 5,
|
||||||
n_bits: 2,
|
n_bits: 2,
|
||||||
with_counts: false,
|
with_counts: false,
|
||||||
evidence: obikindex::layer::IndexMode::Exact,
|
evidence: obikindex::layer::IndexMode::Exact,
|
||||||
block_bits: 0,
|
block_bits: 0,
|
||||||
};
|
};
|
||||||
let index = KmerIndex::create(tmp.path().join("idx"), config, None).expect("create index");
|
let index = KmerIndex::create(tmp.path().join("idx"), config, None).expect("create index");
|
||||||
|
let cache = IndexCache::new(Arc::new(index), Some(vec![0]));
|
||||||
|
|
||||||
let mut kmers: HashMap<CanonicalKmer, Vec<KmerDesc>> = HashMap::new();
|
let mut kmers: HashMap<CanonicalKmer, Vec<KmerDesc>> = HashMap::new();
|
||||||
// Any well-formed canonical k-mer works here — the call must return
|
// Any well-formed canonical k-mer works here — the call must return
|
||||||
// before ever attempting an MPHF lookup, since `index/` doesn't exist.
|
// before ever attempting an MPHF lookup, since the partition has zero
|
||||||
|
// layers.
|
||||||
let kmer = CanonicalKmer::from_raw_unchecked(0u64);
|
let kmer = CanonicalKmer::from_raw_unchecked(0u64);
|
||||||
kmers.insert(kmer, vec![KmerDesc { seq_idx: 0, pos: 0 }]);
|
kmers.insert(kmer, vec![KmerDesc { seq_idx: 0, pos: 0 }]);
|
||||||
|
|
||||||
let stats = index
|
let stats = cache.query_partition_with(0, &kmers, 1, |_event| {
|
||||||
.query_partition_with(0, &kmers, 1, false, |_event| {
|
|
||||||
panic!("on_event must not be called: no index was built");
|
panic!("on_event must not be called: no index was built");
|
||||||
})
|
});
|
||||||
.expect("query_partition_with should not error on a missing index dir");
|
|
||||||
|
|
||||||
assert_eq!(stats, QueryStats::default());
|
assert_eq!(stats, QueryStats::default());
|
||||||
}
|
}
|
||||||
@@ -73,22 +81,23 @@ fn query_partition_with_missing_index_dir_returns_default_stats() {
|
|||||||
#[test]
|
#[test]
|
||||||
fn query_partition_with_empty_kmers_is_a_noop() {
|
fn query_partition_with_empty_kmers_is_a_noop() {
|
||||||
let tmp = tempfile::tempdir().expect("tempdir");
|
let tmp = tempfile::tempdir().expect("tempdir");
|
||||||
|
// Same k/m pair as `batch`'s tests — see the comment in the previous
|
||||||
|
// test for why this must match across every test in this crate.
|
||||||
let config = IndexConfig {
|
let config = IndexConfig {
|
||||||
kmer_size: 21,
|
kmer_size: 11,
|
||||||
minimizer_size: 9,
|
minimizer_size: 5,
|
||||||
n_bits: 2,
|
n_bits: 2,
|
||||||
with_counts: false,
|
with_counts: false,
|
||||||
evidence: obikindex::layer::IndexMode::Exact,
|
evidence: obikindex::layer::IndexMode::Exact,
|
||||||
block_bits: 0,
|
block_bits: 0,
|
||||||
};
|
};
|
||||||
let index = KmerIndex::create(tmp.path().join("idx"), config, None).expect("create index");
|
let index = KmerIndex::create(tmp.path().join("idx"), config, None).expect("create index");
|
||||||
|
let cache = IndexCache::new(Arc::new(index), Some(vec![0]));
|
||||||
|
|
||||||
let kmers: HashMap<CanonicalKmer, Vec<KmerDesc>> = HashMap::new();
|
let kmers: HashMap<CanonicalKmer, Vec<KmerDesc>> = HashMap::new();
|
||||||
let stats = index
|
let stats = cache.query_partition_with(0, &kmers, 1, |_event| {
|
||||||
.query_partition_with(0, &kmers, 1, false, |_event| {
|
|
||||||
panic!("on_event must not be called on an empty kmer map");
|
panic!("on_event must not be called on an empty kmer map");
|
||||||
})
|
});
|
||||||
.expect("query_partition_with on an empty map should not error");
|
|
||||||
|
|
||||||
assert_eq!(stats, QueryStats::default());
|
assert_eq!(stats, QueryStats::default());
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -8,6 +8,7 @@
|
|||||||
|
|
||||||
use std::collections::HashMap;
|
use std::collections::HashMap;
|
||||||
use std::fs;
|
use std::fs;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
use obidebruinj::GraphDeBruijn;
|
use obidebruinj::GraphDeBruijn;
|
||||||
use obikfilter::FilteredPartitionIter;
|
use obikfilter::FilteredPartitionIter;
|
||||||
@@ -24,7 +25,7 @@ pub trait IndexCompact {
|
|||||||
fn compact_layers(&self, on_partition: impl Fn() + Send + Sync) -> OKIResult<()>;
|
fn compact_layers(&self, on_partition: impl Fn() + Send + Sync) -> OKIResult<()>;
|
||||||
}
|
}
|
||||||
|
|
||||||
impl IndexCompact for KmerIndex {
|
impl IndexCompact for Arc<KmerIndex> {
|
||||||
fn compact_layers(&self, on_partition: impl Fn() + Send + Sync) -> OKIResult<()> {
|
fn compact_layers(&self, on_partition: impl Fn() + Send + Sync) -> OKIResult<()> {
|
||||||
let n_genomes = self.meta().genomes().map_err(OKIError::Io)?.len();
|
let n_genomes = self.meta().genomes().map_err(OKIError::Io)?.len();
|
||||||
let use_counts = self.meta().config.with_counts;
|
let use_counts = self.meta().config.with_counts;
|
||||||
@@ -47,7 +48,7 @@ impl IndexCompact for KmerIndex {
|
|||||||
/// existing layers in place. No-op if the partition already has one layer
|
/// existing layers in place. No-op if the partition already has one layer
|
||||||
/// (or none).
|
/// (or none).
|
||||||
fn compact_partition(
|
fn compact_partition(
|
||||||
idx: &KmerIndex,
|
idx: &Arc<KmerIndex>,
|
||||||
i: usize,
|
i: usize,
|
||||||
use_counts: bool,
|
use_counts: bool,
|
||||||
n_genomes: usize,
|
n_genomes: usize,
|
||||||
@@ -63,7 +64,7 @@ fn compact_partition(
|
|||||||
// ── Read: every layer's kmers, combined — no filter, nothing removed ──
|
// ── Read: every layer's kmers, combined — no filter, nothing removed ──
|
||||||
let mut survivors: HashMap<CanonicalKmer, Box<[u32]>> = HashMap::new();
|
let mut survivors: HashMap<CanonicalKmer, Box<[u32]>> = HashMap::new();
|
||||||
{
|
{
|
||||||
let cache = IndexCache::new(idx, Some(vec![i]));
|
let cache = IndexCache::new(Arc::clone(idx), Some(vec![i]));
|
||||||
cache.iter_partition_kmers(i, use_counts, n_genomes, &[], |kmer, row| {
|
cache.iter_partition_kmers(i, use_counts, n_genomes, &[], |kmer, row| {
|
||||||
survivors.insert(kmer, row);
|
survivors.insert(kmer, row);
|
||||||
true
|
true
|
||||||
|
|||||||
+13
-12
@@ -1,5 +1,6 @@
|
|||||||
use std::fs;
|
use std::fs;
|
||||||
use std::path::Path;
|
use std::path::Path;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
use rayon::prelude::*;
|
use rayon::prelude::*;
|
||||||
|
|
||||||
@@ -88,22 +89,22 @@ fn layer_bytes(layer: &KmerLayer) -> LayerBytes {
|
|||||||
/// once up front — this is diagnostic tooling, not a hot path, so trading
|
/// once up front — this is diagnostic tooling, not a hot path, so trading
|
||||||
/// that eager open for not having to hand-reconstruct every layer's file
|
/// that eager open for not having to hand-reconstruct every layer's file
|
||||||
/// paths is the right tradeoff here.
|
/// paths is the right tradeoff here.
|
||||||
pub struct BitsPerKmer<'a> {
|
pub struct BitsPerKmer {
|
||||||
index: &'a KmerIndex,
|
index: Arc<KmerIndex>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl<'a> BitsPerKmer<'a> {
|
impl BitsPerKmer {
|
||||||
pub fn new(index: &'a KmerIndex) -> Self {
|
pub fn new(index: Arc<KmerIndex>) -> Self {
|
||||||
Self { index }
|
Self { index }
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Algorithm for BitsPerKmer<'_> {
|
impl Algorithm for BitsPerKmer {
|
||||||
type Output = IndexBitsPerKmer;
|
type Output = IndexBitsPerKmer;
|
||||||
|
|
||||||
fn run(&mut self) -> obikalgorithm::Result<IndexBitsPerKmer> {
|
fn run(&mut self) -> obikalgorithm::Result<IndexBitsPerKmer> {
|
||||||
let n_genomes = self.index.meta().genomes()?.len().max(1);
|
let n_genomes = self.index.meta().genomes()?.len().max(1);
|
||||||
let cache = IndexCache::new(self.index, None);
|
let cache = IndexCache::new(Arc::clone(&self.index), None);
|
||||||
let layers: Vec<&KmerLayer> = cache.iter().collect();
|
let layers: Vec<&KmerLayer> = cache.iter().collect();
|
||||||
|
|
||||||
let (n_kmers, mphf_b, evidence_b, matrix_b) = layers
|
let (n_kmers, mphf_b, evidence_b, matrix_b) = layers
|
||||||
@@ -148,23 +149,23 @@ impl Algorithm for BitsPerKmer<'_> {
|
|||||||
/// genome has a non-zero value (presence = 1, count > 0) — via
|
/// genome has a non-zero value (presence = 1, count > 0) — via
|
||||||
/// `KmerLayer::col_weights` (`obicompactvec::ColumnWeights`), one call per
|
/// `KmerLayer::col_weights` (`obicompactvec::ColumnWeights`), one call per
|
||||||
/// layer instead of hand-opening each layer's matrix and summing rows.
|
/// layer instead of hand-opening each layer's matrix and summing rows.
|
||||||
pub struct GenomeKmerCounts<'a> {
|
pub struct GenomeKmerCounts {
|
||||||
index: &'a KmerIndex,
|
index: Arc<KmerIndex>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl<'a> GenomeKmerCounts<'a> {
|
impl GenomeKmerCounts {
|
||||||
pub fn new(index: &'a KmerIndex) -> Self {
|
pub fn new(index: Arc<KmerIndex>) -> Self {
|
||||||
Self { index }
|
Self { index }
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Algorithm for GenomeKmerCounts<'_> {
|
impl Algorithm for GenomeKmerCounts {
|
||||||
/// `(total_distinct_kmers, per_genome_kmer_counts)`.
|
/// `(total_distinct_kmers, per_genome_kmer_counts)`.
|
||||||
type Output = (usize, Vec<u64>);
|
type Output = (usize, Vec<u64>);
|
||||||
|
|
||||||
fn run(&mut self) -> obikalgorithm::Result<(usize, Vec<u64>)> {
|
fn run(&mut self) -> obikalgorithm::Result<(usize, Vec<u64>)> {
|
||||||
let n_genomes = self.index.meta().genomes()?.len();
|
let n_genomes = self.index.meta().genomes()?.len();
|
||||||
let cache = IndexCache::new(self.index, None);
|
let cache = IndexCache::new(Arc::clone(&self.index), None);
|
||||||
let layers: Vec<&KmerLayer> = cache.iter().collect();
|
let layers: Vec<&KmerLayer> = cache.iter().collect();
|
||||||
|
|
||||||
let total_kmers: usize = layers.iter().map(|l| l.n()).sum();
|
let total_kmers: usize = layers.iter().map(|l| l.n()).sum();
|
||||||
|
|||||||
Reference in New Issue
Block a user