Push zunrplorkwkt #70

Merged
coissac merged 93 commits from push-zunrplorkwkt into main 2026-08-28 23:15:38 +00:00
5 changed files with 294 additions and 145 deletions
Showing only changes of commit 346095b9eb - Show all commits
+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(())
}