OBIKmers/obikmer
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 27 Wiki Activity

Add cache-optimized batch retrieval and sub-matrix methods

Browse Source 
Introduces batch retrieval and sub-matrix extraction methods across vector, view, reader, and matrix types. These implementations optimize cache locality by sorting requested indices for sequential memory access before applying an inverse permutation to restore original order. Includes allocation-free variants that populate caller-provided buffers. Updates architecture documentation to define sibling annex persistence in iteration order and clarify pipeline separation.
This commit is contained in:
Eric Coissac
2026-08-16 14:01:35 +02:00
parent dae543fdfc
commit b1f54b7d2f
9 changed files with 375 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/obicompactvec/src/bitmatrix/persistent.rs
+51
View File
@@ -126,6 +126,57 @@ impl PersistentBitMatrix {
}
}
/// Extract a sub-matrix containing only the rows at `slots`.
///
/// Returns a column-first `Vec<Vec<bool>>`: the outer `Vec` has one entry per
/// column, and each inner `Vec` contains the values for the requested rows
/// in the same order as `slots`. Column access is sequential to maximize
/// cache efficiency on the underlying mmap.
pub fn sub_matrix(&self, slots: &[usize]) -> Vec<Vec<bool>> {
let n_cols = self.n_cols();
let mut out: Vec<Vec<bool>> = Vec::with_capacity(n_cols);
for c in 0..n_cols {
let mut col_buf = vec![false; slots.len()];
match self {
Self::Columnar(m) => m.col(c).view().fill_batch(slots, &mut col_buf),
Self::Packed(m) => m.col_slice(c).fill_batch(slots, &mut col_buf),
Self::Implicit { .. } => col_buf.iter_mut().for_each(|b| *b = true),
}
out.push(col_buf);
}
out
}
/// Like [`sub_matrix`](Self::sub_matrix), but fills caller-provided column
/// buffers to avoid allocating the outer `Vec`.
///
/// `out` must have length `self.n_cols()`. Each `out[c]` is cleared,
/// resized to `slots.len()`, and filled with the values for column `c`
/// in the same order as `slots`.
pub fn fill_sub_matrix(&self, slots: &[usize], out: &mut [Vec<bool>]) {
assert_eq!(out.len(), self.n_cols());
let n = slots.len();
if n == 0 {
for col in out.iter_mut() { col.clear(); }
return;
}
let mut perm: Vec<usize> = (0..n).collect();
perm.sort_by_key(|&i| slots[i]);
let sorted_slots: Vec<usize> = perm.iter().map(|&i| slots[i]).collect();
for (c, col) in out.iter_mut().enumerate() {
col.resize(n, false);
let mut tmp = vec![false; n];
match self {
Self::Columnar(m) => m.col(c).view().fill_batch_sorted(&sorted_slots, &mut tmp),
Self::Packed(m) => m.col_slice(c).fill_batch_sorted(&sorted_slots, &mut tmp),
Self::Implicit { .. } => tmp.iter_mut().for_each(|b| *b = true),
}
for (i, &orig_idx) in perm.iter().enumerate() {
col[orig_idx] = tmp[i];
}
}
}
pub fn count_ones(&self) -> Array1<u64> {
match self {
Self::Columnar(m) => m.count_ones(),
src/obicompactvec/src/bitvec.rs
+36
View File
@@ -48,6 +48,42 @@ impl PersistentBitVec {
(self.mmap[HEADER_SIZE + (slot >> 3)] >> (slot & 7)) & 1 != 0
}
/// Batch lookup: read bits at `slots` in an order that minimizes
/// cache misses.
///
/// The slots are sorted internally before reading so that accesses to
/// the underlying mmap are as sequential as possible, then the results
/// are reordered to match the input order.
pub fn get_batch(&self, slots: &[usize]) -> Vec<bool> {
let mut out = vec![false; slots.len()];
self.fill_batch(slots, &mut out);
out
}
/// Like [`get_batch`](Self::get_batch), but fills a caller-provided buffer.
pub fn fill_batch(&self, slots: &[usize], out: &mut [bool]) {
assert_eq!(slots.len(), out.len());
let n = slots.len();
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();
let mut tmp = vec![false; n];
self.fill_batch_sorted(&sorted, &mut tmp);
for (i, &orig_idx) in perm.iter().enumerate() {
out[orig_idx] = tmp[i];
}
}
/// Fill `out` assuming `sorted_slots` is already in ascending order.
/// Results are written in `sorted_slots` order (no reordering).
pub(crate) fn fill_batch_sorted(&self, sorted_slots: &[usize], out: &mut [bool]) {
assert_eq!(sorted_slots.len(), out.len());
for (i, &slot) in sorted_slots.iter().enumerate() {
out[i] = self.get(slot);
}
}
// SAFETY: mmap is page-aligned, HEADER_SIZE=16 divisible by 8 → u64-aligned.
fn data_words(&self) -> &[u64] {
let nw = n_words(self.n);
src/obicompactvec/src/intmatrix.rs
+44
View File
@@ -346,6 +346,50 @@ impl PersistentCompactIntMatrix {
pub fn fill_row(&self, slot: usize, buf: &mut [u32]) {
match self { Self::Columnar(m) => m.fill_row(slot, buf), Self::Packed(m) => m.fill_row(slot, buf) }
}
/// Extract a sub-matrix containing only the rows at `slots`.
///
/// Returns a column-first `Vec<Vec<u32>>`: the outer `Vec` has one entry per
/// column, and each inner `Vec` contains the values for the requested rows
/// in the same order as `slots`. Column access is sequential to maximize
/// cache efficiency on the underlying mmap.
pub fn sub_matrix(&self, slots: &[usize]) -> Vec<Vec<u32>> {
let n_cols = self.n_cols();
let mut out: Vec<Vec<u32>> = Vec::with_capacity(n_cols);
for c in 0..n_cols {
let mut col_buf = vec![0u32; slots.len()];
self.col_view(c).fill_batch(slots, &mut col_buf);
out.push(col_buf);
}
out
}
/// Like [`sub_matrix`](Self::sub_matrix), but fills caller-provided column
/// buffers to avoid allocating the outer `Vec`.
///
/// `out` must have length `self.n_cols()`. Each `out[c]` is cleared,
/// resized to `slots.len()`, and filled with the values for column `c`
/// in the same order as `slots`.
pub fn fill_sub_matrix(&self, slots: &[usize], out: &mut [Vec<u32>]) {
assert_eq!(out.len(), self.n_cols());
let n = slots.len();
if n == 0 {
for col in out.iter_mut() { col.clear(); }
return;
}
let mut perm: Vec<usize> = (0..n).collect();
perm.sort_by_key(|&i| slots[i]);
let sorted_slots: Vec<usize> = perm.iter().map(|&i| slots[i]).collect();
for (c, col) in out.iter_mut().enumerate() {
col.resize(n, 0);
let mut tmp = vec![0u32; n];
self.col_view(c).fill_batch_sorted(&sorted_slots, &mut tmp);
for (i, &orig_idx) in perm.iter().enumerate() {
col[orig_idx] = tmp[i];
}
}
}
pub fn sum(&self) -> Array1<u64> {
match self { Self::Columnar(m) => m.sum(), Self::Packed(m) => m.sum() }
}
src/obicompactvec/src/reader.rs
+36
View File
@@ -57,6 +57,42 @@ impl PersistentCompactIntVec {
}
}
/// Batch lookup: read values at `slots` in an order that minimizes
/// cache misses.
///
/// The slots are sorted internally before reading so that accesses to
/// the underlying mmap are as sequential as possible, then the results
/// are reordered to match the input order.
pub fn get_batch(&self, slots: &[usize]) -> Vec<u32> {
let mut out = vec![0u32; slots.len()];
self.fill_batch(slots, &mut out);
out
}
/// Like [`get_batch`](Self::get_batch), but fills a caller-provided buffer.
pub fn fill_batch(&self, slots: &[usize], out: &mut [u32]) {
assert_eq!(slots.len(), out.len());
let n = slots.len();
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();
let mut tmp = vec![0u32; n];
self.fill_batch_sorted(&sorted, &mut tmp);
for (i, &orig_idx) in perm.iter().enumerate() {
out[orig_idx] = tmp[i];
}
}
/// Fill `out` assuming `sorted_slots` is already in ascending order.
/// Results are written in `sorted_slots` order (no reordering).
pub(crate) fn fill_batch_sorted(&self, sorted_slots: &[usize], out: &mut [u32]) {
assert_eq!(sorted_slots.len(), out.len());
for (i, &slot) in sorted_slots.iter().enumerate() {
out[i] = self.get(slot);
}
}
fn overflow_get(&self, slot: usize) -> u32 {
let (pos_start, pos_end) = if self.step == 0 {
(0, self.n_overflow)
src/obicompactvec/src/views.rs
+66
View File
@@ -23,6 +23,38 @@ impl<'a> BitSliceView<'a> {
(self.words[slot >> 6] >> (slot & 63)) & 1 != 0
}
/// Batch lookup: read bits at `slots` with cache-friendly access pattern.
///
/// Slots are sorted internally before reading, then results are reordered
/// to match the input order.
pub fn get_batch(&self, slots: &[usize]) -> Vec<bool> {
let mut out = vec![false; slots.len()];
self.fill_batch(slots, &mut out);
out
}
/// Like [`get_batch`](Self::get_batch), but fills a caller-provided buffer.
pub fn fill_batch(&self, slots: &[usize], out: &mut [bool]) {
assert_eq!(slots.len(), out.len());
let n = slots.len();
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();
let mut tmp = vec![false; n];
self.fill_batch_sorted(&sorted, &mut tmp);
for (i, &orig_idx) in perm.iter().enumerate() {
out[orig_idx] = tmp[i];
}
}
/// Fill `out` assuming `sorted_slots` is already in ascending order.
/// Results are written in `sorted_slots` order (no reordering).
pub(crate) fn fill_batch_sorted(&self, sorted_slots: &[usize], out: &mut [bool]) {
assert_eq!(sorted_slots.len(), out.len());
for (i, &slot) in sorted_slots.iter().enumerate() {
out[i] = self.get(slot);
}
}
pub fn count_ones(&self) -> u64 {
self.words.iter().map(|w| w.count_ones() as u64).sum()
}
@@ -123,6 +155,40 @@ impl<'a> IntSliceView<'a> {
panic!("slot {slot} marked overflow but not found")
}
/// Batch lookup: read values at `slots` with cache-friendly access pattern.
///
/// Slots are sorted internally before reading, then results are reordered
/// to match the input order.
pub fn get_batch(&self, slots: &[usize]) -> Vec<u32> {
let mut out = vec![0u32; slots.len()];
self.fill_batch(slots, &mut out);
out
}
/// Like [`get_batch`](Self::get_batch), but fills a caller-provided buffer.
pub fn fill_batch(&self, slots: &[usize], out: &mut [u32]) {
assert_eq!(slots.len(), out.len());
let n = slots.len();
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();
let mut tmp = vec![0u32; n];
self.fill_batch_sorted(&sorted, &mut tmp);
for (i, &orig_idx) in perm.iter().enumerate() {
out[orig_idx] = tmp[i];
}
}
/// Fill `out` assuming `sorted_slots` is already in ascending order.
/// Results are written in `sorted_slots` order (no reordering).
pub(crate) fn fill_batch_sorted(&self, sorted_slots: &[usize], out: &mut [u32]) {
assert_eq!(sorted_slots.len(), out.len());
for (i, &slot) in sorted_slots.iter().enumerate() {
out[i] = self.get(slot);
}
}
/// Sequential merge scan: yields all n values in slot order.
pub fn iter(&self) -> IntSliceViewIter<'a> {
IntSliceViewIter {
src/obilayeredmap/src/layer.rs
+38
View File
@@ -186,6 +186,25 @@ impl Layer<PersistentCompactIntMatrix> {
PersistentCompactIntMatrix::append_column(&layer_dir.join(COUNTS_DIR), value_of)
.map_err(OLMError::Io)
}
/// Extract a sub-matrix of counts for the rows at `slots`.
///
/// Returns a column-first `Vec<Vec<u32>>`: one inner `Vec` per genome
/// column, containing the counts for the requested slots in order.
/// Column access is sequential for cache efficiency.
pub fn sub_matrix(&self, slots: &[usize]) -> Vec<Vec<u32>> {
self.data.sub_matrix(slots)
}
/// Like [`sub_matrix`](Self::sub_matrix), but fills caller-provided column
/// buffers to avoid allocating the outer `Vec`.
///
/// `out` must have length equal to the number of genome columns. Each
/// `out[c]` is cleared, resized to `slots.len()`, and filled with the
/// counts for column `c` in the same order as `slots`.
pub fn fill_sub_matrix(&self, slots: &[usize], out: &mut [Vec<u32>]) {
self.data.fill_sub_matrix(slots, out)
}
}
// ── Mode 3 — presence/absence matrix ─────────────────────────────────────────
@@ -199,6 +218,25 @@ impl Layer<PersistentBitMatrix> {
.map_err(OLMError::Io)
}
/// Extract a sub-matrix of presence/absence for the rows at `slots`.
///
/// Returns a column-first `Vec<Vec<bool>>`: one inner `Vec` per genome
/// column, containing the presence bits for the requested slots in order.
/// Column access is sequential for cache efficiency.
pub fn sub_matrix(&self, slots: &[usize]) -> Vec<Vec<bool>> {
self.data.sub_matrix(slots)
}
/// Like [`sub_matrix`](Self::sub_matrix), but fills caller-provided column
/// buffers to avoid allocating the outer `Vec`.
///
/// `out` must have length equal to the number of genome columns. Each
/// `out[c]` is cleared, resized to `slots.len()`, and filled with the
/// presence bits for column `c` in the same order as `slots`.
pub fn fill_sub_matrix(&self, slots: &[usize], out: &mut [Vec<bool>]) {
self.data.fill_sub_matrix(slots, out)
}
pub fn build_presence(
out_dir: &Path,
block_bits: u8,