Push zpwxxpnpktps #67

Merged
coissac merged 46 commits from push-zpwxxpnpktps into main 2026-08-17 09:41:42 +00:00
4 changed files with 297 additions and 102 deletions
Showing only changes of commit 693c18bfa7 - Show all commits
+161 -57
View File
@@ -14,16 +14,20 @@ const MAGIC: [u8; 4] = *b"PBIV";
const HEADER_SIZE: usize = 16; const HEADER_SIZE: usize = 16;
#[inline] #[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] #[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 ────────────────────────────────────────────────────────── // ── PersistentBitVec ──────────────────────────────────────────────────────────
pub struct PersistentBitVec { pub struct PersistentBitVec {
mmap: Mmap, mmap: Mmap,
n: usize, n: usize,
path: PathBuf, path: PathBuf,
} }
@@ -31,18 +35,31 @@ impl PersistentBitVec {
pub fn open(path: &Path) -> io::Result<Self> { pub fn open(path: &Path) -> io::Result<Self> {
let mmap = unsafe { Mmap::map(&File::open(path)?)? }; let mmap = unsafe { Mmap::map(&File::open(path)?)? };
if mmap.len() < HEADER_SIZE { 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 { if &mmap[0..4] != &MAGIC {
return Err(io::Error::new(io::ErrorKind::InvalidData, "bad PBIV 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; 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 path(&self) -> &Path {
pub fn len(&self) -> usize { self.n } &self.path
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 { pub fn get(&self, slot: usize) -> bool {
(self.mmap[HEADER_SIZE + (slot >> 3)] >> (slot & 7)) & 1 != 0 (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]) { pub fn fill_batch(&self, slots: &[usize], out: &mut [bool]) {
assert_eq!(slots.len(), out.len()); assert_eq!(slots.len(), out.len());
let n = slots.len(); let n = slots.len();
if n == 0 { return; } if n == 0 {
return;
}
let mut perm: Vec<usize> = (0..n).collect(); let mut perm: Vec<usize> = (0..n).collect();
perm.sort_by_key(|&i| slots[i]); perm.sort_by_key(|&i| slots[i]);
let sorted: Vec<usize> = perm.iter().map(|&i| slots[i]).collect(); 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. // SAFETY: mmap is page-aligned, HEADER_SIZE=16 divisible by 8 → u64-aligned.
fn data_words(&self) -> &[u64] { 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; let ptr = self.mmap[HEADER_SIZE..].as_ptr() as *const u64;
unsafe { std::slice::from_raw_parts(ptr, nw) } unsafe { std::slice::from_raw_parts(ptr, nw) }
} }
@@ -95,10 +114,16 @@ impl PersistentBitVec {
BitSliceView::new(self.data_words(), self.n) 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_ones(&self) -> u64 {
pub fn count_zeros(&self) -> u64 { self.view().count_zeros() } self.view().count_ones()
}
pub fn count_zeros(&self) -> u64 {
self.view().count_zeros()
}
pub fn partial_jaccard_dist(&self, other: &PersistentBitVec) -> (u64, u64) { pub fn partial_jaccard_dist(&self, other: &PersistentBitVec) -> (u64, u64) {
self.view().partial_jaccard_dist(other.view()) self.view().partial_jaccard_dist(other.view())
@@ -111,22 +136,28 @@ impl PersistentBitVec {
} }
pub fn iter(&self) -> BitIter<'_> { 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 { impl<'a> IntoIterator for &'a PersistentBitVec {
type Item = bool; type Item = bool;
type IntoIter = BitIter<'a>; type IntoIter = BitIter<'a>;
fn into_iter(self) -> BitIter<'a> { self.iter() } fn into_iter(self) -> BitIter<'a> {
self.iter()
}
} }
// ── BitIter ─────────────────────────────────────────────────────────────────── // ── BitIter ───────────────────────────────────────────────────────────────────
pub struct BitIter<'a> { pub struct BitIter<'a> {
words: &'a [u64], words: &'a [u64],
slot: usize, slot: usize,
n: usize, n: usize,
} }
impl ExactSizeIterator for BitIter<'_> {} impl ExactSizeIterator for BitIter<'_> {}
@@ -134,7 +165,9 @@ impl ExactSizeIterator for BitIter<'_> {}
impl Iterator for BitIter<'_> { impl Iterator for BitIter<'_> {
type Item = bool; type Item = bool;
fn next(&mut self) -> Option<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; let v = (self.words[self.slot >> 6] >> (self.slot & 63)) & 1 != 0;
self.slot += 1; self.slot += 1;
Some(v) Some(v)
@@ -149,7 +182,7 @@ impl Iterator for BitIter<'_> {
pub struct PersistentBitVecBuilder { pub struct PersistentBitVecBuilder {
mmap: MmapMut, mmap: MmapMut,
n: usize, n: usize,
path: PathBuf, path: PathBuf,
} }
@@ -157,7 +190,10 @@ impl PersistentBitVecBuilder {
pub fn new(n: usize, path: &Path) -> io::Result<Self> { pub fn new(n: usize, path: &Path) -> io::Result<Self> {
let file_size = HEADER_SIZE + n_bytes_for_words(n); let file_size = HEADER_SIZE + n_bytes_for_words(n);
let mut file = OpenOptions::new() let mut file = OpenOptions::new()
.read(true).write(true).create(true).truncate(true) .read(true)
.write(true)
.create(true)
.truncate(true)
.open(path)?; .open(path)?;
file.write_all(&MAGIC)?; file.write_all(&MAGIC)?;
file.write_all(&[0u8; 4])?; file.write_all(&[0u8; 4])?;
@@ -165,20 +201,31 @@ impl PersistentBitVecBuilder {
file.seek(SeekFrom::Start(0))?; file.seek(SeekFrom::Start(0))?;
file.set_len(file_size as u64)?; file.set_len(file_size as u64)?;
let mmap = unsafe { MmapMut::map_mut(&file)? }; 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> { 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_size = HEADER_SIZE + n_bytes_for_words(n);
let file = OpenOptions::new() let file = OpenOptions::new()
.read(true).write(true).create(true).truncate(true) .read(true)
.write(true)
.create(true)
.truncate(true)
.open(path)?; .open(path)?;
file.set_len(file_size as u64)?; file.set_len(file_size as u64)?;
let mut mmap = unsafe { MmapMut::map_mut(&file)? }; let mut mmap = unsafe { MmapMut::map_mut(&file)? };
mmap[0..4].copy_from_slice(&MAGIC); mmap[0..4].copy_from_slice(&MAGIC);
mmap[8..16].copy_from_slice(&(n as u64).to_le_bytes()); mmap[8..16].copy_from_slice(&(n as u64).to_le_bytes());
mmap[HEADER_SIZE..HEADER_SIZE + bytes.len()].copy_from_slice(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`. /// 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 /// 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. /// 0xFF bytes in a single sequential pass, with no intermediate all-zeros state.
pub fn new_ones(n: usize, path: &Path) -> io::Result<Self> { 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 file_size = HEADER_SIZE + nw * 8;
let mut file = OpenOptions::new() let mut file = OpenOptions::new()
.read(true).write(true).create(true).truncate(true) .read(true)
.write(true)
.create(true)
.truncate(true)
.open(path)?; .open(path)?;
file.write_all(&MAGIC)?; file.write_all(&MAGIC)?;
file.write_all(&[0u8; 4])?; file.write_all(&[0u8; 4])?;
@@ -201,11 +251,15 @@ impl PersistentBitVecBuilder {
// Clear padding bits in the last word so trailing bits are always 0. // Clear padding bits in the last word so trailing bits are always 0.
let rem = n % 64; let rem = n % 64;
if rem != 0 { 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) }; let words = unsafe { std::slice::from_raw_parts_mut(ptr, nw) };
words[nw - 1] &= (1u64 << rem) - 1; 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> { 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 file = OpenOptions::new().read(true).write(true).open(path)?;
let mmap = unsafe { MmapMut::map_mut(&file)? }; let mmap = unsafe { MmapMut::map_mut(&file)? };
let n = source.len(); 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 n = source.len();
let file_size = HEADER_SIZE + n_bytes_for_words(n); let file_size = HEADER_SIZE + n_bytes_for_words(n);
let mut file = OpenOptions::new() let mut file = OpenOptions::new()
.read(true).write(true).create(true).truncate(true) .read(true)
.write(true)
.create(true)
.truncate(true)
.open(path)?; .open(path)?;
file.write_all(&MAGIC)?; file.write_all(&MAGIC)?;
file.write_all(&[0u8; 4])?; file.write_all(&[0u8; 4])?;
@@ -229,22 +294,32 @@ impl PersistentBitVecBuilder {
file.set_len(file_size as u64)?; file.set_len(file_size as u64)?;
let mut mmap = unsafe { MmapMut::map_mut(&file)? }; 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 ptr = mmap[HEADER_SIZE..].as_mut_ptr() as *mut u64;
let words = unsafe { std::slice::from_raw_parts_mut(ptr, nw) }; let words = unsafe { std::slice::from_raw_parts_mut(ptr, nw) };
for (slot, count) in source.iter().enumerate() { 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> { pub fn build_from_presence(source: &PersistentCompactIntVec, path: &Path) -> io::Result<Self> {
Self::build_from_counts(source, 1, path) Self::build_from_counts(source, 1, path)
} }
pub fn len(&self) -> usize { self.n } pub fn len(&self) -> usize {
pub fn is_empty(&self) -> bool { self.n == 0 } self.n
}
pub fn is_empty(&self) -> bool {
self.n == 0
}
pub fn get(&self, slot: usize) -> bool { pub fn get(&self, slot: usize) -> bool {
(self.mmap[HEADER_SIZE + (slot >> 3)] >> (slot & 7)) & 1 != 0 (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) { pub fn set(&mut self, slot: usize, value: bool) {
let bit = 1u64 << (slot & 63); let bit = 1u64 << (slot & 63);
if value { self.data_words_mut()[slot >> 6] |= bit; } if value {
else { self.data_words_mut()[slot >> 6] &= !bit; } self.data_words_mut()[slot >> 6] |= bit;
} else {
self.data_words_mut()[slot >> 6] &= !bit;
}
} }
fn data_words(&self) -> &[u64] { 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; let ptr = self.mmap[HEADER_SIZE..].as_ptr() as *const u64;
unsafe { std::slice::from_raw_parts(ptr, nw) } unsafe { std::slice::from_raw_parts(ptr, nw) }
} }
// SAFETY: same alignment argument as PersistentBitVec::data_words. // SAFETY: same alignment argument as PersistentBitVec::data_words.
fn data_words_mut(&mut self) -> &mut [u64] { 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; let ptr = self.mmap[HEADER_SIZE..].as_mut_ptr() as *mut u64;
unsafe { std::slice::from_raw_parts_mut(ptr, nw) } unsafe { std::slice::from_raw_parts_mut(ptr, nw) }
} }
@@ -273,7 +351,9 @@ impl PersistentBitVecBuilder {
BitSliceView::new(self.data_words(), self.n) 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<'_>) { pub fn copy_from(&mut self, src: BitSliceView<'_>) {
assert_eq!(self.n, src.len(), "BitSliceView length mismatch"); assert_eq!(self.n, src.len(), "BitSliceView length mismatch");
@@ -282,25 +362,35 @@ impl PersistentBitVecBuilder {
pub fn and(&mut self, other: BitSliceView<'_>) { pub fn and(&mut self, other: BitSliceView<'_>) {
assert_eq!(self.n, other.len(), "BitSliceView length mismatch"); 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<'_>) { pub fn or(&mut self, other: BitSliceView<'_>) {
assert_eq!(self.n, other.len(), "BitSliceView length mismatch"); 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<'_>) { pub fn xor(&mut self, other: BitSliceView<'_>) {
assert_eq!(self.n, other.len(), "BitSliceView length mismatch"); 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) { pub fn not(&mut self) {
let rem = self.n % 64; let rem = self.n % 64;
let words = self.data_words_mut(); 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 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 words = self.data_words_mut();
let nw = n_words(n); let nw = n_words(n);
for wi in 0..nw { for wi in 0..nw {
let base = wi * 64; let base = wi * 64;
let limit = (base + 64).min(n); let limit = (base + 64).min(n);
let mut mask = 0u64; let mut mask = 0u64;
for bit in 0..(limit - base) { for bit in 0..(limit - base) {
let b = primary[base + bit]; 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; words[wi] |= mask;
} }
for (slot, val) in col.overflow_entries() { 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 words = self.data_words_mut();
let nw = n_words(n); let nw = n_words(n);
for wi in 0..nw { for wi in 0..nw {
let base = wi * 64; let base = wi * 64;
let limit = (base + 64).min(n); let limit = (base + 64).min(n);
let mut mask = 0u64; let mut mask = 0u64;
for bit in 0..(limit - base) { for bit in 0..(limit - base) {
let b = primary[base + bit]; 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; words[wi] &= !mask;
} }
for (slot, val) in col.overflow_entries() { 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 words = self.data_words_mut();
let nw = n_words(n); let nw = n_words(n);
for wi in 0..nw { for wi in 0..nw {
let base = wi * 64; let base = wi * 64;
let limit = (base + 64).min(n); let limit = (base + 64).min(n);
let mut mask = 0u64; let mut mask = 0u64;
for bit in 0..(limit - base) { for bit in 0..(limit - base) {
let b = primary[base + bit]; 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; words[wi] ^= mask;
} }
for (slot, val) in col.overflow_entries() { 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() 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> { pub fn finish(self) -> io::Result<PersistentBitVec> {
let path = self.path.clone(); let path = self.path.clone();
+6 -14
View File
@@ -132,20 +132,12 @@ fn build_layer_sibling_annex(
let n = mphf.n(); let n = mphf.n();
// Whether this layer's family-member fields can carry a real layer // Whether this layer's family-member fields can carry a real layer
// number (fast path, exploited by a future scan-time reader) or must // number (exploited by `family_scan::scan_layer_families`'s fast path)
// stay presence-only (today's behaviour) — a single fact about the // or must stay presence-only — a single fact about the whole index,
// whole index, decided once here, never stored: `n_layers` is // decided once by `cache` (built from the same `PartitionMeta` this
// guaranteed identical across every partition (a merge adds one layer // function would otherwise re-derive) so the writer and every reader
// to all of them at once), and `FamilyMask`'s 3-bit field only has room // agree; see `PartitionCache::fast_mode`'s docs.
// for layers `0..=6`. let fast_mode = cache.fast_mode();
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
);
}
// ── The annex file itself is the reconciliation state, indexed by // ── The annex file itself is the reconciliation state, indexed by
// this layer's k-mer iteration order (the physical layout of // this layer's k-mer iteration order (the physical layout of
+107 -17
View File
@@ -117,18 +117,28 @@ pub(super) struct PartitionCache {
/// both [`obikindex::KmerIndex::build_sibling_annex`] and /// both [`obikindex::KmerIndex::build_sibling_annex`] and
/// [`obikindex::KmerIndex::sibling_annex_stats`]. /// [`obikindex::KmerIndex::sibling_annex_stats`].
mats: Vec<Vec<Mat>>, 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 { impl PartitionCache {
pub(super) fn build(partition: &KmerPartition, n_parts: usize, with_counts: bool) -> OKIResult<Self> { 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 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() .into_par_iter()
.map(|part| -> OKIResult<Vec<Mat>> { .map(|part| -> OKIResult<(Vec<Mat>, usize)> {
let index_dir = partition.part_dir(part).join(INDEX_SUBDIR); let index_dir = partition.part_dir(part).join(INDEX_SUBDIR);
if !index_dir.exists() { if !index_dir.exists() {
pb.inc(1); pb.inc(1);
return Ok(Vec::new()); return Ok((Vec::new(), 0));
} }
let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?; let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?;
let mut mats = Vec::with_capacity(meta.n_layers); let mut mats = Vec::with_capacity(meta.n_layers);
@@ -144,11 +154,27 @@ impl PartitionCache {
mats.push(mat); mats.push(mat);
} }
pb.inc(1); pb.inc(1);
Ok(mats) Ok((mats, meta.n_layers))
}) })
.collect::<OKIResult<Vec<_>>>()?; .collect::<OKIResult<Vec<_>>>()?;
pb.finish_and_clear(); 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 /// Existence lookup of `variant` in partition `dest_partition`: tries
@@ -186,10 +212,16 @@ impl PartitionCache {
/// `PersistentBitMatrix::get` dominated wall-clock time, mostly blocked /// `PersistentBitMatrix::get` dominated wall-clock time, mostly blocked
/// on page faults, even after every partition/batch locality fix above /// on page faults, even after every partition/batch locality fix above
/// this in the traversal. /// 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( pub(super) fn find_presence_batch(
&self, &self,
dest_partition: usize, dest_partition: usize,
queries: &[(CanonicalKmer, usize, u8)], queries: &[(CanonicalKmer, usize, u8, u8)],
n_genomes: usize, n_genomes: usize,
mut on_hit: impl FnMut(usize, u8, 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 // First hit wins, same semantics as the old per-query loop (a
// variant present in an earlier layer shadows later ones). // variant present in an earlier layer shadows later ones).
let mut by_layer: Vec<Vec<(usize, usize, u8)>> = vec![Vec::new(); mats.len()]; 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() { for (li, mat) in mats.iter().enumerate() {
if let Some(slot) = mat.find_slot(variant) { if let Some(slot) = mat.find_slot(variant) {
by_layer[li].push((slot, family_idx, base)); by_layer[li].push((slot, family_idx, base));
@@ -211,17 +243,75 @@ impl PartitionCache {
if hits.is_empty() { if hits.is_empty() {
continue; 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 mat = &mats[li];
let n_cols = mat.n_cols().min(n_genomes); let variants: Vec<CanonicalKmer> = entries.iter().map(|&(variant, _, _)| variant).collect();
let slots: Vec<usize> = hits.iter().map(|&(slot, _, _)| slot).collect(); let slots = mat.index_batch(&variants);
let mut carries: Vec<Vec<bool>> = (0..n_cols).map(|_| Vec::new()).collect(); let hits: Vec<(usize, usize, u8)> = slots
mat.fill_sub_matrix_carries(&slots, &mut carries); .into_iter()
for (g, col) in carries.iter().enumerate() { .zip(entries.iter())
for (&(_, family_idx, base), &carries_it) in hits.iter().zip(col.iter()) { .map(|(slot, &(_, family_idx, base))| (slot, family_idx, base))
if carries_it { .collect();
on_hit(family_idx, base, g); 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);
} }
} }
} }
+23 -14
View File
@@ -136,8 +136,11 @@ struct GeneratedBatch {
masks: Vec<FamilyMask>, masks: Vec<FamilyMask>,
/// Flat `slots.len() * n_genomes` — `genome_mask[i * n_genomes + g]`. /// Flat `slots.len() * n_genomes` — `genome_mask[i * n_genomes + g]`.
genome_mask: Vec<u8>, genome_mask: Vec<u8>,
/// `outgoing[dest_partition]` = `(variant, family_idx_in_batch, base)`. /// `outgoing[dest_partition]` = `(variant, family_idx_in_batch, base,
outgoing: Vec<Vec<(CanonicalKmer, usize, u8)>>, /// 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, _permit: ThrottleGuard,
} }
@@ -210,27 +213,27 @@ pub(super) fn scan_layer_families(
let mut masks = Vec::with_capacity(n); let mut masks = Vec::with_capacity(n);
let mut bases = Vec::with_capacity(n); let mut bases = Vec::with_capacity(n);
let mut genome_mask = vec![0u8; n * ctx.n_genomes]; 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 // Pass 1: cheap, no matrix access — own base and this
// batch's cross-partition queries. Kmer and mask already in // batch's cross-partition queries. Kmer and mask already in
// hand from `iter_minorants_batch`, no second annex read, // 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() { for (i, entry) in batch.entries.iter().enumerate() {
let (kmer, mask) = (entry.kmer, entry.mask); let (kmer, mask) = (entry.kmer, entry.mask);
masks.push(mask); masks.push(mask);
let base = central_base(kmer, ctx.k); let base = central_base(kmer, ctx.k);
bases.push(base); bases.push(base);
for other in kmer.central_canonical_neighbors() { for (member, layer) in mask.family_members(kmer, ctx.k) {
if other == kmer { if member == kmer {
continue; // local — resolved below straight from `mat`, no lookup continue; // local — resolved below straight from `mat`, no lookup
} }
let b = central_base(other, ctx.k); let b = central_base(member, ctx.k);
if !mask.has(b) { let dest = member.partition(ctx.n_parts);
continue; outgoing[dest].push((member, i, b, layer.unwrap_or(0)));
}
let dest = other.partition(ctx.n_parts);
outgoing[dest].push((other, i, b));
} }
} }
@@ -280,10 +283,16 @@ pub(super) fn scan_layer_families(
// 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();
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)| {
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); 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 {