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 obikseq::CanonicalKmer;
use crate::layer::error::OLMResult; use crate::layer::error::OLMResult;
use crate::layer::typed_layer::{LayerContent, TypedLayer, COUNTS_DIR, PRESENCE_DIR};
use crate::layer::meta::IndexMode; 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 /// One layer, at any point in its life — see the module docs. Only
/// [`Empty`](Layer::Empty) and the two ready-to-read states /// [`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 /// the path accessors panics on this variant: calling them means the
/// caller assumed a layer was ready when it wasn't, an implementation /// caller assumed a layer was ready when it wasn't, an implementation
/// error to surface loudly, not paper over with a default value. /// error to surface loudly, not paper over with a default value.
Empty { dir: PathBuf }, Empty {
dir: PathBuf,
},
Count(TypedLayer<PersistentCompactIntMatrix>), Count(TypedLayer<PersistentCompactIntMatrix>),
Presence(TypedLayer<PersistentBitMatrix>), Presence(TypedLayer<PersistentBitMatrix>),
} }
@@ -53,7 +57,9 @@ impl Layer {
/// building a new layer. /// building a new layer.
pub fn create(dir: &Path) -> std::io::Result<Self> { pub fn create(dir: &Path) -> std::io::Result<Self> {
std::fs::create_dir_all(dir)?; 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 /// Open one layer, auto-detecting count vs. presence from what is
@@ -66,7 +72,8 @@ impl Layer {
/// here. /// here.
pub fn open(layer_dir: &Path, mode: &IndexMode, with_counts: bool) -> OLMResult<Self> { pub fn open(layer_dir: &Path, mode: &IndexMode, with_counts: bool) -> OLMResult<Self> {
if with_counts && layer_dir.join(COUNTS_DIR).exists() { 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) 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 mphf_path(&self) -> PathBuf {
pub fn unitigs_path(&self) -> PathBuf { self.dir().join(UNITIGS_FILE) } self.dir().join(MPHF_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 unitigs_path(&self) -> PathBuf {
pub fn counts_dir(&self) -> PathBuf { self.dir().join(COUNTS_DIR) } self.dir().join(UNITIGS_FILE)
pub fn presence_dir(&self) -> PathBuf { self.dir().join(PRESENCE_DIR) } }
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 ─────────────────────────────────────────────── // ── Ready-only surface ───────────────────────────────────────────────
@@ -135,11 +154,11 @@ impl Layer {
/// Raw MPHF batch lookup: kmer → slot, no membership check — for /// Raw MPHF batch lookup: kmer → slot, no membership check — for
/// callers that already know every kmer is a member of *this* layer, so /// callers that already know every kmer is a member of *this* layer, so
/// the evidence check `find_slot`/`find` would perform is redundant. /// 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 { match self {
Layer::Count(l) => l.index_batch(kmers), Layer::Count(l) => l.hash_batch(kmers),
Layer::Presence(l) => l.index_batch(kmers), Layer::Presence(l) => l.hash_batch(kmers),
Layer::Empty { .. } => panic!("Layer::index_batch() called on an Empty layer"), 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)); 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")
}
} }
} }
} }
+118 -78
View File
@@ -15,9 +15,9 @@ use crate::layer::evidence::{Evidence, EvidenceWriter};
use crate::layer::fingerprint::{FingerprintVec, FingerprintVecWriter}; use crate::layer::fingerprint::{FingerprintVec, FingerprintVecWriter};
use crate::layer::meta::IndexMode; use crate::layer::meta::IndexMode;
pub(crate) const MPHF_FILE: &str = "mphf.bin"; pub(crate) const MPHF_FILE: &str = "mphf.bin";
pub(crate) const UNITIGS_FILE: &str = "unitigs.bin"; pub(crate) const UNITIGS_FILE: &str = "unitigs.bin";
pub(crate) const EVIDENCE_FILE: &str = "evidence.bin"; pub(crate) const EVIDENCE_FILE: &str = "evidence.bin";
pub(crate) const FINGERPRINT_FILE: &str = "fingerprint.bin"; pub(crate) const FINGERPRINT_FILE: &str = "fingerprint.bin";
/// Owned MPHF — used only at build time (construction + store). /// Owned MPHF — used only at build time (construction + store).
@@ -30,9 +30,20 @@ type MphfEps = PtrHash<u64, CubicEps, CachelineEfVec<&'static [CachelineEf]>, Xx
// ── LayerEvidence ───────────────────────────────────────────────────────────── // ── LayerEvidence ─────────────────────────────────────────────────────────────
enum LayerEvidence { enum LayerEvidence {
Exact { evidence: Evidence, unitigs: Arc<UnitigFileReader> }, Exact {
Approx { fingerprint: FingerprintVec, unitigs: Arc<UnitigFileReader>, unitigs_path: PathBuf }, evidence: Evidence,
Hybrid { evidence: Evidence, unitigs: Arc<UnitigFileReader>, fingerprint: FingerprintVec }, unitigs: Arc<UnitigFileReader>,
},
Approx {
fingerprint: FingerprintVec,
unitigs: Arc<UnitigFileReader>,
unitigs_path: PathBuf,
},
Hybrid {
evidence: Evidence,
unitigs: Arc<UnitigFileReader>,
fingerprint: FingerprintVec,
},
} }
// ── EvidenceKind ────────────────────────────────────────────────────────────── // ── EvidenceKind ──────────────────────────────────────────────────────────────
@@ -51,14 +62,15 @@ impl EvidenceKind {
/// `evidence.bin`/`fingerprint.bin` alone determines the mode — same /// `evidence.bin`/`fingerprint.bin` alone determines the mode — same
/// signal `MphfLayer::open` uses, just without opening either file. /// signal `MphfLayer::open` uses, just without opening either file.
pub fn detect(layer_dir: &Path) -> OLMResult<EvidenceKind> { pub fn detect(layer_dir: &Path) -> OLMResult<EvidenceKind> {
let has_evidence = layer_dir.join(EVIDENCE_FILE).exists(); let has_evidence = layer_dir.join(EVIDENCE_FILE).exists();
let has_fingerprint = layer_dir.join(FINGERPRINT_FILE).exists(); let has_fingerprint = layer_dir.join(FINGERPRINT_FILE).exists();
match (has_evidence, has_fingerprint) { match (has_evidence, has_fingerprint) {
(true, false) => Ok(EvidenceKind::Exact), (true, false) => Ok(EvidenceKind::Exact),
(false, true) => Ok(EvidenceKind::Approx), (false, true) => Ok(EvidenceKind::Approx),
(true, true) => Ok(EvidenceKind::Hybrid), (true, true) => Ok(EvidenceKind::Hybrid),
(false, false) => Err(OLMError::InvalidLayer(format!( (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()
))), ))),
} }
} }
@@ -74,8 +86,8 @@ impl EvidenceKind {
/// O(n) sequential scan on Approx layers. /// O(n) sequential scan on Approx layers.
pub struct MphfLayer { pub struct MphfLayer {
mphf: MemCase<MphfEps>, mphf: MemCase<MphfEps>,
ev: LayerEvidence, ev: LayerEvidence,
n: usize, n: usize,
} }
impl MphfLayer { impl MphfLayer {
@@ -96,14 +108,28 @@ impl MphfLayer {
let n = fingerprint.n(); let n = fingerprint.n();
let unitigs = Arc::new(UnitigFileReader::open(&dir.join(UNITIGS_FILE))?); let unitigs = Arc::new(UnitigFileReader::open(&dir.join(UNITIGS_FILE))?);
let unitigs_path = 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 { .. } => { IndexMode::Hybrid { .. } => {
let evidence = Evidence::open(&dir.join(EVIDENCE_FILE))?; let evidence = Evidence::open(&dir.join(EVIDENCE_FILE))?;
let fingerprint = FingerprintVec::open(&dir.join(FINGERPRINT_FILE))?; let fingerprint = FingerprintVec::open(&dir.join(FINGERPRINT_FILE))?;
let n = evidence.len(); let n = evidence.len();
let unitigs = Arc::new(UnitigFileReader::open(&dir.join(UNITIGS_FILE))?); 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 }) Ok(Self { mphf, ev, n })
@@ -118,9 +144,13 @@ impl MphfLayer {
#[inline] #[inline]
pub fn find(&self, kmer: CanonicalKmer) -> Option<usize> { pub fn find(&self, kmer: CanonicalKmer) -> Option<usize> {
let slot = self.mphf.index(&kmer.raw()); let slot = self.mphf.index(&kmer.raw());
if slot >= self.n { return None; } if slot >= self.n {
return None;
}
match &self.ev { match &self.ev {
LayerEvidence::Exact { evidence, unitigs, .. } => { LayerEvidence::Exact {
evidence, unitigs, ..
} => {
let (chunk_id, rank) = evidence.decode(slot); let (chunk_id, rank) = evidence.decode(slot);
if unitigs.verify_canonical_kmer(chunk_id as usize, rank as usize, kmer) { if unitigs.verify_canonical_kmer(chunk_id as usize, rank as usize, kmer) {
Some(slot) Some(slot)
@@ -128,9 +158,13 @@ impl MphfLayer {
None None
} }
} }
LayerEvidence::Approx { fingerprint, .. } | LayerEvidence::Approx { fingerprint, .. }
LayerEvidence::Hybrid { fingerprint, .. } => { | LayerEvidence::Hybrid { fingerprint, .. } => {
if fingerprint.matches(slot, kmer.seq_hash()) { Some(slot) } else { None } if fingerprint.matches(slot, kmer.seq_hash()) {
Some(slot)
} else {
None
}
} }
} }
} }
@@ -142,10 +176,16 @@ impl MphfLayer {
/// that owns the slot, then exact comparison. /// that owns the slot, then exact comparison.
pub fn find_strict(&self, kmer: CanonicalKmer) -> Option<usize> { pub fn find_strict(&self, kmer: CanonicalKmer) -> Option<usize> {
let slot = self.mphf.index(&kmer.raw()); let slot = self.mphf.index(&kmer.raw());
if slot >= self.n { return None; } if slot >= self.n {
return None;
}
match &self.ev { match &self.ev {
LayerEvidence::Exact { evidence, unitigs, .. } | LayerEvidence::Exact {
LayerEvidence::Hybrid { evidence, unitigs, .. } => { evidence, unitigs, ..
}
| LayerEvidence::Hybrid {
evidence, unitigs, ..
} => {
let (chunk_id, rank) = evidence.decode(slot); let (chunk_id, rank) = evidence.decode(slot);
if unitigs.verify_canonical_kmer(chunk_id as usize, rank as usize, kmer) { if unitigs.verify_canonical_kmer(chunk_id as usize, rank as usize, kmer) {
Some(slot) Some(slot)
@@ -165,38 +205,16 @@ impl MphfLayer {
} }
} }
/// Reconstruct the canonical k-mer stored at `slot` — the inverse of /// Number of slots in the MPHF layer — equal to the number of stored k-mers.
/// [`find`](Self::find)/[`find_strict`](Self::find_strict) (k-mer → slot). pub fn n(&self) -> usize {
/// O(1) on `Exact`/`Hybrid` layers: `evidence.decode(slot)` gives self.n
/// `(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;
}
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 /// This already-open layer's evidence mode — reads the discriminant
/// already in memory, no disk access. /// already in memory, no disk access.
pub fn evidence_kind(&self) -> EvidenceKind { pub fn evidence_kind(&self) -> EvidenceKind {
match &self.ev { match &self.ev {
LayerEvidence::Exact { .. } => EvidenceKind::Exact, LayerEvidence::Exact { .. } => EvidenceKind::Exact,
LayerEvidence::Approx { .. } => EvidenceKind::Approx, LayerEvidence::Approx { .. } => EvidenceKind::Approx,
LayerEvidence::Hybrid { .. } => EvidenceKind::Hybrid, LayerEvidence::Hybrid { .. } => EvidenceKind::Hybrid,
} }
@@ -208,7 +226,7 @@ impl MphfLayer {
/// check against the layer's evidence. The returned slot is only meaningful /// 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 /// when the kmer is actually present in the layer; callers that need an
/// existence test should use [`find`](Self::find) instead. /// 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()) 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 /// 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 /// input slice. Like [`index`](Self::index), no membership check is
/// performed. /// 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() kmers.iter().map(|k| self.mphf.index(&k.raw())).collect()
} }
@@ -237,7 +255,9 @@ impl MphfLayer {
LayerEvidence::Approx { unitigs, .. } => unitigs, LayerEvidence::Approx { unitigs, .. } => unitigs,
LayerEvidence::Hybrid { 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 /// 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` /// 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. /// kmers. The final batch may be shorter when the layer is exhausted.
pub fn iter_kmers_batch(&self, n: usize) -> KmerBatchIter { 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 /// 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 /// `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` /// in that batch within the full layer sequence — i.e. a multiple of `n`
/// except for the final (possibly shorter) batch. /// 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; let mut offset = 0usize;
self.iter_kmers_batch(n).map(move |batch| { self.iter_kmers_batch(n).map(move |batch| {
let base = offset; let base = offset;
@@ -320,11 +346,7 @@ impl Iterator for KmerBatchIter {
break; break;
} }
} }
if batch.is_empty() { if batch.is_empty() { None } else { Some(batch) }
None
} else {
Some(batch)
}
} }
} }
@@ -373,7 +395,7 @@ impl MphfLayer {
let mphf: Mphf = Mphf::load_full(&dir.join(MPHF_FILE)) let mphf: Mphf = Mphf::load_full(&dir.join(MPHF_FILE))
.map_err(|e| OLMError::InvalidLayer(e.to_string()))?; .map_err(|e| OLMError::InvalidLayer(e.to_string()))?;
let mut ev = EvidenceWriter::new(n); let mut ev = EvidenceWriter::new(n);
let mut seen = vec![0u8; (n + 7) / 8]; let mut seen = vec![0u8; (n + 7) / 8];
for (kmer, chunk_id, rank) in unitigs.iter_indexed_canonical_kmers() { for (kmer, chunk_id, rank) in unitigs.iter_indexed_canonical_kmers() {
@@ -382,7 +404,7 @@ impl MphfLayer {
return Err(OLMError::Mphf("slot out of bounds".into())); return Err(OLMError::Mphf("slot out of bounds".into()));
} }
let byte = slot / 8; let byte = slot / 8;
let bit = 1u8 << (slot % 8); let bit = 1u8 << (slot % 8);
if seen[byte] & bit != 0 { if seen[byte] & bit != 0 {
return Err(OLMError::Mphf("duplicate slot".into())); return Err(OLMError::Mphf("duplicate slot".into()));
} }
@@ -398,7 +420,9 @@ impl MphfLayer {
/// Build `fingerprint.bin` from `unitigs.bin` + `mphf.bin`. /// Build `fingerprint.bin` from `unitigs.bin` + `mphf.bin`.
pub fn build_approx_evidence(dir: &Path, b: u8, z: u8) -> OLMResult<usize> { pub fn build_approx_evidence(dir: &Path, b: u8, z: u8) -> OLMResult<usize> {
if b == 0 || b > 64 { 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 { if z == 0 {
return Err(OLMError::InvalidLayer("z must be ≥ 1".into())); return Err(OLMError::InvalidLayer("z must be ≥ 1".into()));
@@ -452,9 +476,8 @@ impl MphfLayer {
// ── Empty layer ─────────────────────────────────────────────────────── // ── Empty layer ───────────────────────────────────────────────────────
if n == 0 { if n == 0 {
let mphf: Mphf = let mphf: Mphf = Mphf::try_new(&[] as &[u64], PtrHashParams::<CubicEps>::default())
Mphf::try_new(&[] as &[u64], PtrHashParams::<CubicEps>::default()) .ok_or_else(|| OLMError::Mphf("construction failed".into()))?;
.ok_or_else(|| OLMError::Mphf("construction failed".into()))?;
mphf.store(&dir.join(MPHF_FILE)) mphf.store(&dir.join(MPHF_FILE))
.map_err(|e| OLMError::InvalidLayer(e.to_string()))?; .map_err(|e| OLMError::InvalidLayer(e.to_string()))?;
match mode { match mode {
@@ -474,9 +497,11 @@ impl MphfLayer {
// ── Pass 1: MPHF via clonable mmap iterator ─────────────────────────── // ── Pass 1: MPHF via clonable mmap iterator ───────────────────────────
let keys = CanonicalKmerIter::new(&unitig_path).map_err(sk_to_olm)?; let keys = CanonicalKmerIter::new(&unitig_path).map_err(sk_to_olm)?;
let mphf: Mphf = let mphf: Mphf = Mphf::new_from_par_iter(
Mphf::new_from_par_iter(n, keys.map(|k| k.raw()).par_bridge(), n,
PtrHashParams::<CubicEps>::default()); keys.map(|k| k.raw()).par_bridge(),
PtrHashParams::<CubicEps>::default(),
);
mphf.store(&dir.join(MPHF_FILE)) mphf.store(&dir.join(MPHF_FILE))
.map_err(|e| OLMError::InvalidLayer(e.to_string()))?; .map_err(|e| OLMError::InvalidLayer(e.to_string()))?;
@@ -489,9 +514,14 @@ impl MphfLayer {
let mut ev = EvidenceWriter::new(n); let mut ev = EvidenceWriter::new(n);
for (kmer, chunk_id, rank) in unitigs2.iter_indexed_canonical_kmers() { for (kmer, chunk_id, rank) in unitigs2.iter_indexed_canonical_kmers() {
let slot = mphf.index(&kmer.raw()); let slot = mphf.index(&kmer.raw());
if slot >= n { return Err(OLMError::Mphf("slot out of bounds".into())); } if slot >= n {
let byte = slot / 8; let bit = 1u8 << (slot % 8); return Err(OLMError::Mphf("slot out of bounds".into()));
if seen[byte] & bit != 0 { return Err(OLMError::Mphf("duplicate slot".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; seen[byte] |= bit;
ev.set(slot, chunk_id as u32, rank as u8); ev.set(slot, chunk_id as u32, rank as u8);
fill_slot(slot, kmer)?; fill_slot(slot, kmer)?;
@@ -504,9 +534,14 @@ impl MphfLayer {
let mut fw = FingerprintVecWriter::new(n, *b); let mut fw = FingerprintVecWriter::new(n, *b);
for (kmer, _, _) in unitigs2.iter_indexed_canonical_kmers() { for (kmer, _, _) in unitigs2.iter_indexed_canonical_kmers() {
let slot = mphf.index(&kmer.raw()); let slot = mphf.index(&kmer.raw());
if slot >= n { return Err(OLMError::Mphf("slot out of bounds".into())); } if slot >= n {
let byte = slot / 8; let bit = 1u8 << (slot % 8); return Err(OLMError::Mphf("slot out of bounds".into()));
if seen[byte] & bit != 0 { return Err(OLMError::Mphf("duplicate slot".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; seen[byte] |= bit;
fw.set(slot, kmer.seq_hash()); fw.set(slot, kmer.seq_hash());
fill_slot(slot, kmer)?; fill_slot(slot, kmer)?;
@@ -519,9 +554,14 @@ impl MphfLayer {
let mut fw = FingerprintVecWriter::new(n, *b); let mut fw = FingerprintVecWriter::new(n, *b);
for (kmer, chunk_id, rank) in unitigs2.iter_indexed_canonical_kmers() { for (kmer, chunk_id, rank) in unitigs2.iter_indexed_canonical_kmers() {
let slot = mphf.index(&kmer.raw()); let slot = mphf.index(&kmer.raw());
if slot >= n { return Err(OLMError::Mphf("slot out of bounds".into())); } if slot >= n {
let byte = slot / 8; let bit = 1u8 << (slot % 8); return Err(OLMError::Mphf("slot out of bounds".into()));
if seen[byte] & bit != 0 { return Err(OLMError::Mphf("duplicate slot".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; seen[byte] |= bit;
ev.set(slot, chunk_id as u32, rank as u8); ev.set(slot, chunk_id as u32, rank as u8);
fw.set(slot, kmer.seq_hash()); fw.set(slot, kmer.seq_hash());
+41 -21
View File
@@ -3,10 +3,8 @@ use std::fs;
use std::path::{Path, PathBuf}; use std::path::{Path, PathBuf};
use obicompactvec::{ use obicompactvec::{
BinaryMatrix, BinaryMatrix, PersistentBitMatrix, PersistentBitMatrixBuilder, PersistentCompactIntMatrix,
PersistentBitMatrix, PersistentBitMatrixBuilder, PersistentCompactIntMatrixBuilder, PersistentSparseBitMatrix,
PersistentCompactIntMatrix, PersistentCompactIntMatrixBuilder,
PersistentSparseBitMatrix,
}; };
use obikseq::CanonicalKmer; use obikseq::CanonicalKmer;
use obiskio::{UnitigFileReader, UnitigFileWriter}; use obiskio::{UnitigFileReader, UnitigFileWriter};
@@ -16,7 +14,7 @@ use crate::layer::meta::IndexMode;
use crate::layer::mphf_layer::MphfLayer; use crate::layer::mphf_layer::MphfLayer;
pub(crate) use crate::layer::mphf_layer::UNITIGS_FILE; pub(crate) use crate::layer::mphf_layer::UNITIGS_FILE;
pub(crate) const COUNTS_DIR: &str = "counts"; pub(crate) const COUNTS_DIR: &str = "counts";
pub(crate) const PRESENCE_DIR: &str = "presence"; pub(crate) const PRESENCE_DIR: &str = "presence";
// ── Trait ───────────────────────────────────────────────────────────────────── // ── Trait ─────────────────────────────────────────────────────────────────────
@@ -82,7 +80,9 @@ pub fn open_data<D: LayerData>(root: &Path, i: usize) -> OLMResult<D> {
impl LayerData for () { impl LayerData for () {
type Item = (); type Item = ();
fn open(_layer_dir: &Path) -> OLMResult<Self> { Ok(()) } fn open(_layer_dir: &Path) -> OLMResult<Self> {
Ok(())
}
fn read(&self, _slot: usize) {} fn read(&self, _slot: usize) {}
} }
@@ -91,7 +91,9 @@ impl LayerData for PersistentCompactIntMatrix {
fn open(layer_dir: &Path) -> OLMResult<Self> { fn open(layer_dir: &Path) -> OLMResult<Self> {
PersistentCompactIntMatrix::open(layer_dir).map_err(OLMError::Io) 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 { impl LayerData for PersistentBitMatrix {
@@ -99,7 +101,9 @@ impl LayerData for PersistentBitMatrix {
fn open(layer_dir: &Path) -> OLMResult<Self> { fn open(layer_dir: &Path) -> OLMResult<Self> {
PersistentBitMatrix::open(layer_dir).map_err(OLMError::Io) 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 { impl LayerData for PersistentSparseBitMatrix {
@@ -107,7 +111,9 @@ impl LayerData for PersistentSparseBitMatrix {
fn open(layer_dir: &Path) -> OLMResult<Self> { fn open(layer_dir: &Path) -> OLMResult<Self> {
PersistentSparseBitMatrix::open(&layer_dir.join(PRESENCE_DIR)).map_err(OLMError::Io) 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 ───────────────────────────────────────────────────────────── // ── LayerContent ─────────────────────────────────────────────────────────────
@@ -202,7 +208,10 @@ impl<D: LayerData> TypedLayer<D> {
} }
pub fn query(&self, kmer: CanonicalKmer) -> Option<Hit<D::Item>> { 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 /// MPHF + evidence membership check only — no data read. For callers
@@ -214,16 +223,18 @@ impl<D: LayerData> TypedLayer<D> {
self.mphf.find(kmer) 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. /// Raw MPHF lookup: kmer → slot, no membership check.
pub fn index(&self, kmer: CanonicalKmer) -> usize { pub fn hash(&self, kmer: CanonicalKmer) -> usize {
self.mphf.index(kmer) self.mphf.hash(kmer)
} }
/// Batch raw MPHF lookup: kmers → slots, no membership check. /// Batch raw MPHF lookup: kmers → slots, no membership check.
pub fn index_batch(&self, kmers: &[CanonicalKmer]) -> Vec<usize> { pub fn hash_batch(&self, kmers: &[CanonicalKmer]) -> Vec<usize> {
self.mphf.index_batch(kmers) self.mphf.hash_batch(kmers)
} }
/// Iterate over all canonical kmers in the layer, in deterministic order. /// 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 /// Iterate over batches, each paired with the zero-based index of the
/// first kmer in the batch. /// 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) self.mphf.enumerate_kmers_batch(n)
} }
@@ -301,8 +315,12 @@ impl TypedLayer<()> {
pub fn init_presence_matrix(layer_dir: &Path, n_kmers: usize) -> OLMResult<()> { pub fn init_presence_matrix(layer_dir: &Path, n_kmers: usize) -> OLMResult<()> {
let presence_dir = layer_dir.join(PRESENCE_DIR); let presence_dir = layer_dir.join(PRESENCE_DIR);
fs::create_dir_all(&presence_dir).map_err(OLMError::Io)?; fs::create_dir_all(&presence_dir).map_err(OLMError::Io)?;
let mut mb = PersistentBitMatrixBuilder::new(n_kmers, &presence_dir).map_err(OLMError::Io)?; let mut mb =
mb.add_col_ones().map_err(OLMError::Io)?.close().map_err(OLMError::Io)?; 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) mb.close().map_err(OLMError::Io)
} }
} }
@@ -318,8 +336,8 @@ impl TypedLayer<PersistentCompactIntMatrix> {
) -> OLMResult<usize> { ) -> OLMResult<usize> {
let n = UnitigFileReader::open_sequential(&out_dir.join(UNITIGS_FILE))?.n_kmers(); let n = UnitigFileReader::open_sequential(&out_dir.join(UNITIGS_FILE))?.n_kmers();
let counts_dir = out_dir.join(COUNTS_DIR); let counts_dir = out_dir.join(COUNTS_DIR);
let mut mb = PersistentCompactIntMatrixBuilder::new(n, &counts_dir) let mut mb =
.map_err(OLMError::Io)?; PersistentCompactIntMatrixBuilder::new(n, &counts_dir).map_err(OLMError::Io)?;
let mut col = mb.add_col().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| { let n_built = MphfLayer::build(out_dir, block_bits, mode, &mut |slot, kmer| {
col.set(slot, count_of(kmer)); col.set(slot, count_of(kmer));
@@ -336,7 +354,9 @@ impl TypedLayer<PersistentCompactIntMatrix> {
mode: &IndexMode, mode: &IndexMode,
counts: &HashMap<CanonicalKmer, u32>, counts: &HashMap<CanonicalKmer, u32>,
) -> OLMResult<usize> { ) -> 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 rayon::prelude::*;
use obikseq::CanonicalKmer;
use obikindex::layer::Layer; use obikindex::layer::Layer;
use obikseq::CanonicalKmer;
use obisys::progress_bar; use obisys::progress_bar;
use obikindex::{KmerIndex, OKIResult}; use obikindex::{KmerIndex, OKIResult};
@@ -45,11 +45,7 @@ pub(super) struct PartitionCache {
} }
impl PartitionCache { impl PartitionCache {
pub(super) fn build( pub(super) fn build(index: &KmerIndex, n_parts: usize, with_counts: bool) -> OKIResult<Self> {
index: &KmerIndex,
n_parts: usize,
with_counts: bool,
) -> OKIResult<Self> {
let pb = progress_bar("open_partitions", n_parts as u64, "partitions"); let pb = progress_bar("open_partitions", n_parts as u64, "partitions");
let built: Vec<(Vec<Layer>, usize)> = (0..n_parts) let built: Vec<(Vec<Layer>, usize)> = (0..n_parts)
.into_par_iter() .into_par_iter()
@@ -200,7 +196,7 @@ impl PartitionCache {
let mat = &mats[li]; let mat = &mats[li];
let variants: Vec<CanonicalKmer> = let variants: Vec<CanonicalKmer> =
entries.iter().map(|&(variant, _, _)| variant).collect(); 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 let hits: Vec<(usize, usize, u8)> = slots
.into_iter() .into_iter()
.zip(entries.iter()) .zip(entries.iter())
+43 -23
View File
@@ -53,19 +53,19 @@ use std::sync::atomic::{AtomicU8, Ordering};
use rayon::prelude::*; use rayon::prelude::*;
use obikseq::CanonicalKmer;
use obikindex::layer::meta::PartitionMeta; use obikindex::layer::meta::PartitionMeta;
use obikseq::CanonicalKmer;
use obipipeline::{ThrottleGuard, throttle}; use obipipeline::{ThrottleGuard, throttle};
use obikindex::{OKIError, OKIResult};
use obikindex::KmerIndex; use obikindex::KmerIndex;
use obikindex::{OKIError, OKIResult};
use obikindex::layer::Layer; use obikindex::layer::Layer;
use super::cache::PartitionCache; use super::cache::PartitionCache;
use super::helpers::central_base; use super::helpers::central_base;
use super::iter::{SiblingEntry, SiblingLayerExt}; 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- /// 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 /// 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, selection: &Selection,
mut on_family: impl FnMut(usize, FamilyMask, &[u8]), mut on_family: impl FnMut(usize, FamilyMask, &[u8]),
) -> OKIResult<()> { ) -> 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 meta = PartitionMeta::load(index_dir).map_err(olm_to_ok)?;
let annex = Arc::new(SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?); 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 mat = Layer::open(layer_dir, &meta.mode, with_counts).map_err(olm_to_ok)?;
let n_cols = mat.n_cols().min(n_genomes); 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 // Streamed straight from `iter_minorants_batch` (zips this layer's own
// `iter_kmers()` with the annex, both in iteration order — never an // `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 // (see the module docs). `.scan()` computes each batch's starting
// family index lazily, mirroring what the eager `chunks()`+running // family index lazily, mirroring what the eager `chunks()`+running
// `offset` used to do. // `offset` used to do.
let batches = ctx.mat.iter_minorants_batch(annex, FAMILY_BATCH).scan(0usize, |offset, batch| { let batches =
let start = *offset; ctx.mat
*offset += batch.len(); .iter_minorants_batch(annex, FAMILY_BATCH)
Some((start, batch)) .scan(0usize, |offset, batch| {
}); let start = *offset;
*offset += batch.len();
Some((start, batch))
});
let n_workers = obisys::effective_parallelism(); let n_workers = obisys::effective_parallelism();
let capacity = 4; let capacity = 4;
@@ -288,7 +299,7 @@ pub(super) fn scan_layer_families(
// — a pure MPHF lookup, no evidence check, since these are // — a pure MPHF lookup, no evidence check, since these are
// this layer's own kmers, known members by construction. // this layer's own kmers, known members by construction.
let kmers: Vec<CanonicalKmer> = batch.entries.iter().map(|e| e.kmer).collect(); 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(); let mut carries: Vec<Vec<bool>> = (0..ctx.n_cols).map(|_| Vec::new()).collect();
ctx.mat.fill_sub_matrix_carries(&slots, &mut carries); ctx.mat.fill_sub_matrix_carries(&slots, &mut carries);
for (g, col) in carries.iter().enumerate() { 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 // (one batch at a time, never several concurrently — see the
// module docs), each thread owning one partition's queries // module docs), each thread owning one partition's queries
// contiguously until this batch is done. // 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(); let fast_mode = cache.fast_mode();
batch.outgoing.par_iter().enumerate().filter(|(_, q)| !q.is_empty()).for_each(|(dest, queries)| { batch
let on_hit = |i: usize, base: u8, g: usize| { .outgoing
genome_mask[i * n_genomes + g].fetch_or(1 << base, Ordering::Relaxed); .par_iter()
}; .enumerate()
if fast_mode { .filter(|(_, q)| !q.is_empty())
cache.find_presence_batch_fast(dest, queries, n_genomes, on_hit); .for_each(|(dest, queries)| {
} else { let on_hit = |i: usize, base: u8, g: usize| {
cache.find_presence_batch(dest, queries, n_genomes, on_hit); 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 { for i in 0..n {
let family_idx = batch.start_family_idx + i; let family_idx = batch.start_family_idx + i;
@@ -347,7 +364,10 @@ pub(super) fn scan_layer_families(
next_expected += n; 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(()) Ok(())
} }