OBIKmers/obikmer
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 34 Wiki Activity

Rename index methods to hash and expose kmer iteration APIs

Browse Source 
The layer API has been updated to rename `index` and `index_batch` to `hash` and `hash_batch`, aligning with raw MPHF slot hashing behavior. New public methods have been added to directly expose kmer hashing and iteration, enforcing strict state validation on empty layers. All other modifications are consistent code formatting adjustments without functional changes.
This commit is contained in:
Eric Coissac
2026-08-21 10:14:22 +02:00
parent 0299b608e2
commit 346095b9eb
5 changed files with 294 additions and 145 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/obikphylo/src/siblings/cache.rs
+3 -7
View File
@@ -1,7 +1,7 @@
use rayon::prelude::*;
use obikseq::CanonicalKmer;
use obikindex::layer::Layer;
use obikseq::CanonicalKmer;
use obisys::progress_bar;
use obikindex::{KmerIndex, OKIResult};
@@ -45,11 +45,7 @@ pub(super) struct PartitionCache {
}
impl PartitionCache {
pub(super) fn build(
index: &KmerIndex,
n_parts: usize,
with_counts: bool,
) -> OKIResult<Self> {
pub(super) fn build(index: &KmerIndex, n_parts: usize, with_counts: bool) -> OKIResult<Self> {
let pb = progress_bar("open_partitions", n_parts as u64, "partitions");
let built: Vec<(Vec<Layer>, usize)> = (0..n_parts)
.into_par_iter()
@@ -200,7 +196,7 @@ impl PartitionCache {
let mat = &mats[li];
let variants: Vec<CanonicalKmer> =
entries.iter().map(|&(variant, _, _)| variant).collect();
let slots = mat.index_batch(&variants);
let slots = mat.hash_batch(&variants);
let hits: Vec<(usize, usize, u8)> = slots
.into_iter()
.zip(entries.iter())
src/obikphylo/src/siblings/family_scan.rs
+43 -23
View File
@@ -53,19 +53,19 @@ use std::sync::atomic::{AtomicU8, Ordering};
use rayon::prelude::*;
use obikseq::CanonicalKmer;
use obikindex::layer::meta::PartitionMeta;
use obikseq::CanonicalKmer;
use obipipeline::{ThrottleGuard, throttle};
use obikindex::{OKIError, OKIResult};
use obikindex::KmerIndex;
use obikindex::{OKIError, OKIResult};
use obikindex::layer::Layer;
use super::cache::PartitionCache;
use super::helpers::central_base;
use super::iter::{SiblingEntry, SiblingLayerExt};
use super::{olm_to_ok, FamilyMask, SiblingAnnex, ANNEX_FILE_NAME};
use super::{ANNEX_FILE_NAME, FamilyMask, SiblingAnnex, olm_to_ok};
/// Families per batch — see the module docs for the memory-vs-per-partition-
/// density trade-off this picks a point on. At ~90 genomes and a few
@@ -191,14 +191,22 @@ pub(super) fn scan_layer_families(
selection: &Selection,
mut on_family: impl FnMut(usize, FamilyMask, &[u8]),
) -> OKIResult<()> {
let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir");
let index_dir = layer_dir
.parent()
.expect("layer_dir has a parent index dir");
let meta = PartitionMeta::load(index_dir).map_err(olm_to_ok)?;
let annex = Arc::new(SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?);
let mat = Layer::open(layer_dir, &meta.mode, with_counts).map_err(olm_to_ok)?;
let n_cols = mat.n_cols().min(n_genomes);
let ctx = Arc::new(LayerCtx { mat, n_parts, n_genomes, n_cols, k });
let ctx = Arc::new(LayerCtx {
mat,
n_parts,
n_genomes,
n_cols,
k,
});
// Streamed straight from `iter_minorants_batch` (zips this layer's own
// `iter_kmers()` with the annex, both in iteration order — never an
@@ -209,11 +217,14 @@ pub(super) fn scan_layer_families(
// (see the module docs). `.scan()` computes each batch's starting
// family index lazily, mirroring what the eager `chunks()`+running
// `offset` used to do.
let batches = ctx.mat.iter_minorants_batch(annex, FAMILY_BATCH).scan(0usize, |offset, batch| {
let start = *offset;
*offset += batch.len();
Some((start, batch))
});
let batches =
ctx.mat
.iter_minorants_batch(annex, FAMILY_BATCH)
.scan(0usize, |offset, batch| {
let start = *offset;
*offset += batch.len();
Some((start, batch))
});
let n_workers = obisys::effective_parallelism();
let capacity = 4;
@@ -288,7 +299,7 @@ pub(super) fn scan_layer_families(
// — a pure MPHF lookup, no evidence check, since these are
// this layer's own kmers, known members by construction.
let kmers: Vec<CanonicalKmer> = batch.entries.iter().map(|e| e.kmer).collect();
let slots = ctx.mat.index_batch(&kmers);
let slots = ctx.mat.hash_batch(&kmers);
let mut carries: Vec<Vec<bool>> = (0..ctx.n_cols).map(|_| Vec::new()).collect();
ctx.mat.fill_sub_matrix_carries(&slots, &mut carries);
for (g, col) in carries.iter().enumerate() {
@@ -321,18 +332,24 @@ pub(super) fn scan_layer_families(
// (one batch at a time, never several concurrently — see the
// module docs), each thread owning one partition's queries
// contiguously until this batch is done.
let genome_mask: Vec<AtomicU8> = batch.genome_mask.into_iter().map(AtomicU8::new).collect();
let genome_mask: Vec<AtomicU8> =
batch.genome_mask.into_iter().map(AtomicU8::new).collect();
let fast_mode = cache.fast_mode();
batch.outgoing.par_iter().enumerate().filter(|(_, q)| !q.is_empty()).for_each(|(dest, queries)| {
let on_hit = |i: usize, base: u8, g: usize| {
genome_mask[i * n_genomes + g].fetch_or(1 << base, Ordering::Relaxed);
};
if fast_mode {
cache.find_presence_batch_fast(dest, queries, n_genomes, on_hit);
} else {
cache.find_presence_batch(dest, queries, n_genomes, on_hit);
}
});
batch
.outgoing
.par_iter()
.enumerate()
.filter(|(_, q)| !q.is_empty())
.for_each(|(dest, queries)| {
let on_hit = |i: usize, base: u8, g: usize| {
genome_mask[i * n_genomes + g].fetch_or(1 << base, Ordering::Relaxed);
};
if fast_mode {
cache.find_presence_batch_fast(dest, queries, n_genomes, on_hit);
} else {
cache.find_presence_batch(dest, queries, n_genomes, on_hit);
}
});
for i in 0..n {
let family_idx = batch.start_family_idx + i;
@@ -347,7 +364,10 @@ pub(super) fn scan_layer_families(
next_expected += n;
}
}
debug_assert!(pending.is_empty(), "every generated batch must have been replayed");
debug_assert!(
pending.is_empty(),
"every generated batch must have been replayed"
);
Ok(())
}