Rename index methods to hash and expose kmer iteration APIs

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
+89 -16
View File
@@ -28,9 +28,11 @@ use obicompactvec::{PersistentBitMatrix, PersistentCompactIntMatrix};
use obikseq::CanonicalKmer;
use crate::layer::error::OLMResult;
use crate::layer::typed_layer::{LayerContent, TypedLayer, COUNTS_DIR, PRESENCE_DIR};
use crate::layer::meta::IndexMode;
use crate::layer::mphf_layer::{EvidenceKind, EVIDENCE_FILE, FINGERPRINT_FILE, MPHF_FILE, UNITIGS_FILE};
use crate::layer::mphf_layer::{
EVIDENCE_FILE, EvidenceKind, FINGERPRINT_FILE, MPHF_FILE, UNITIGS_FILE,
};
use crate::layer::typed_layer::{COUNTS_DIR, LayerContent, PRESENCE_DIR, TypedLayer};
/// One layer, at any point in its life — see the module docs. Only
/// [`Empty`](Layer::Empty) and the two ready-to-read states
@@ -42,7 +44,9 @@ pub enum Layer {
/// the path accessors panics on this variant: calling them means the
/// caller assumed a layer was ready when it wasn't, an implementation
/// error to surface loudly, not paper over with a default value.
Empty { dir: PathBuf },
Empty {
dir: PathBuf,
},
Count(TypedLayer<PersistentCompactIntMatrix>),
Presence(TypedLayer<PersistentBitMatrix>),
}
@@ -53,7 +57,9 @@ impl Layer {
/// building a new layer.
pub fn create(dir: &Path) -> std::io::Result<Self> {
std::fs::create_dir_all(dir)?;
Ok(Layer::Empty { dir: dir.to_owned() })
Ok(Layer::Empty {
dir: dir.to_owned(),
})
}
/// Open one layer, auto-detecting count vs. presence from what is
@@ -66,7 +72,8 @@ impl Layer {
/// here.
pub fn open(layer_dir: &Path, mode: &IndexMode, with_counts: bool) -> OLMResult<Self> {
if with_counts && layer_dir.join(COUNTS_DIR).exists() {
return TypedLayer::<PersistentCompactIntMatrix>::open(layer_dir, mode).map(Layer::Count);
return TypedLayer::<PersistentCompactIntMatrix>::open(layer_dir, mode)
.map(Layer::Count);
}
TypedLayer::<PersistentBitMatrix>::open(layer_dir, mode).map(Layer::Presence)
}
@@ -91,12 +98,24 @@ impl Layer {
}
}
pub fn mphf_path(&self) -> PathBuf { self.dir().join(MPHF_FILE) }
pub fn unitigs_path(&self) -> PathBuf { self.dir().join(UNITIGS_FILE) }
pub fn evidence_path(&self) -> PathBuf { self.dir().join(EVIDENCE_FILE) }
pub fn fingerprint_path(&self) -> PathBuf { self.dir().join(FINGERPRINT_FILE) }
pub fn counts_dir(&self) -> PathBuf { self.dir().join(COUNTS_DIR) }
pub fn presence_dir(&self) -> PathBuf { self.dir().join(PRESENCE_DIR) }
pub fn mphf_path(&self) -> PathBuf {
self.dir().join(MPHF_FILE)
}
pub fn unitigs_path(&self) -> PathBuf {
self.dir().join(UNITIGS_FILE)
}
pub fn evidence_path(&self) -> PathBuf {
self.dir().join(EVIDENCE_FILE)
}
pub fn fingerprint_path(&self) -> PathBuf {
self.dir().join(FINGERPRINT_FILE)
}
pub fn counts_dir(&self) -> PathBuf {
self.dir().join(COUNTS_DIR)
}
pub fn presence_dir(&self) -> PathBuf {
self.dir().join(PRESENCE_DIR)
}
// ── Ready-only surface ───────────────────────────────────────────────
@@ -135,11 +154,11 @@ impl Layer {
/// Raw MPHF batch lookup: kmer → slot, no membership check — for
/// callers that already know every kmer is a member of *this* layer, so
/// the evidence check `find_slot`/`find` would perform is redundant.
pub fn index_batch(&self, kmers: &[CanonicalKmer]) -> Vec<usize> {
pub fn hash_batch(&self, kmers: &[CanonicalKmer]) -> Vec<usize> {
match self {
Layer::Count(l) => l.index_batch(kmers),
Layer::Presence(l) => l.index_batch(kmers),
Layer::Empty { .. } => panic!("Layer::index_batch() called on an Empty layer"),
Layer::Count(l) => l.hash_batch(kmers),
Layer::Presence(l) => l.hash_batch(kmers),
Layer::Empty { .. } => panic!("Layer::hash_batch() called on an Empty layer"),
}
}
@@ -168,7 +187,61 @@ impl Layer {
o.extend(c.iter().map(|&v| v != 0));
}
}
Layer::Empty { .. } => panic!("Layer::fill_sub_matrix_carries() called on an Empty layer"),
Layer::Empty { .. } => {
panic!("Layer::fill_sub_matrix_carries() called on an Empty layer")
}
}
}
/// Raw MPHF lookup: kmer → slot, no membership check.
pub fn hash(&self, kmer: CanonicalKmer) -> usize {
match self {
Layer::Count(l) => l.hash(kmer),
Layer::Presence(l) => l.hash(kmer),
Layer::Empty { .. } => panic!("Layer::hash() called on an Empty layer"),
}
}
/// Iterate over all canonical kmers in the layer, in deterministic order.
pub fn iter_kmers(&self) -> crate::layer::mphf_layer::KmerIter {
match self {
Layer::Count(l) => l.iter_kmers(),
Layer::Presence(l) => l.iter_kmers(),
Layer::Empty { .. } => panic!("Layer::iter_kmers() called on an Empty layer"),
}
}
/// Iterate over all canonical kmers, each paired with its zero-based
/// sequence index in `unitigs.bin`.
pub fn enumerate_kmers(&self) -> std::iter::Enumerate<crate::layer::mphf_layer::KmerIter> {
match self {
Layer::Count(l) => l.enumerate_kmers(),
Layer::Presence(l) => l.enumerate_kmers(),
Layer::Empty { .. } => panic!("Layer::enumerate_kmers() called on an Empty layer"),
}
}
/// Iterate over the layer's canonical kmers in batches of `n`.
pub fn iter_kmers_batch(&self, n: usize) -> crate::layer::mphf_layer::KmerBatchIter {
match self {
Layer::Count(l) => l.iter_kmers_batch(n),
Layer::Presence(l) => l.iter_kmers_batch(n),
Layer::Empty { .. } => panic!("Layer::iter_kmers_batch() called on an Empty layer"),
}
}
/// Iterate over batches, each paired with the zero-based index of the
/// first kmer in the batch.
pub fn enumerate_kmers_batch(
&self,
n: usize,
) -> Box<dyn Iterator<Item = (usize, Vec<CanonicalKmer>)> + Send + 'static> {
match self {
Layer::Count(l) => Box::new(l.enumerate_kmers_batch(n)),
Layer::Presence(l) => Box::new(l.enumerate_kmers_batch(n)),
Layer::Empty { .. } => {
panic!("Layer::enumerate_kmers_batch() called on an Empty layer")
}
}
}
}
+104 -64
View File
@@ -30,9 +30,20 @@ type MphfEps = PtrHash<u64, CubicEps, CachelineEfVec<&'static [CachelineEf]>, Xx
// ── LayerEvidence ─────────────────────────────────────────────────────────────
enum LayerEvidence {
Exact { evidence: Evidence, unitigs: Arc<UnitigFileReader> },
Approx { fingerprint: FingerprintVec, unitigs: Arc<UnitigFileReader>, unitigs_path: PathBuf },
Hybrid { evidence: Evidence, unitigs: Arc<UnitigFileReader>, fingerprint: FingerprintVec },
Exact {
evidence: Evidence,
unitigs: Arc<UnitigFileReader>,
},
Approx {
fingerprint: FingerprintVec,
unitigs: Arc<UnitigFileReader>,
unitigs_path: PathBuf,
},
Hybrid {
evidence: Evidence,
unitigs: Arc<UnitigFileReader>,
fingerprint: FingerprintVec,
},
}
// ── EvidenceKind ──────────────────────────────────────────────────────────────
@@ -58,7 +69,8 @@ impl EvidenceKind {
(false, true) => Ok(EvidenceKind::Approx),
(true, true) => Ok(EvidenceKind::Hybrid),
(false, false) => Err(OLMError::InvalidLayer(format!(
"no evidence.bin or fingerprint.bin in {}", layer_dir.display()
"no evidence.bin or fingerprint.bin in {}",
layer_dir.display()
))),
}
}
@@ -96,14 +108,28 @@ impl MphfLayer {
let n = fingerprint.n();
let unitigs = Arc::new(UnitigFileReader::open(&dir.join(UNITIGS_FILE))?);
let unitigs_path = dir.join(UNITIGS_FILE);
(LayerEvidence::Approx { fingerprint, unitigs, unitigs_path }, n)
(
LayerEvidence::Approx {
fingerprint,
unitigs,
unitigs_path,
},
n,
)
}
IndexMode::Hybrid { .. } => {
let evidence = Evidence::open(&dir.join(EVIDENCE_FILE))?;
let fingerprint = FingerprintVec::open(&dir.join(FINGERPRINT_FILE))?;
let n = evidence.len();
let unitigs = Arc::new(UnitigFileReader::open(&dir.join(UNITIGS_FILE))?);
(LayerEvidence::Hybrid { evidence, unitigs, fingerprint }, n)
(
LayerEvidence::Hybrid {
evidence,
unitigs,
fingerprint,
},
n,
)
}
};
Ok(Self { mphf, ev, n })
@@ -118,9 +144,13 @@ impl MphfLayer {
#[inline]
pub fn find(&self, kmer: CanonicalKmer) -> Option<usize> {
let slot = self.mphf.index(&kmer.raw());
if slot >= self.n { return None; }
if slot >= self.n {
return None;
}
match &self.ev {
LayerEvidence::Exact { evidence, unitigs, .. } => {
LayerEvidence::Exact {
evidence, unitigs, ..
} => {
let (chunk_id, rank) = evidence.decode(slot);
if unitigs.verify_canonical_kmer(chunk_id as usize, rank as usize, kmer) {
Some(slot)
@@ -128,9 +158,13 @@ impl MphfLayer {
None
}
}
LayerEvidence::Approx { fingerprint, .. } |
LayerEvidence::Hybrid { fingerprint, .. } => {
if fingerprint.matches(slot, kmer.seq_hash()) { Some(slot) } else { None }
LayerEvidence::Approx { fingerprint, .. }
| LayerEvidence::Hybrid { fingerprint, .. } => {
if fingerprint.matches(slot, kmer.seq_hash()) {
Some(slot)
} else {
None
}
}
}
}
@@ -142,10 +176,16 @@ impl MphfLayer {
/// that owns the slot, then exact comparison.
pub fn find_strict(&self, kmer: CanonicalKmer) -> Option<usize> {
let slot = self.mphf.index(&kmer.raw());
if slot >= self.n { return None; }
if slot >= self.n {
return None;
}
match &self.ev {
LayerEvidence::Exact { evidence, unitigs, .. } |
LayerEvidence::Hybrid { evidence, unitigs, .. } => {
LayerEvidence::Exact {
evidence, unitigs, ..
}
| LayerEvidence::Hybrid {
evidence, unitigs, ..
} => {
let (chunk_id, rank) = evidence.decode(slot);
if unitigs.verify_canonical_kmer(chunk_id as usize, rank as usize, kmer) {
Some(slot)
@@ -165,32 +205,10 @@ impl MphfLayer {
}
}
/// Reconstruct the canonical k-mer stored at `slot` — the inverse of
/// [`find`](Self::find)/[`find_strict`](Self::find_strict) (k-mer → slot).
/// O(1) on `Exact`/`Hybrid` layers: `evidence.decode(slot)` gives
/// `(chunk_id, rank)` directly (no MPHF hashing, no file scan), then
/// `unitigs.canonical_raw_kmer` is a direct-access read. `None` on
/// `Approx` layers — fingerprints alone can't recover a k-mer, and
/// falling back to a full sequential scan here would silently make one
/// call cost O(n); callers needing this on an `Approx` layer should
/// scan `unitigs.bin` themselves and decide how to handle that cost
/// explicitly.
pub fn kmer_at(&self, slot: usize) -> Option<CanonicalKmer> {
if slot >= self.n {
return None;
/// Number of slots in the MPHF layer — equal to the number of stored k-mers.
pub fn n(&self) -> usize {
self.n
}
match &self.ev {
LayerEvidence::Exact { evidence, unitigs, .. } |
LayerEvidence::Hybrid { evidence, unitigs, .. } => {
let (chunk_id, rank) = evidence.decode(slot);
let raw = unitigs.canonical_raw_kmer(chunk_id as usize, rank as usize);
Some(CanonicalKmer::from_raw_unchecked(raw))
}
LayerEvidence::Approx { .. } => None,
}
}
pub fn n(&self) -> usize { self.n }
/// This already-open layer's evidence mode — reads the discriminant
/// already in memory, no disk access.
@@ -208,7 +226,7 @@ impl MphfLayer {
/// check against the layer's evidence. The returned slot is only meaningful
/// when the kmer is actually present in the layer; callers that need an
/// existence test should use [`find`](Self::find) instead.
pub fn index(&self, kmer: CanonicalKmer) -> usize {
pub fn hash(&self, kmer: CanonicalKmer) -> usize {
self.mphf.index(&kmer.raw())
}
@@ -217,7 +235,7 @@ impl MphfLayer {
/// Returns a [`Vec`] of slots, one per input kmer, in the same order as the
/// input slice. Like [`index`](Self::index), no membership check is
/// performed.
pub fn index_batch(&self, kmers: &[CanonicalKmer]) -> Vec<usize> {
pub fn hash_batch(&self, kmers: &[CanonicalKmer]) -> Vec<usize> {
kmers.iter().map(|k| self.mphf.index(&k.raw())).collect()
}
@@ -237,7 +255,9 @@ impl MphfLayer {
LayerEvidence::Approx { unitigs, .. } => unitigs,
LayerEvidence::Hybrid { unitigs, .. } => unitigs,
};
KmerIter { inner: Box::new(reader.iter_indexed_canonical_kmers_owned()) }
KmerIter {
inner: Box::new(reader.iter_indexed_canonical_kmers_owned()),
}
}
/// Iterate over all canonical kmers, each paired with its zero-based
@@ -255,7 +275,10 @@ impl MphfLayer {
/// Each call to [`next`](Iterator::next) returns a [`Vec`] of up to `n`
/// kmers. The final batch may be shorter when the layer is exhausted.
pub fn iter_kmers_batch(&self, n: usize) -> KmerBatchIter {
KmerBatchIter { inner: self.iter_kmers(), batch_size: n }
KmerBatchIter {
inner: self.iter_kmers(),
batch_size: n,
}
}
/// Iterate over batches, each paired with the zero-based index of the
@@ -265,7 +288,10 @@ impl MphfLayer {
/// `batch_start_index` is the iteration-order offset of the first kmer
/// in that batch within the full layer sequence — i.e. a multiple of `n`
/// except for the final (possibly shorter) batch.
pub fn enumerate_kmers_batch(&self, n: usize) -> impl Iterator<Item = (usize, Vec<CanonicalKmer>)> + Send + 'static {
pub fn enumerate_kmers_batch(
&self,
n: usize,
) -> impl Iterator<Item = (usize, Vec<CanonicalKmer>)> + Send + 'static {
let mut offset = 0usize;
self.iter_kmers_batch(n).map(move |batch| {
let base = offset;
@@ -320,11 +346,7 @@ impl Iterator for KmerBatchIter {
break;
}
}
if batch.is_empty() {
None
} else {
Some(batch)
}
if batch.is_empty() { None } else { Some(batch) }
}
}
@@ -398,7 +420,9 @@ impl MphfLayer {
/// Build `fingerprint.bin` from `unitigs.bin` + `mphf.bin`.
pub fn build_approx_evidence(dir: &Path, b: u8, z: u8) -> OLMResult<usize> {
if b == 0 || b > 64 {
return Err(OLMError::InvalidLayer("fingerprint width must be 1..=64".into()));
return Err(OLMError::InvalidLayer(
"fingerprint width must be 1..=64".into(),
));
}
if z == 0 {
return Err(OLMError::InvalidLayer("z must be ≥ 1".into()));
@@ -452,8 +476,7 @@ impl MphfLayer {
// ── Empty layer ───────────────────────────────────────────────────────
if n == 0 {
let mphf: Mphf =
Mphf::try_new(&[] as &[u64], PtrHashParams::<CubicEps>::default())
let mphf: Mphf = Mphf::try_new(&[] as &[u64], PtrHashParams::<CubicEps>::default())
.ok_or_else(|| OLMError::Mphf("construction failed".into()))?;
mphf.store(&dir.join(MPHF_FILE))
.map_err(|e| OLMError::InvalidLayer(e.to_string()))?;
@@ -474,9 +497,11 @@ impl MphfLayer {
// ── Pass 1: MPHF via clonable mmap iterator ───────────────────────────
let keys = CanonicalKmerIter::new(&unitig_path).map_err(sk_to_olm)?;
let mphf: Mphf =
Mphf::new_from_par_iter(n, keys.map(|k| k.raw()).par_bridge(),
PtrHashParams::<CubicEps>::default());
let mphf: Mphf = Mphf::new_from_par_iter(
n,
keys.map(|k| k.raw()).par_bridge(),
PtrHashParams::<CubicEps>::default(),
);
mphf.store(&dir.join(MPHF_FILE))
.map_err(|e| OLMError::InvalidLayer(e.to_string()))?;
@@ -489,9 +514,14 @@ impl MphfLayer {
let mut ev = EvidenceWriter::new(n);
for (kmer, chunk_id, rank) in unitigs2.iter_indexed_canonical_kmers() {
let slot = mphf.index(&kmer.raw());
if slot >= n { return Err(OLMError::Mphf("slot out of bounds".into())); }
let byte = slot / 8; let bit = 1u8 << (slot % 8);
if seen[byte] & bit != 0 { return Err(OLMError::Mphf("duplicate slot".into())); }
if slot >= n {
return Err(OLMError::Mphf("slot out of bounds".into()));
}
let byte = slot / 8;
let bit = 1u8 << (slot % 8);
if seen[byte] & bit != 0 {
return Err(OLMError::Mphf("duplicate slot".into()));
}
seen[byte] |= bit;
ev.set(slot, chunk_id as u32, rank as u8);
fill_slot(slot, kmer)?;
@@ -504,9 +534,14 @@ impl MphfLayer {
let mut fw = FingerprintVecWriter::new(n, *b);
for (kmer, _, _) in unitigs2.iter_indexed_canonical_kmers() {
let slot = mphf.index(&kmer.raw());
if slot >= n { return Err(OLMError::Mphf("slot out of bounds".into())); }
let byte = slot / 8; let bit = 1u8 << (slot % 8);
if seen[byte] & bit != 0 { return Err(OLMError::Mphf("duplicate slot".into())); }
if slot >= n {
return Err(OLMError::Mphf("slot out of bounds".into()));
}
let byte = slot / 8;
let bit = 1u8 << (slot % 8);
if seen[byte] & bit != 0 {
return Err(OLMError::Mphf("duplicate slot".into()));
}
seen[byte] |= bit;
fw.set(slot, kmer.seq_hash());
fill_slot(slot, kmer)?;
@@ -519,9 +554,14 @@ impl MphfLayer {
let mut fw = FingerprintVecWriter::new(n, *b);
for (kmer, chunk_id, rank) in unitigs2.iter_indexed_canonical_kmers() {
let slot = mphf.index(&kmer.raw());
if slot >= n { return Err(OLMError::Mphf("slot out of bounds".into())); }
let byte = slot / 8; let bit = 1u8 << (slot % 8);
if seen[byte] & bit != 0 { return Err(OLMError::Mphf("duplicate slot".into())); }
if slot >= n {
return Err(OLMError::Mphf("slot out of bounds".into()));
}
let byte = slot / 8;
let bit = 1u8 << (slot % 8);
if seen[byte] & bit != 0 {
return Err(OLMError::Mphf("duplicate slot".into()));
}
seen[byte] |= bit;
ev.set(slot, chunk_id as u32, rank as u8);
fw.set(slot, kmer.seq_hash());
+40 -20
View File
@@ -3,10 +3,8 @@ use std::fs;
use std::path::{Path, PathBuf};
use obicompactvec::{
BinaryMatrix,
PersistentBitMatrix, PersistentBitMatrixBuilder,
PersistentCompactIntMatrix, PersistentCompactIntMatrixBuilder,
PersistentSparseBitMatrix,
BinaryMatrix, PersistentBitMatrix, PersistentBitMatrixBuilder, PersistentCompactIntMatrix,
PersistentCompactIntMatrixBuilder, PersistentSparseBitMatrix,
};
use obikseq::CanonicalKmer;
use obiskio::{UnitigFileReader, UnitigFileWriter};
@@ -82,7 +80,9 @@ pub fn open_data<D: LayerData>(root: &Path, i: usize) -> OLMResult<D> {
impl LayerData for () {
type Item = ();
fn open(_layer_dir: &Path) -> OLMResult<Self> { Ok(()) }
fn open(_layer_dir: &Path) -> OLMResult<Self> {
Ok(())
}
fn read(&self, _slot: usize) {}
}
@@ -91,7 +91,9 @@ impl LayerData for PersistentCompactIntMatrix {
fn open(layer_dir: &Path) -> OLMResult<Self> {
PersistentCompactIntMatrix::open(layer_dir).map_err(OLMError::Io)
}
fn read(&self, slot: usize) -> Box<[u32]> { self.row(slot) }
fn read(&self, slot: usize) -> Box<[u32]> {
self.row(slot)
}
}
impl LayerData for PersistentBitMatrix {
@@ -99,7 +101,9 @@ impl LayerData for PersistentBitMatrix {
fn open(layer_dir: &Path) -> OLMResult<Self> {
PersistentBitMatrix::open(layer_dir).map_err(OLMError::Io)
}
fn read(&self, slot: usize) -> Box<[bool]> { self.row(slot) }
fn read(&self, slot: usize) -> Box<[bool]> {
self.row(slot)
}
}
impl LayerData for PersistentSparseBitMatrix {
@@ -107,7 +111,9 @@ impl LayerData for PersistentSparseBitMatrix {
fn open(layer_dir: &Path) -> OLMResult<Self> {
PersistentSparseBitMatrix::open(&layer_dir.join(PRESENCE_DIR)).map_err(OLMError::Io)
}
fn read(&self, slot: usize) -> Box<[bool]> { self.row(slot) }
fn read(&self, slot: usize) -> Box<[bool]> {
self.row(slot)
}
}
// ── LayerContent ─────────────────────────────────────────────────────────────
@@ -202,7 +208,10 @@ impl<D: LayerData> TypedLayer<D> {
}
pub fn query(&self, kmer: CanonicalKmer) -> Option<Hit<D::Item>> {
self.mphf.find(kmer).map(|slot| Hit { slot, data: self.data.read(slot) })
self.mphf.find(kmer).map(|slot| Hit {
slot,
data: self.data.read(slot),
})
}
/// MPHF + evidence membership check only — no data read. For callers
@@ -214,16 +223,18 @@ impl<D: LayerData> TypedLayer<D> {
self.mphf.find(kmer)
}
pub fn n(&self) -> usize { self.mphf.n() }
pub fn n(&self) -> usize {
self.mphf.n()
}
/// Raw MPHF lookup: kmer → slot, no membership check.
pub fn index(&self, kmer: CanonicalKmer) -> usize {
self.mphf.index(kmer)
pub fn hash(&self, kmer: CanonicalKmer) -> usize {
self.mphf.hash(kmer)
}
/// Batch raw MPHF lookup: kmers → slots, no membership check.
pub fn index_batch(&self, kmers: &[CanonicalKmer]) -> Vec<usize> {
self.mphf.index_batch(kmers)
pub fn hash_batch(&self, kmers: &[CanonicalKmer]) -> Vec<usize> {
self.mphf.hash_batch(kmers)
}
/// Iterate over all canonical kmers in the layer, in deterministic order.
@@ -244,7 +255,10 @@ impl<D: LayerData> TypedLayer<D> {
/// Iterate over batches, each paired with the zero-based index of the
/// first kmer in the batch.
pub fn enumerate_kmers_batch(&self, n: usize) -> impl Iterator<Item = (usize, Vec<CanonicalKmer>)> + Send + 'static {
pub fn enumerate_kmers_batch(
&self,
n: usize,
) -> impl Iterator<Item = (usize, Vec<CanonicalKmer>)> + Send + 'static {
self.mphf.enumerate_kmers_batch(n)
}
@@ -301,8 +315,12 @@ impl TypedLayer<()> {
pub fn init_presence_matrix(layer_dir: &Path, n_kmers: usize) -> OLMResult<()> {
let presence_dir = layer_dir.join(PRESENCE_DIR);
fs::create_dir_all(&presence_dir).map_err(OLMError::Io)?;
let mut mb = PersistentBitMatrixBuilder::new(n_kmers, &presence_dir).map_err(OLMError::Io)?;
mb.add_col_ones().map_err(OLMError::Io)?.close().map_err(OLMError::Io)?;
let mut mb =
PersistentBitMatrixBuilder::new(n_kmers, &presence_dir).map_err(OLMError::Io)?;
mb.add_col_ones()
.map_err(OLMError::Io)?
.close()
.map_err(OLMError::Io)?;
mb.close().map_err(OLMError::Io)
}
}
@@ -318,8 +336,8 @@ impl TypedLayer<PersistentCompactIntMatrix> {
) -> OLMResult<usize> {
let n = UnitigFileReader::open_sequential(&out_dir.join(UNITIGS_FILE))?.n_kmers();
let counts_dir = out_dir.join(COUNTS_DIR);
let mut mb = PersistentCompactIntMatrixBuilder::new(n, &counts_dir)
.map_err(OLMError::Io)?;
let mut mb =
PersistentCompactIntMatrixBuilder::new(n, &counts_dir).map_err(OLMError::Io)?;
let mut col = mb.add_col().map_err(OLMError::Io)?;
let n_built = MphfLayer::build(out_dir, block_bits, mode, &mut |slot, kmer| {
col.set(slot, count_of(kmer));
@@ -336,7 +354,9 @@ impl TypedLayer<PersistentCompactIntMatrix> {
mode: &IndexMode,
counts: &HashMap<CanonicalKmer, u32>,
) -> OLMResult<usize> {
Self::build(out_dir, block_bits, mode, |kmer| counts.get(&kmer).copied().unwrap_or(0))
Self::build(out_dir, block_bits, mode, |kmer| {
counts.get(&kmer).copied().unwrap_or(0)
})
}
}
+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())
+30 -10
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,7 +217,10 @@ 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 batches =
ctx.mat
.iter_minorants_batch(annex, FAMILY_BATCH)
.scan(0usize, |offset, batch| {
let start = *offset;
*offset += batch.len();
Some((start, batch))
@@ -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,9 +332,15 @@ 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)| {
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);
};
@@ -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(())
}