introduce fast mode for optimized sibling presence checks
Centralize the layer count validation into PartitionCache and track it via a new fast_mode flag. Extend query tuples to include a pre-resolved destination layer index, enabling a fast-path batch lookup that bypasses per-layer probing when enabled. Refactor neighbor iteration and hit resolution to eliminate duplication and conditionally dispatch to the optimized path based on the cache state.
This commit is contained in:
Eric Coissac
+161
-57
@@ -14,16 +14,20 @@ const MAGIC: [u8; 4] = *b"PBIV";
|
||||
const HEADER_SIZE: usize = 16;
|
||||
|
||||
#[inline]
|
||||
pub(crate) fn n_words(n: usize) -> usize { n.div_ceil(64) }
|
||||
pub(crate) fn n_words(n: usize) -> usize {
|
||||
n.div_ceil(64)
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn n_bytes_for_words(n: usize) -> usize { n_words(n) * 8 }
|
||||
fn n_bytes_for_words(n: usize) -> usize {
|
||||
n_words(n) * 8
|
||||
}
|
||||
|
||||
// ── PersistentBitVec ──────────────────────────────────────────────────────────
|
||||
|
||||
pub struct PersistentBitVec {
|
||||
mmap: Mmap,
|
||||
n: usize,
|
||||
n: usize,
|
||||
path: PathBuf,
|
||||
}
|
||||
|
||||
@@ -31,18 +35,31 @@ impl PersistentBitVec {
|
||||
pub fn open(path: &Path) -> io::Result<Self> {
|
||||
let mmap = unsafe { Mmap::map(&File::open(path)?)? };
|
||||
if mmap.len() < HEADER_SIZE {
|
||||
return Err(io::Error::new(io::ErrorKind::InvalidData, "PBIV file too short"));
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::InvalidData,
|
||||
"PBIV file too short",
|
||||
));
|
||||
}
|
||||
if &mmap[0..4] != &MAGIC {
|
||||
return Err(io::Error::new(io::ErrorKind::InvalidData, "bad PBIV magic"));
|
||||
}
|
||||
let n = u64::from_le_bytes(mmap[8..16].try_into().unwrap()) as usize;
|
||||
Ok(Self { mmap, n, path: path.to_path_buf() })
|
||||
Ok(Self {
|
||||
mmap,
|
||||
n,
|
||||
path: path.to_path_buf(),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn path(&self) -> &Path { &self.path }
|
||||
pub fn len(&self) -> usize { self.n }
|
||||
pub fn is_empty(&self) -> bool { self.n == 0 }
|
||||
pub fn path(&self) -> &Path {
|
||||
&self.path
|
||||
}
|
||||
pub fn len(&self) -> usize {
|
||||
self.n
|
||||
}
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.n == 0
|
||||
}
|
||||
|
||||
pub fn get(&self, slot: usize) -> bool {
|
||||
(self.mmap[HEADER_SIZE + (slot >> 3)] >> (slot & 7)) & 1 != 0
|
||||
@@ -64,7 +81,9 @@ impl PersistentBitVec {
|
||||
pub fn fill_batch(&self, slots: &[usize], out: &mut [bool]) {
|
||||
assert_eq!(slots.len(), out.len());
|
||||
let n = slots.len();
|
||||
if n == 0 { return; }
|
||||
if n == 0 {
|
||||
return;
|
||||
}
|
||||
let mut perm: Vec<usize> = (0..n).collect();
|
||||
perm.sort_by_key(|&i| slots[i]);
|
||||
let sorted: Vec<usize> = perm.iter().map(|&i| slots[i]).collect();
|
||||
@@ -86,7 +105,7 @@ impl PersistentBitVec {
|
||||
|
||||
// SAFETY: mmap is page-aligned, HEADER_SIZE=16 divisible by 8 → u64-aligned.
|
||||
fn data_words(&self) -> &[u64] {
|
||||
let nw = n_words(self.n);
|
||||
let nw = n_words(self.n);
|
||||
let ptr = self.mmap[HEADER_SIZE..].as_ptr() as *const u64;
|
||||
unsafe { std::slice::from_raw_parts(ptr, nw) }
|
||||
}
|
||||
@@ -95,10 +114,16 @@ impl PersistentBitVec {
|
||||
BitSliceView::new(self.data_words(), self.n)
|
||||
}
|
||||
|
||||
pub fn words(&self) -> &[u64] { self.data_words() }
|
||||
pub fn words(&self) -> &[u64] {
|
||||
self.data_words()
|
||||
}
|
||||
|
||||
pub fn count_ones(&self) -> u64 { self.view().count_ones() }
|
||||
pub fn count_zeros(&self) -> u64 { self.view().count_zeros() }
|
||||
pub fn count_ones(&self) -> u64 {
|
||||
self.view().count_ones()
|
||||
}
|
||||
pub fn count_zeros(&self) -> u64 {
|
||||
self.view().count_zeros()
|
||||
}
|
||||
|
||||
pub fn partial_jaccard_dist(&self, other: &PersistentBitVec) -> (u64, u64) {
|
||||
self.view().partial_jaccard_dist(other.view())
|
||||
@@ -111,22 +136,28 @@ impl PersistentBitVec {
|
||||
}
|
||||
|
||||
pub fn iter(&self) -> BitIter<'_> {
|
||||
BitIter { words: self.data_words(), slot: 0, n: self.n }
|
||||
BitIter {
|
||||
words: self.data_words(),
|
||||
slot: 0,
|
||||
n: self.n,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a> IntoIterator for &'a PersistentBitVec {
|
||||
type Item = bool;
|
||||
type IntoIter = BitIter<'a>;
|
||||
fn into_iter(self) -> BitIter<'a> { self.iter() }
|
||||
fn into_iter(self) -> BitIter<'a> {
|
||||
self.iter()
|
||||
}
|
||||
}
|
||||
|
||||
// ── BitIter ───────────────────────────────────────────────────────────────────
|
||||
|
||||
pub struct BitIter<'a> {
|
||||
words: &'a [u64],
|
||||
slot: usize,
|
||||
n: usize,
|
||||
slot: usize,
|
||||
n: usize,
|
||||
}
|
||||
|
||||
impl ExactSizeIterator for BitIter<'_> {}
|
||||
@@ -134,7 +165,9 @@ impl ExactSizeIterator for BitIter<'_> {}
|
||||
impl Iterator for BitIter<'_> {
|
||||
type Item = bool;
|
||||
fn next(&mut self) -> Option<bool> {
|
||||
if self.slot >= self.n { return None; }
|
||||
if self.slot >= self.n {
|
||||
return None;
|
||||
}
|
||||
let v = (self.words[self.slot >> 6] >> (self.slot & 63)) & 1 != 0;
|
||||
self.slot += 1;
|
||||
Some(v)
|
||||
@@ -149,7 +182,7 @@ impl Iterator for BitIter<'_> {
|
||||
|
||||
pub struct PersistentBitVecBuilder {
|
||||
mmap: MmapMut,
|
||||
n: usize,
|
||||
n: usize,
|
||||
path: PathBuf,
|
||||
}
|
||||
|
||||
@@ -157,7 +190,10 @@ impl PersistentBitVecBuilder {
|
||||
pub fn new(n: usize, path: &Path) -> io::Result<Self> {
|
||||
let file_size = HEADER_SIZE + n_bytes_for_words(n);
|
||||
let mut file = OpenOptions::new()
|
||||
.read(true).write(true).create(true).truncate(true)
|
||||
.read(true)
|
||||
.write(true)
|
||||
.create(true)
|
||||
.truncate(true)
|
||||
.open(path)?;
|
||||
file.write_all(&MAGIC)?;
|
||||
file.write_all(&[0u8; 4])?;
|
||||
@@ -165,20 +201,31 @@ impl PersistentBitVecBuilder {
|
||||
file.seek(SeekFrom::Start(0))?;
|
||||
file.set_len(file_size as u64)?;
|
||||
let mmap = unsafe { MmapMut::map_mut(&file)? };
|
||||
Ok(Self { mmap, n, path: path.to_path_buf() })
|
||||
Ok(Self {
|
||||
mmap,
|
||||
n,
|
||||
path: path.to_path_buf(),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn from_raw_bytes(bytes: &[u8], n: usize, path: &Path) -> io::Result<Self> {
|
||||
let file_size = HEADER_SIZE + n_bytes_for_words(n);
|
||||
let file = OpenOptions::new()
|
||||
.read(true).write(true).create(true).truncate(true)
|
||||
.read(true)
|
||||
.write(true)
|
||||
.create(true)
|
||||
.truncate(true)
|
||||
.open(path)?;
|
||||
file.set_len(file_size as u64)?;
|
||||
let mut mmap = unsafe { MmapMut::map_mut(&file)? };
|
||||
mmap[0..4].copy_from_slice(&MAGIC);
|
||||
mmap[8..16].copy_from_slice(&(n as u64).to_le_bytes());
|
||||
mmap[HEADER_SIZE..HEADER_SIZE + bytes.len()].copy_from_slice(bytes);
|
||||
Ok(Self { mmap, n, path: path.to_path_buf() })
|
||||
Ok(Self {
|
||||
mmap,
|
||||
n,
|
||||
path: path.to_path_buf(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Create an all-ones bit vector of length `n` at `path`.
|
||||
@@ -186,10 +233,13 @@ impl PersistentBitVecBuilder {
|
||||
/// More efficient than `new(n, path)` + `not()`: the data is written as
|
||||
/// 0xFF bytes in a single sequential pass, with no intermediate all-zeros state.
|
||||
pub fn new_ones(n: usize, path: &Path) -> io::Result<Self> {
|
||||
let nw = n_words(n);
|
||||
let nw = n_words(n);
|
||||
let file_size = HEADER_SIZE + nw * 8;
|
||||
let mut file = OpenOptions::new()
|
||||
.read(true).write(true).create(true).truncate(true)
|
||||
let mut file = OpenOptions::new()
|
||||
.read(true)
|
||||
.write(true)
|
||||
.create(true)
|
||||
.truncate(true)
|
||||
.open(path)?;
|
||||
file.write_all(&MAGIC)?;
|
||||
file.write_all(&[0u8; 4])?;
|
||||
@@ -201,11 +251,15 @@ impl PersistentBitVecBuilder {
|
||||
// Clear padding bits in the last word so trailing bits are always 0.
|
||||
let rem = n % 64;
|
||||
if rem != 0 {
|
||||
let ptr = mmap[HEADER_SIZE..].as_mut_ptr() as *mut u64;
|
||||
let ptr = mmap[HEADER_SIZE..].as_mut_ptr() as *mut u64;
|
||||
let words = unsafe { std::slice::from_raw_parts_mut(ptr, nw) };
|
||||
words[nw - 1] &= (1u64 << rem) - 1;
|
||||
}
|
||||
Ok(Self { mmap, n, path: path.to_path_buf() })
|
||||
Ok(Self {
|
||||
mmap,
|
||||
n,
|
||||
path: path.to_path_buf(),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn build_from(source: &PersistentBitVec, path: &Path) -> io::Result<Self> {
|
||||
@@ -213,14 +267,25 @@ impl PersistentBitVecBuilder {
|
||||
let file = OpenOptions::new().read(true).write(true).open(path)?;
|
||||
let mmap = unsafe { MmapMut::map_mut(&file)? };
|
||||
let n = source.len();
|
||||
Ok(Self { mmap, n, path: path.to_path_buf() })
|
||||
Ok(Self {
|
||||
mmap,
|
||||
n,
|
||||
path: path.to_path_buf(),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn build_from_counts(source: &PersistentCompactIntVec, threshold: u32, path: &Path) -> io::Result<Self> {
|
||||
pub fn build_from_counts(
|
||||
source: &PersistentCompactIntVec,
|
||||
threshold: u32,
|
||||
path: &Path,
|
||||
) -> io::Result<Self> {
|
||||
let n = source.len();
|
||||
let file_size = HEADER_SIZE + n_bytes_for_words(n);
|
||||
let mut file = OpenOptions::new()
|
||||
.read(true).write(true).create(true).truncate(true)
|
||||
.read(true)
|
||||
.write(true)
|
||||
.create(true)
|
||||
.truncate(true)
|
||||
.open(path)?;
|
||||
file.write_all(&MAGIC)?;
|
||||
file.write_all(&[0u8; 4])?;
|
||||
@@ -229,22 +294,32 @@ impl PersistentBitVecBuilder {
|
||||
file.set_len(file_size as u64)?;
|
||||
let mut mmap = unsafe { MmapMut::map_mut(&file)? };
|
||||
{
|
||||
let nw = n_words(n);
|
||||
let nw = n_words(n);
|
||||
let ptr = mmap[HEADER_SIZE..].as_mut_ptr() as *mut u64;
|
||||
let words = unsafe { std::slice::from_raw_parts_mut(ptr, nw) };
|
||||
for (slot, count) in source.iter().enumerate() {
|
||||
if count >= threshold { words[slot >> 6] |= 1u64 << (slot & 63); }
|
||||
if count >= threshold {
|
||||
words[slot >> 6] |= 1u64 << (slot & 63);
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(Self { mmap, n, path: path.to_path_buf() })
|
||||
Ok(Self {
|
||||
mmap,
|
||||
n,
|
||||
path: path.to_path_buf(),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn build_from_presence(source: &PersistentCompactIntVec, path: &Path) -> io::Result<Self> {
|
||||
Self::build_from_counts(source, 1, path)
|
||||
}
|
||||
|
||||
pub fn len(&self) -> usize { self.n }
|
||||
pub fn is_empty(&self) -> bool { self.n == 0 }
|
||||
pub fn len(&self) -> usize {
|
||||
self.n
|
||||
}
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.n == 0
|
||||
}
|
||||
|
||||
pub fn get(&self, slot: usize) -> bool {
|
||||
(self.mmap[HEADER_SIZE + (slot >> 3)] >> (slot & 7)) & 1 != 0
|
||||
@@ -252,19 +327,22 @@ impl PersistentBitVecBuilder {
|
||||
|
||||
pub fn set(&mut self, slot: usize, value: bool) {
|
||||
let bit = 1u64 << (slot & 63);
|
||||
if value { self.data_words_mut()[slot >> 6] |= bit; }
|
||||
else { self.data_words_mut()[slot >> 6] &= !bit; }
|
||||
if value {
|
||||
self.data_words_mut()[slot >> 6] |= bit;
|
||||
} else {
|
||||
self.data_words_mut()[slot >> 6] &= !bit;
|
||||
}
|
||||
}
|
||||
|
||||
fn data_words(&self) -> &[u64] {
|
||||
let nw = n_words(self.n);
|
||||
let nw = n_words(self.n);
|
||||
let ptr = self.mmap[HEADER_SIZE..].as_ptr() as *const u64;
|
||||
unsafe { std::slice::from_raw_parts(ptr, nw) }
|
||||
}
|
||||
|
||||
// SAFETY: same alignment argument as PersistentBitVec::data_words.
|
||||
fn data_words_mut(&mut self) -> &mut [u64] {
|
||||
let nw = n_words(self.n);
|
||||
let nw = n_words(self.n);
|
||||
let ptr = self.mmap[HEADER_SIZE..].as_mut_ptr() as *mut u64;
|
||||
unsafe { std::slice::from_raw_parts_mut(ptr, nw) }
|
||||
}
|
||||
@@ -273,7 +351,9 @@ impl PersistentBitVecBuilder {
|
||||
BitSliceView::new(self.data_words(), self.n)
|
||||
}
|
||||
|
||||
pub fn words(&self) -> &[u64] { self.data_words() }
|
||||
pub fn words(&self) -> &[u64] {
|
||||
self.data_words()
|
||||
}
|
||||
|
||||
pub fn copy_from(&mut self, src: BitSliceView<'_>) {
|
||||
assert_eq!(self.n, src.len(), "BitSliceView length mismatch");
|
||||
@@ -282,25 +362,35 @@ impl PersistentBitVecBuilder {
|
||||
|
||||
pub fn and(&mut self, other: BitSliceView<'_>) {
|
||||
assert_eq!(self.n, other.len(), "BitSliceView length mismatch");
|
||||
for (w, &o) in self.data_words_mut().iter_mut().zip(other.words()) { *w &= o; }
|
||||
for (w, &o) in self.data_words_mut().iter_mut().zip(other.words()) {
|
||||
*w &= o;
|
||||
}
|
||||
}
|
||||
|
||||
pub fn or(&mut self, other: BitSliceView<'_>) {
|
||||
assert_eq!(self.n, other.len(), "BitSliceView length mismatch");
|
||||
for (w, &o) in self.data_words_mut().iter_mut().zip(other.words()) { *w |= o; }
|
||||
for (w, &o) in self.data_words_mut().iter_mut().zip(other.words()) {
|
||||
*w |= o;
|
||||
}
|
||||
}
|
||||
|
||||
pub fn xor(&mut self, other: BitSliceView<'_>) {
|
||||
assert_eq!(self.n, other.len(), "BitSliceView length mismatch");
|
||||
for (w, &o) in self.data_words_mut().iter_mut().zip(other.words()) { *w ^= o; }
|
||||
for (w, &o) in self.data_words_mut().iter_mut().zip(other.words()) {
|
||||
*w ^= o;
|
||||
}
|
||||
}
|
||||
|
||||
pub fn not(&mut self) {
|
||||
let rem = self.n % 64;
|
||||
let rem = self.n % 64;
|
||||
let words = self.data_words_mut();
|
||||
for w in words.iter_mut() { *w ^= u64::MAX; }
|
||||
for w in words.iter_mut() {
|
||||
*w ^= u64::MAX;
|
||||
}
|
||||
if rem != 0 {
|
||||
if let Some(last) = words.last_mut() { *last &= (1u64 << rem) - 1; }
|
||||
if let Some(last) = words.last_mut() {
|
||||
*last &= (1u64 << rem) - 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -312,17 +402,21 @@ impl PersistentBitVecBuilder {
|
||||
let words = self.data_words_mut();
|
||||
let nw = n_words(n);
|
||||
for wi in 0..nw {
|
||||
let base = wi * 64;
|
||||
let base = wi * 64;
|
||||
let limit = (base + 64).min(n);
|
||||
let mut mask = 0u64;
|
||||
for bit in 0..(limit - base) {
|
||||
let b = primary[base + bit];
|
||||
if b < 255 && pred(b as u32) { mask |= 1u64 << bit; }
|
||||
if b < 255 && pred(b as u32) {
|
||||
mask |= 1u64 << bit;
|
||||
}
|
||||
}
|
||||
words[wi] |= mask;
|
||||
}
|
||||
for (slot, val) in col.overflow_entries() {
|
||||
if pred(val) { words[slot >> 6] |= 1u64 << (slot & 63); }
|
||||
if pred(val) {
|
||||
words[slot >> 6] |= 1u64 << (slot & 63);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -334,17 +428,21 @@ impl PersistentBitVecBuilder {
|
||||
let words = self.data_words_mut();
|
||||
let nw = n_words(n);
|
||||
for wi in 0..nw {
|
||||
let base = wi * 64;
|
||||
let base = wi * 64;
|
||||
let limit = (base + 64).min(n);
|
||||
let mut mask = 0u64;
|
||||
for bit in 0..(limit - base) {
|
||||
let b = primary[base + bit];
|
||||
if b < 255 && !pred(b as u32) { mask |= 1u64 << bit; }
|
||||
if b < 255 && !pred(b as u32) {
|
||||
mask |= 1u64 << bit;
|
||||
}
|
||||
}
|
||||
words[wi] &= !mask;
|
||||
}
|
||||
for (slot, val) in col.overflow_entries() {
|
||||
if !pred(val) { words[slot >> 6] &= !(1u64 << (slot & 63)); }
|
||||
if !pred(val) {
|
||||
words[slot >> 6] &= !(1u64 << (slot & 63));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -356,17 +454,21 @@ impl PersistentBitVecBuilder {
|
||||
let words = self.data_words_mut();
|
||||
let nw = n_words(n);
|
||||
for wi in 0..nw {
|
||||
let base = wi * 64;
|
||||
let base = wi * 64;
|
||||
let limit = (base + 64).min(n);
|
||||
let mut mask = 0u64;
|
||||
for bit in 0..(limit - base) {
|
||||
let b = primary[base + bit];
|
||||
if b < 255 && pred(b as u32) { mask |= 1u64 << bit; }
|
||||
if b < 255 && pred(b as u32) {
|
||||
mask |= 1u64 << bit;
|
||||
}
|
||||
}
|
||||
words[wi] ^= mask;
|
||||
}
|
||||
for (slot, val) in col.overflow_entries() {
|
||||
if pred(val) { words[slot >> 6] ^= 1u64 << (slot & 63); }
|
||||
if pred(val) {
|
||||
words[slot >> 6] ^= 1u64 << (slot & 63);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -374,7 +476,9 @@ impl PersistentBitVecBuilder {
|
||||
self.view().iter()
|
||||
}
|
||||
|
||||
pub fn close(self) -> io::Result<()> { self.mmap.flush() }
|
||||
pub fn close(self) -> io::Result<()> {
|
||||
self.mmap.flush()
|
||||
}
|
||||
|
||||
pub fn finish(self) -> io::Result<PersistentBitVec> {
|
||||
let path = self.path.clone();
|
||||
|
||||
@@ -132,20 +132,12 @@ fn build_layer_sibling_annex(
|
||||
let n = mphf.n();
|
||||
|
||||
// Whether this layer's family-member fields can carry a real layer
|
||||
// number (fast path, exploited by a future scan-time reader) or must
|
||||
// stay presence-only (today's behaviour) — a single fact about the
|
||||
// whole index, decided once here, never stored: `n_layers` is
|
||||
// guaranteed identical across every partition (a merge adds one layer
|
||||
// to all of them at once), and `FamilyMask`'s 3-bit field only has room
|
||||
// for layers `0..=6`.
|
||||
let fast_mode = meta.n_layers <= 7;
|
||||
if !fast_mode {
|
||||
tracing::warn!(
|
||||
"layer {l} ({layer_dir:?}): index has {} layers (>7) — sibling-annex fast layer \
|
||||
lookup disabled for this build; consider compacting this index",
|
||||
meta.n_layers
|
||||
);
|
||||
}
|
||||
// number (exploited by `family_scan::scan_layer_families`'s fast path)
|
||||
// or must stay presence-only — a single fact about the whole index,
|
||||
// decided once by `cache` (built from the same `PartitionMeta` this
|
||||
// function would otherwise re-derive) so the writer and every reader
|
||||
// agree; see `PartitionCache::fast_mode`'s docs.
|
||||
let fast_mode = cache.fast_mode();
|
||||
|
||||
// ── The annex file itself is the reconciliation state, indexed by
|
||||
// this layer's k-mer iteration order (the physical layout of
|
||||
|
||||
@@ -117,18 +117,28 @@ pub(super) struct PartitionCache {
|
||||
/// both [`obikindex::KmerIndex::build_sibling_annex`] and
|
||||
/// [`obikindex::KmerIndex::sibling_annex_stats`].
|
||||
mats: Vec<Vec<Mat>>,
|
||||
/// Whether every `FamilyMask` field in this index's sibling annexes can
|
||||
/// be trusted as a real layer index (`n_layers <= 7`, the same decision
|
||||
/// `build_layer_sibling_annex` makes when writing them) — computed once
|
||||
/// here, from `PartitionMeta::n_layers` (not `mats[p].len()`, which can
|
||||
/// undercount if a layer failed to open), reused by both the annex
|
||||
/// writer (`build.rs`) and reader (`family_scan.rs`) so the two never
|
||||
/// disagree. `n_layers` is guaranteed identical across every partition
|
||||
/// (a merge adds one layer to all of them at once), so the first
|
||||
/// non-empty partition's count speaks for the whole index.
|
||||
fast_mode: bool,
|
||||
}
|
||||
|
||||
impl PartitionCache {
|
||||
pub(super) fn build(partition: &KmerPartition, n_parts: usize, with_counts: bool) -> OKIResult<Self> {
|
||||
let pb = progress_bar("open_partitions", n_parts as u64, "partitions");
|
||||
let built: Vec<Vec<Mat>> = (0..n_parts)
|
||||
let built: Vec<(Vec<Mat>, usize)> = (0..n_parts)
|
||||
.into_par_iter()
|
||||
.map(|part| -> OKIResult<Vec<Mat>> {
|
||||
.map(|part| -> OKIResult<(Vec<Mat>, usize)> {
|
||||
let index_dir = partition.part_dir(part).join(INDEX_SUBDIR);
|
||||
if !index_dir.exists() {
|
||||
pb.inc(1);
|
||||
return Ok(Vec::new());
|
||||
return Ok((Vec::new(), 0));
|
||||
}
|
||||
let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?;
|
||||
let mut mats = Vec::with_capacity(meta.n_layers);
|
||||
@@ -144,11 +154,27 @@ impl PartitionCache {
|
||||
mats.push(mat);
|
||||
}
|
||||
pb.inc(1);
|
||||
Ok(mats)
|
||||
Ok((mats, meta.n_layers))
|
||||
})
|
||||
.collect::<OKIResult<Vec<_>>>()?;
|
||||
pb.finish_and_clear();
|
||||
Ok(Self { mats: built })
|
||||
|
||||
let n_layers = built.iter().map(|&(_, n)| n).find(|&n| n > 0).unwrap_or(0);
|
||||
let fast_mode = n_layers <= 7;
|
||||
if !fast_mode {
|
||||
tracing::warn!(
|
||||
"index has {n_layers} layers (>7) — sibling-annex fast layer lookup disabled; \
|
||||
consider compacting this index"
|
||||
);
|
||||
}
|
||||
|
||||
let mats = built.into_iter().map(|(m, _)| m).collect();
|
||||
Ok(Self { mats, fast_mode })
|
||||
}
|
||||
|
||||
/// See the [`fast_mode`](Self::fast_mode) field docs.
|
||||
pub(super) fn fast_mode(&self) -> bool {
|
||||
self.fast_mode
|
||||
}
|
||||
|
||||
/// Existence lookup of `variant` in partition `dest_partition`: tries
|
||||
@@ -186,10 +212,16 @@ impl PartitionCache {
|
||||
/// `PersistentBitMatrix::get` dominated wall-clock time, mostly blocked
|
||||
/// on page faults, even after every partition/batch locality fix above
|
||||
/// this in the traversal.
|
||||
/// `queries`' 4th element (the annex-recorded layer) is ignored here —
|
||||
/// every layer is probed with evidence regardless, since this is the
|
||||
/// `!fast_mode` fallback where that field isn't trustworthy (see
|
||||
/// [`find_presence_batch_fast`](Self::find_presence_batch_fast)). Same
|
||||
/// tuple shape as the fast path purely so callers don't need two
|
||||
/// query representations.
|
||||
pub(super) fn find_presence_batch(
|
||||
&self,
|
||||
dest_partition: usize,
|
||||
queries: &[(CanonicalKmer, usize, u8)],
|
||||
queries: &[(CanonicalKmer, usize, u8, u8)],
|
||||
n_genomes: usize,
|
||||
mut on_hit: impl FnMut(usize, u8, usize),
|
||||
) {
|
||||
@@ -198,7 +230,7 @@ impl PartitionCache {
|
||||
// First hit wins, same semantics as the old per-query loop (a
|
||||
// variant present in an earlier layer shadows later ones).
|
||||
let mut by_layer: Vec<Vec<(usize, usize, u8)>> = vec![Vec::new(); mats.len()];
|
||||
for &(variant, family_idx, base) in queries {
|
||||
for &(variant, family_idx, base, _layer) in queries {
|
||||
for (li, mat) in mats.iter().enumerate() {
|
||||
if let Some(slot) = mat.find_slot(variant) {
|
||||
by_layer[li].push((slot, family_idx, base));
|
||||
@@ -211,17 +243,75 @@ impl PartitionCache {
|
||||
if hits.is_empty() {
|
||||
continue;
|
||||
}
|
||||
resolve_layer_hits(&mats[li], &hits, n_genomes, &mut on_hit);
|
||||
}
|
||||
}
|
||||
|
||||
/// Same contract as [`find_presence_batch`](Self::find_presence_batch),
|
||||
/// but only valid when [`fast_mode`](Self::fast_mode) is true for this
|
||||
/// whole index: every query's 4th element is trusted outright as its
|
||||
/// real destination layer (recorded once, evidence-checked, at
|
||||
/// annex-build time by [`find`](Self::find)) — no per-layer probing, a
|
||||
/// direct `index_batch` (no evidence) straight to that layer instead.
|
||||
/// Safe for the same reason `build_layer_sibling_annex`'s own
|
||||
/// iteration-pipeline lookups are: this exact variant's membership in
|
||||
/// this exact layer was already positively established once: replaying
|
||||
/// it via `index`/`index_batch` here isn't a fresh, unverified
|
||||
/// assumption, it's reusing a fact already paid for.
|
||||
pub(super) fn find_presence_batch_fast(
|
||||
&self,
|
||||
dest_partition: usize,
|
||||
queries: &[(CanonicalKmer, usize, u8, u8)],
|
||||
n_genomes: usize,
|
||||
mut on_hit: impl FnMut(usize, u8, usize),
|
||||
) {
|
||||
let Some(mats) = self.mats.get(dest_partition) else { return };
|
||||
|
||||
let mut by_layer: Vec<Vec<(CanonicalKmer, usize, u8)>> = vec![Vec::new(); mats.len()];
|
||||
for &(variant, family_idx, base, layer) in queries {
|
||||
if let Some(bucket) = by_layer.get_mut(layer as usize) {
|
||||
bucket.push((variant, family_idx, base));
|
||||
}
|
||||
}
|
||||
|
||||
for (li, entries) in by_layer.into_iter().enumerate() {
|
||||
if entries.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let mat = &mats[li];
|
||||
let n_cols = mat.n_cols().min(n_genomes);
|
||||
let slots: Vec<usize> = hits.iter().map(|&(slot, _, _)| slot).collect();
|
||||
let mut carries: Vec<Vec<bool>> = (0..n_cols).map(|_| Vec::new()).collect();
|
||||
mat.fill_sub_matrix_carries(&slots, &mut carries);
|
||||
for (g, col) in carries.iter().enumerate() {
|
||||
for (&(_, family_idx, base), &carries_it) in hits.iter().zip(col.iter()) {
|
||||
if carries_it {
|
||||
on_hit(family_idx, base, g);
|
||||
}
|
||||
}
|
||||
let variants: Vec<CanonicalKmer> = entries.iter().map(|&(variant, _, _)| variant).collect();
|
||||
let slots = mat.index_batch(&variants);
|
||||
let hits: Vec<(usize, usize, u8)> = slots
|
||||
.into_iter()
|
||||
.zip(entries.iter())
|
||||
.map(|(slot, &(_, family_idx, base))| (slot, family_idx, base))
|
||||
.collect();
|
||||
resolve_layer_hits(mat, &hits, n_genomes, &mut on_hit);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Shared tail of both [`PartitionCache::find_presence_batch`] and
|
||||
/// [`PartitionCache::find_presence_batch_fast`]: given one layer's already-
|
||||
/// resolved `(slot, family_idx, base)` hits, sweep genome-major (via
|
||||
/// `fill_sub_matrix_carries`, buffer-reusing, no per-call allocation) and
|
||||
/// call `on_hit` for every carry. The two callers differ only in *how*
|
||||
/// `hits` gets its slots (evidence-probed vs. trusted `index_batch`) — this
|
||||
/// is everything after that.
|
||||
fn resolve_layer_hits(
|
||||
mat: &Mat,
|
||||
hits: &[(usize, usize, u8)],
|
||||
n_genomes: usize,
|
||||
on_hit: &mut impl FnMut(usize, u8, usize),
|
||||
) {
|
||||
let n_cols = mat.n_cols().min(n_genomes);
|
||||
let slots: Vec<usize> = hits.iter().map(|&(slot, _, _)| slot).collect();
|
||||
let mut carries: Vec<Vec<bool>> = (0..n_cols).map(|_| Vec::new()).collect();
|
||||
mat.fill_sub_matrix_carries(&slots, &mut carries);
|
||||
for (g, col) in carries.iter().enumerate() {
|
||||
for (&(_, family_idx, base), &carries_it) in hits.iter().zip(col.iter()) {
|
||||
if carries_it {
|
||||
on_hit(family_idx, base, g);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -136,8 +136,11 @@ struct GeneratedBatch {
|
||||
masks: Vec<FamilyMask>,
|
||||
/// Flat `slots.len() * n_genomes` — `genome_mask[i * n_genomes + g]`.
|
||||
genome_mask: Vec<u8>,
|
||||
/// `outgoing[dest_partition]` = `(variant, family_idx_in_batch, base)`.
|
||||
outgoing: Vec<Vec<(CanonicalKmer, usize, u8)>>,
|
||||
/// `outgoing[dest_partition]` = `(variant, family_idx_in_batch, base,
|
||||
/// layer)` — `layer` is the annex-recorded destination layer
|
||||
/// (`FamilyMask::family_members`'s `layer_value`, `0` if absent),
|
||||
/// trustworthy only when `PartitionCache::fast_mode()` is true.
|
||||
outgoing: Vec<Vec<(CanonicalKmer, usize, u8, u8)>>,
|
||||
_permit: ThrottleGuard,
|
||||
}
|
||||
|
||||
@@ -210,27 +213,27 @@ pub(super) fn scan_layer_families(
|
||||
let mut masks = Vec::with_capacity(n);
|
||||
let mut bases = Vec::with_capacity(n);
|
||||
let mut genome_mask = vec![0u8; n * ctx.n_genomes];
|
||||
let mut outgoing: Vec<Vec<(CanonicalKmer, usize, u8)>> = (0..ctx.n_parts).map(|_| Vec::new()).collect();
|
||||
let mut outgoing: Vec<Vec<(CanonicalKmer, usize, u8, u8)>> = (0..ctx.n_parts).map(|_| Vec::new()).collect();
|
||||
|
||||
// Pass 1: cheap, no matrix access — own base and this
|
||||
// batch's cross-partition queries. Kmer and mask already in
|
||||
// hand from `iter_minorants_batch`, no second annex read,
|
||||
// no slot lookup needed for this pass.
|
||||
// no slot lookup needed for this pass. `family_members`
|
||||
// (not a hand-rolled `central_canonical_neighbors` +
|
||||
// `mask.has` loop) already filters to present members and
|
||||
// hands back each one's annex-recorded layer alongside.
|
||||
for (i, entry) in batch.entries.iter().enumerate() {
|
||||
let (kmer, mask) = (entry.kmer, entry.mask);
|
||||
masks.push(mask);
|
||||
let base = central_base(kmer, ctx.k);
|
||||
bases.push(base);
|
||||
for other in kmer.central_canonical_neighbors() {
|
||||
if other == kmer {
|
||||
for (member, layer) in mask.family_members(kmer, ctx.k) {
|
||||
if member == kmer {
|
||||
continue; // local — resolved below straight from `mat`, no lookup
|
||||
}
|
||||
let b = central_base(other, ctx.k);
|
||||
if !mask.has(b) {
|
||||
continue;
|
||||
}
|
||||
let dest = other.partition(ctx.n_parts);
|
||||
outgoing[dest].push((other, i, b));
|
||||
let b = central_base(member, ctx.k);
|
||||
let dest = member.partition(ctx.n_parts);
|
||||
outgoing[dest].push((member, i, b, layer.unwrap_or(0)));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -280,10 +283,16 @@ pub(super) fn scan_layer_families(
|
||||
// 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 fast_mode = cache.fast_mode();
|
||||
batch.outgoing.par_iter().enumerate().filter(|(_, q)| !q.is_empty()).for_each(|(dest, queries)| {
|
||||
cache.find_presence_batch(dest, queries, n_genomes, |i, base, g| {
|
||||
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 {
|
||||
|
||||
Reference in New Issue
Block a user