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

refactor: update CI toolchain setup and optimize parallel indexing
CI / build (push) Successful in 4m56s
Details
CI / build (pull_request) Successful in 4m11s
Details

Browse Source 
Update CI workflows to explicitly install the Rust toolchain via rustup and configure musl targets for deterministic static builds in Docker containers. Bump obikmer dependency to 0.1.3. Refactor obicompactvec to reduce peak memory usage by computing column sizes from filesystem metadata, add atomic writes, and implement cleanup guards. Replace parallel iteration patterns in obikindex with a structured PartitionRunner pipeline for simplified error handling and progress tracking.
This commit is contained in:
Eric Coissac
2026-06-22 10:46:24 +02:00
parent 9578f991f4 a522c0907e
commit 4e625afaba
83 changed files with 5959 additions and 1010 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/Cargo.lock
Generated
+6 -1
View File
@@ -1704,7 +1704,7 @@ dependencies = [
[[package]]
name = "obikmer"
version = "0.1.0"
version = "0.1.3"
dependencies = [
"clap",
"csv",
@@ -1722,6 +1722,7 @@ dependencies = [
"obiskbuilder",
"obiskio",
"obisys",
"obitaxonomy",
"pprof",
"rayon",
"serde_json",
@@ -1853,6 +1854,10 @@ dependencies = [
"tracing",
]
[[package]]
name = "obitaxonomy"
version = "0.1.0"
[[package]]
name = "object"
version = "0.37.3"
src/Cargo.toml
+1 -1
View File
@@ -1,5 +1,5 @@
[workspace]
resolver = "3"
members = ["obikseq", "obiread", "obiskbuilder", "obifastwrite", "obikmer","obikrope","obipipeline", "obikpartitionner","obiskio","obidebruinj","obilayeredmap", "obicompactvec", "obisys", "obikindex"]
members = ["obikseq", "obiread", "obiskbuilder", "obifastwrite", "obikmer","obikrope","obipipeline", "obikpartitionner","obiskio","obidebruinj","obilayeredmap", "obicompactvec", "obisys", "obikindex", "obitaxonomy"]
[profile.release]
debug = 1
src/obicompactvec/Cargo.toml
+1 -1
View File
@@ -7,6 +7,6 @@ edition = "2024"
memmap2 = "0.9"
ndarray = "0.16"
rayon = "1"
tempfile = "3"
[dev-dependencies]
tempfile = "3"
src/obicompactvec/src/bitmatrix.rs
+154 -85
View File
@@ -7,8 +7,12 @@ use ndarray::{Array1, Array2};
use rayon::prelude::*;
use crate::bitvec::{PersistentBitVec, PersistentBitVecBuilder};
use crate::colgroup::{ColGroup, MatrixGroupOps};
use crate::layer_meta::LayerMeta;
use crate::meta::MatrixMeta;
use crate::tempbitvec::{TempBitVec, TempBitVecBuilder};
use crate::tempintvec::{TempCompactIntVec, TempCompactIntVecBuilder};
use crate::views::BitSliceView;
fn col_path(dir: &Path, col: usize) -> PathBuf {
dir.join(format!("col_{col:06}.pbiv"))
@@ -54,34 +58,11 @@ impl ColumnarBitMatrix {
}
pub(crate) fn partial_jaccard_dist_matrix(&self) -> (Array2<u64>, Array2<u64>) {
let n = self.n_cols();
let results: Vec<(usize, usize, u64, u64)> = upper_pairs(n)
.into_par_iter()
.map(|(i, j)| {
let (inter, union) = self.col(i).partial_jaccard_dist(self.col(j));
(i, j, inter, union)
})
.collect();
let mut inter_m = Array2::zeros((n, n));
let mut union_m = Array2::zeros((n, n));
for (i, j, inter, union) in results {
inter_m[[i, j]] = inter; inter_m[[j, i]] = inter;
union_m[[i, j]] = union; union_m[[j, i]] = union;
}
(inter_m, union_m)
pairwise2_matrix(self.n_cols(), |i, j| self.col(i).partial_jaccard_dist(self.col(j)))
}
pub(crate) fn partial_hamming_dist_matrix(&self) -> Array2<u64> {
self.pairwise_u64(|i, j| self.col(i).hamming_dist(self.col(j)))
}
fn pairwise_u64(&self, f: impl Fn(usize, usize) -> u64 + Sync) -> Array2<u64> {
let n = self.n_cols();
let results: Vec<(usize, usize, u64)> = upper_pairs(n)
.into_par_iter()
.map(|(i, j)| (i, j, f(i, j)))
.collect();
fill_symmetric(n, results.into_iter().map(|(i, j, v)| (i, j, v, v)))
pairwise_matrix(self.n_cols(), |i, j| self.col(i).hamming_dist(self.col(j)))
}
pub(crate) fn append_column(dir: &Path, value_of: impl Fn(usize) -> bool) -> io::Result<()> {
@@ -147,84 +128,46 @@ impl PackedBitMatrix {
}).collect()
}
#[inline]
fn col_bytes(&self, c: usize) -> &[u8] {
let start = self.data_offsets[c];
let len = (self.n_rows + 7) / 8;
&self.mmap[start..start + len]
&self.mmap[start..start + self.n_rows.div_ceil(8)]
}
fn count_ones_col(&self, c: usize) -> u64 {
let bytes = self.col_bytes(c);
let full = self.n_rows / 8;
let rem = self.n_rows % 8;
let mut n: u64 = bytes[..full].iter().map(|b| b.count_ones() as u64).sum();
if rem > 0 { n += (bytes[full] & ((1u8 << rem) - 1)).count_ones() as u64; }
n
fn col_words(&self, c: usize) -> &[u64] {
let nw = self.n_rows.div_ceil(64);
// SAFETY: data_offsets[c] is always 8-byte aligned.
// PBMX header = 24 + n_cols×8 (multiple of 8); each PBIV blob =
// 16 + nwords×8 (multiple of 8); mmap base is page-aligned.
let ptr = self.mmap[self.data_offsets[c]..].as_ptr() as *const u64;
unsafe { std::slice::from_raw_parts(ptr, nw) }
}
fn pair_op(&self, i: usize, j: usize, and_or: bool) -> u64 {
let ai = self.col_bytes(i);
let aj = self.col_bytes(j);
let full = self.n_rows / 8;
let rem = self.n_rows % 8;
let mut n: u64 = ai[..full].iter().zip(aj[..full].iter())
.map(|(a, b)| if and_or { a & b } else { a ^ b }.count_ones() as u64)
.sum();
if rem > 0 {
let mask = (1u8 << rem) - 1;
let last = if and_or { ai[full] & aj[full] } else { ai[full] ^ aj[full] };
n += (last & mask).count_ones() as u64;
}
n
pub(crate) fn col_slice(&self, c: usize) -> BitSliceView<'_> {
BitSliceView::new(self.col_words(c), self.n_rows)
}
fn partial_jaccard_col(&self, i: usize, j: usize) -> (u64, u64) {
let ai = self.col_bytes(i);
let aj = self.col_bytes(j);
let full = self.n_rows / 8;
let rem = self.n_rows % 8;
let (mut inter, mut union) = ai[..full].iter().zip(aj[..full].iter())
.fold((0u64, 0u64), |(inter, union), (a, b)| {
(inter + (a & b).count_ones() as u64,
union + (a | b).count_ones() as u64)
});
if rem > 0 {
let mask = (1u8 << rem) - 1;
inter += ((ai[full] & aj[full]) & mask).count_ones() as u64;
union += ((ai[full] | aj[full]) & mask).count_ones() as u64;
}
(inter, union)
pub(crate) fn col_persist(&self, c: usize, path: &Path) -> io::Result<PersistentBitVecBuilder> {
PersistentBitVecBuilder::from_raw_bytes(self.col_bytes(c), self.n_rows, path)
}
pub(crate) fn count_ones(&self) -> Array1<u64> {
Array1::from_vec(
(0..self.n_cols).into_par_iter().map(|c| self.count_ones_col(c)).collect()
(0..self.n_cols).into_par_iter()
.map(|c| self.col_slice(c).count_ones())
.collect()
)
}
pub(crate) fn partial_jaccard_dist_matrix(&self) -> (Array2<u64>, Array2<u64>) {
let n = self.n_cols;
let results: Vec<(usize, usize, u64, u64)> = upper_pairs(n)
.into_par_iter()
.map(|(i, j)| { let (inter, union) = self.partial_jaccard_col(i, j); (i, j, inter, union) })
.collect();
let mut inter_m = Array2::zeros((n, n));
let mut union_m = Array2::zeros((n, n));
for (i, j, inter, union) in results {
inter_m[[i, j]] = inter; inter_m[[j, i]] = inter;
union_m[[i, j]] = union; union_m[[j, i]] = union;
}
(inter_m, union_m)
pairwise2_matrix(self.n_cols, |i, j| {
self.col_slice(i).partial_jaccard_dist(self.col_slice(j))
})
}
pub(crate) fn partial_hamming_dist_matrix(&self) -> Array2<u64> {
let n = self.n_cols;
let results: Vec<(usize, usize, u64)> = upper_pairs(n)
.into_par_iter()
.map(|(i, j)| (i, j, self.pair_op(i, j, false)))
.collect();
fill_symmetric(n, results.into_iter().map(|(i, j, v)| (i, j, v, v)))
pairwise_matrix(self.n_cols, |i, j| {
self.col_slice(i).hamming_dist(self.col_slice(j))
})
}
}
@@ -343,6 +286,24 @@ impl PersistentBitMatrix {
}
}
pub fn col_view(&self, c: usize) -> BitSliceView<'_> {
match self {
Self::Columnar(m) => m.col(c).view(),
Self::Packed(m) => m.col_slice(c),
Self::Implicit { .. } => panic!("col_view() not available on Implicit PersistentBitMatrix"),
}
}
pub fn col_persist(&self, c: usize, path: &Path) -> io::Result<PersistentBitVecBuilder> {
match self {
Self::Columnar(m) => PersistentBitVecBuilder::build_from(m.col(c), path),
Self::Packed(m) => m.col_persist(c, path),
Self::Implicit { n_rows, .. } => {
PersistentBitVecBuilder::new_ones(*n_rows, path)
}
}
}
pub fn row(&self, slot: usize) -> Box<[bool]> {
match self {
Self::Columnar(m) => m.row(slot),
@@ -439,12 +400,93 @@ impl PersistentBitMatrixBuilder {
PersistentBitVecBuilder::new(self.n, &path)
}
pub fn add_col_ones(&mut self) -> io::Result<PersistentBitVecBuilder> {
let path = col_path(&self.dir, self.n_cols);
self.n_cols += 1;
PersistentBitVecBuilder::new_ones(self.n, &path)
}
pub fn add_col_from(&mut self, src: &TempBitVec) -> io::Result<()> {
src.make_persistent(&col_path(&self.dir, self.n_cols))?;
self.n_cols += 1;
Ok(())
}
pub fn add_col_from_int(&mut self, src: &TempCompactIntVec) -> io::Result<()> {
let path = col_path(&self.dir, self.n_cols);
self.n_cols += 1;
let mut b = PersistentBitVecBuilder::new(self.n, &path)?;
b.or_where(src.view(), |v| v > 0);
b.close()
}
pub fn close(self) -> io::Result<()> {
MatrixMeta { n: self.n, n_cols: self.n_cols }.save(&self.dir)
}
}
// ── Helpers ───────────────────────────────────────────────────────────────────
// ── MatrixGroupOps ────────────────────────────────────────────────────────────
impl MatrixGroupOps for PersistentBitMatrix {
fn partial_group_presence_count(&self, g: &ColGroup, _threshold: u32) -> io::Result<TempCompactIntVec> {
// Bit matrices store 0/1 — threshold is structurally always 1.
let n = self.n();
if g.indices.len() < 255 {
let mut builder = TempCompactIntVecBuilder::new(n)?;
for &c in &g.indices {
builder.inc_present_fast(self.col_view(c));
}
builder.freeze()
} else {
let mut result = TempCompactIntVecBuilder::new(n)?;
for chunk in g.indices.chunks(254) {
let mut chunk_b = TempCompactIntVecBuilder::new(n)?;
for &c in chunk {
chunk_b.inc_present_fast(self.col_view(c));
}
let frozen = chunk_b.freeze()?;
result.add(frozen.view());
}
result.freeze()
}
}
fn partial_group_sum(&self, g: &ColGroup) -> io::Result<TempCompactIntVec> {
// For bit matrices, sum = count of 1-bits — identical to presence_count.
self.partial_group_presence_count(g, 1)
}
fn partial_group_any(&self, g: &ColGroup, _threshold: u32) -> io::Result<TempBitVec> {
let n = self.n();
let mut result = TempBitVecBuilder::new(n)?;
for &c in &g.indices {
result.or(self.col_view(c));
}
result.freeze()
}
fn partial_group_min(&self, g: &ColGroup) -> io::Result<TempCompactIntVec> {
// min of 0/1 values = AND: 1 only if ALL columns are 1
let n = self.n();
let mut result = TempCompactIntVecBuilder::new(n)?;
if let Some((&first, rest)) = g.indices.split_first() {
result.inc_present_fast(self.col_view(first));
for &c in rest { result.mask_with(self.col_view(c)); }
}
result.freeze()
}
fn partial_group_max(&self, g: &ColGroup) -> io::Result<TempCompactIntVec> {
// max of 0/1 values = OR: 1 if any column is 1
let any = self.partial_group_any(g, 1)?;
let n = any.len();
let mut result = TempCompactIntVecBuilder::new(n)?;
result.inc_present(any.view());
result.freeze()
}
}
// ── Shared matrix helpers (also used by intmatrix.rs) ─────────────────────────
fn upper_pairs(n: usize) -> Vec<(usize, usize)> {
(0..n).flat_map(|i| (i + 1..n).map(move |j| (i, j))).collect()
@@ -456,3 +498,30 @@ where T: Clone + Default {
for (i, j, vij, vji) in vals { m[[i, j]] = vij; m[[j, i]] = vji; }
m
}
/// Compute a symmetric `n×n` matrix in parallel by evaluating `f(i,j)` for
/// all upper-triangle pairs. `T: Copy` avoids the `.clone()` needed for the
/// lower-triangle mirror.
pub(crate) fn pairwise_matrix<T>(n: usize, f: impl Fn(usize, usize) -> T + Sync) -> Array2<T>
where T: Copy + Default + Send {
let results: Vec<(usize, usize, T)> = upper_pairs(n)
.into_par_iter().map(|(i, j)| (i, j, f(i, j))).collect();
fill_symmetric(n, results.into_iter().map(|(i, j, v)| (i, j, v, v)))
}
/// Same as `pairwise_matrix` but `f` returns two values that fill two
/// symmetric matrices simultaneously (e.g. intersection + union for Jaccard).
pub(crate) fn pairwise2_matrix<T>(n: usize, f: impl Fn(usize, usize) -> (T, T) + Sync) -> (Array2<T>, Array2<T>)
where T: Copy + Default + Send {
let results: Vec<(usize, usize, T, T)> = upper_pairs(n)
.into_par_iter()
.map(|(i, j)| { let (a, b) = f(i, j); (i, j, a, b) })
.collect();
let mut m0 = Array2::from_elem((n, n), T::default());
let mut m1 = Array2::from_elem((n, n), T::default());
for (i, j, a, b) in results {
m0[[i, j]] = a; m0[[j, i]] = a;
m1[[i, j]] = b; m1[[j, i]] = b;
}
(m0, m1)
}
src/obicompactvec/src/bitvec.rs
+221 -179
View File
@@ -5,29 +5,25 @@ use std::path::{Path, PathBuf};
use memmap2::{Mmap, MmapMut};
use crate::reader::PersistentCompactIntVec;
use crate::views::{BitSliceIter, BitSliceView, IntSliceView};
const MAGIC: [u8; 4] = *b"PBIV";
// Header: magic(4) + _pad(4) + n(8) = 16 bytes.
// Data starts at offset 16, which is divisible by 8 → u64-aligned
// (mmap base is page-aligned, 16 % 8 == 0).
// Data starts at offset 16, u64-aligned (mmap base is page-aligned, 16 % 8 == 0).
const HEADER_SIZE: usize = 16;
#[inline]
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 }
// ── Reader ────────────────────────────────────────────────────────────────────
// ── PersistentBitVec ──────────────────────────────────────────────────────────
pub struct PersistentBitVec {
mmap: Mmap,
n: usize,
n: usize,
path: PathBuf,
}
@@ -35,157 +31,145 @@ 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
}
// Used by iter() and get(): exact byte window, no padding.
fn data_bytes(&self) -> &[u8] {
&self.mmap[HEADER_SIZE..HEADER_SIZE + self.n.div_ceil(8)]
}
// Bulk word view. SAFETY: mmap is page-aligned, HEADER_SIZE=16 is divisible by 8,
// so &mmap[HEADER_SIZE] is u64-aligned. Slice length is n_words * 8 bytes.
// 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) }
}
pub fn count_ones(&self) -> u64 {
// Padding bits in the last word are 0, so no masking needed.
self.data_words()
.iter()
.map(|w| w.count_ones() as u64)
.sum()
pub fn view(&self) -> BitSliceView<'_> {
BitSliceView::new(self.data_words(), self.n)
}
pub fn count_zeros(&self) -> u64 {
self.n as u64 - self.count_ones()
}
pub fn words(&self) -> &[u64] { self.data_words() }
pub fn jaccard_dist(&self, other: &PersistentBitVec) -> f64 {
let (inter, union) = self.partial_jaccard_dist(other);
if union == 0 {
return 0.0;
}
1.0 - inter as f64 / union as f64
}
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) {
assert_eq!(self.n, other.n, "length mismatch");
self.data_words()
.iter()
.zip(other.data_words())
.fold((0u64, 0u64), |(i, u), (&a, &b)| {
(
i + (a & b).count_ones() as u64,
u + (a | b).count_ones() as u64,
)
})
self.view().partial_jaccard_dist(other.view())
}
pub fn jaccard_dist(&self, other: &PersistentBitVec) -> f64 {
self.view().jaccard_dist(other.view())
}
pub fn hamming_dist(&self, other: &PersistentBitVec) -> u64 {
assert_eq!(self.n, other.n, "length mismatch");
self.data_words()
.iter()
.zip(other.data_words())
.map(|(&a, &b)| (a ^ b).count_ones() as u64)
.sum()
self.view().hamming_dist(other.view())
}
pub fn iter(&self) -> BitIter<'_> {
BitIter {
bytes: self.data_bytes(),
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> {
bytes: &'a [u8],
slot: usize,
n: usize,
words: &'a [u64],
slot: usize,
n: usize,
}
impl ExactSizeIterator for BitIter<'_> {}
impl Iterator for BitIter<'_> {
type Item = bool;
fn next(&mut self) -> Option<bool> {
if self.slot >= self.n {
return None;
}
let v = (self.bytes[self.slot >> 3] >> (self.slot & 7)) & 1 != 0;
if self.slot >= self.n { return None; }
let v = (self.words[self.slot >> 6] >> (self.slot & 63)) & 1 != 0;
self.slot += 1;
Some(v)
}
fn size_hint(&self) -> (usize, Option<usize>) {
let rem = self.n - self.slot;
(rem, Some(rem))
}
}
// ── Builder ───────────────────────────────────────────────────────────────────
// ── PersistentBitVecBuilder ───────────────────────────────────────────────────
pub struct PersistentBitVecBuilder {
mmap: MmapMut,
n: usize,
n: usize,
path: PathBuf,
}
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])?; // padding
file.write_all(&[0u8; 4])?;
file.write_all(&(n as u64).to_le_bytes())?;
file.seek(SeekFrom::Start(0))?;
file.set_len(file_size as u64)?;
let mmap = unsafe { MmapMut::map_mut(&file)? };
Ok(Self { mmap, n })
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)
.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() })
}
/// Create an all-ones bit vector of length `n` at `path`.
///
/// 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 file_size = HEADER_SIZE + nw * 8;
let mut file = OpenOptions::new()
.read(true).write(true).create(true).truncate(true)
.open(path)?;
file.write_all(&MAGIC)?;
file.write_all(&[0u8; 4])?;
file.write_all(&(n as u64).to_le_bytes())?;
file.write_all(&vec![0xFFu8; nw * 8])?;
file.seek(SeekFrom::Start(0))?;
file.set_len(file_size as u64)?;
let mut mmap = unsafe { MmapMut::map_mut(&file)? };
// 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 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() })
}
pub fn build_from(source: &PersistentBitVec, path: &Path) -> io::Result<Self> {
@@ -193,86 +177,14 @@ 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 })
Ok(Self { mmap, n, path: path.to_path_buf() })
}
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
}
pub fn set(&mut self, slot: usize, value: bool) {
let byte = HEADER_SIZE + (slot >> 3);
let bit = 1u8 << (slot & 7);
if value {
self.mmap[byte] |= bit;
} else {
self.mmap[byte] &= !bit;
}
}
// SAFETY: same alignment argument as PersistentBitVec::data_words.
fn data_words_mut(&mut self) -> &mut [u64] {
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) }
}
pub fn and(&mut self, other: &PersistentBitVec) {
assert_eq!(self.n, other.n, "length mismatch");
for (sw, &ow) in self.data_words_mut().iter_mut().zip(other.data_words()) {
*sw &= ow;
}
}
pub fn or(&mut self, other: &PersistentBitVec) {
assert_eq!(self.n, other.n, "length mismatch");
for (sw, &ow) in self.data_words_mut().iter_mut().zip(other.data_words()) {
*sw |= ow;
}
}
pub fn xor(&mut self, other: &PersistentBitVec) {
assert_eq!(self.n, other.n, "length mismatch");
for (sw, &ow) in self.data_words_mut().iter_mut().zip(other.data_words()) {
*sw ^= ow;
}
}
pub fn not(&mut self) {
let rem = self.n % 64;
let words = self.data_words_mut();
for w in words.iter_mut() {
*w ^= u64::MAX;
}
// Zero padding bits in the last word so count_ones / jaccard remain correct.
if rem != 0 {
if let Some(last) = words.last_mut() {
*last &= (1u64 << rem) - 1;
}
}
}
/// Convert a count vector to a bit vector: bit set iff count >= threshold.
/// Fills u64 words directly from the count iterator — O(n), no bit-level set() overhead.
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])?;
@@ -280,27 +192,157 @@ impl PersistentBitVecBuilder {
file.seek(SeekFrom::Start(0))?;
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 })
Ok(Self { mmap, n, path: path.to_path_buf() })
}
/// Convert a count vector to a presence/absence bit vector (threshold = 1).
pub fn build_from_presence(source: &PersistentCompactIntVec, path: &Path) -> io::Result<Self> {
Self::build_from_counts(source, 1, path)
}
pub fn close(self) -> io::Result<()> {
self.mmap.flush()
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
}
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; }
}
fn data_words(&self) -> &[u64] {
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 ptr = self.mmap[HEADER_SIZE..].as_mut_ptr() as *mut u64;
unsafe { std::slice::from_raw_parts_mut(ptr, nw) }
}
pub fn view(&self) -> BitSliceView<'_> {
BitSliceView::new(self.data_words(), self.n)
}
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");
self.data_words_mut().copy_from_slice(src.words());
}
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; }
}
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; }
}
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; }
}
pub fn not(&mut self) {
let rem = self.n % 64;
let words = self.data_words_mut();
for w in words.iter_mut() { *w ^= u64::MAX; }
if rem != 0 {
if let Some(last) = words.last_mut() { *last &= (1u64 << rem) - 1; }
}
}
/// OR in bits at slots where `pred(col[slot])` is true.
pub fn or_where(&mut self, col: IntSliceView<'_>, pred: impl Fn(u32) -> bool) {
assert_eq!(self.n, col.len(), "IntSliceView length mismatch");
let n = self.n;
let primary = col.primary_bytes();
let words = self.data_words_mut();
let nw = n_words(n);
for wi in 0..nw {
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; }
}
words[wi] |= mask;
}
for (slot, val) in col.overflow_entries() {
if pred(val) { words[slot >> 6] |= 1u64 << (slot & 63); }
}
}
/// Clear bits at slots where `pred(col[slot])` is false.
pub fn and_where(&mut self, col: IntSliceView<'_>, pred: impl Fn(u32) -> bool) {
assert_eq!(self.n, col.len(), "IntSliceView length mismatch");
let n = self.n;
let primary = col.primary_bytes();
let words = self.data_words_mut();
let nw = n_words(n);
for wi in 0..nw {
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; }
}
words[wi] &= !mask;
}
for (slot, val) in col.overflow_entries() {
if !pred(val) { words[slot >> 6] &= !(1u64 << (slot & 63)); }
}
}
/// Toggle bits at slots where `pred(col[slot])` is true.
pub fn xor_where(&mut self, col: IntSliceView<'_>, pred: impl Fn(u32) -> bool) {
assert_eq!(self.n, col.len(), "IntSliceView length mismatch");
let n = self.n;
let primary = col.primary_bytes();
let words = self.data_words_mut();
let nw = n_words(n);
for wi in 0..nw {
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; }
}
words[wi] ^= mask;
}
for (slot, val) in col.overflow_entries() {
if pred(val) { words[slot >> 6] ^= 1u64 << (slot & 63); }
}
}
pub fn iter(&self) -> BitSliceIter<'_> {
self.view().iter()
}
pub fn close(self) -> io::Result<()> { self.mmap.flush() }
pub fn finish(self) -> io::Result<PersistentBitVec> {
let path = self.path.clone();
self.close()?;
PersistentBitVec::open(&path)
}
}
src/obicompactvec/src/builder.rs
+195 -69
View File
@@ -5,71 +5,57 @@ use std::path::{Path, PathBuf};
use memmap2::MmapMut;
use crate::format::{HEADER_SIZE, OVERFLOW_ENTRY_SIZE, finalize_pciv};
use crate::format::{byte_count_nonzero, byte_sum, HEADER_SIZE, finalize_pciv, parse_overflow_entry};
use crate::reader::PersistentCompactIntVec;
use crate::views::{BitSliceView, IntSliceView};
pub struct PersistentCompactIntVecBuilder {
path: PathBuf,
mmap: MmapMut,
n: usize,
path: PathBuf,
mmap: MmapMut,
n: usize,
overflow: HashMap<usize, u32>,
}
impl PersistentCompactIntVecBuilder {
/// Create a new, zero-filled PCIV at `path`. Primary is mmapped immediately.
pub fn new(n: usize, path: &Path) -> io::Result<Self> {
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((HEADER_SIZE + n) as u64)?;
let mmap = unsafe { MmapMut::map_mut(&file)? };
Ok(Self {
path: path.to_path_buf(),
mmap,
n,
overflow: HashMap::new(),
})
Ok(Self { path: path.to_path_buf(), mmap, n, overflow: HashMap::new() })
}
pub fn from_raw_primary(primary: &[u8], overflow: HashMap<usize, u32>, path: &Path) -> io::Result<Self> {
let n = primary.len();
let file = OpenOptions::new()
.read(true).write(true).create(true).truncate(true)
.open(path)?;
file.set_len((HEADER_SIZE + n) as u64)?;
let mut mmap = unsafe { MmapMut::map_mut(&file)? };
mmap[HEADER_SIZE..HEADER_SIZE + n].copy_from_slice(primary);
Ok(Self { path: path.to_path_buf(), mmap, n, overflow })
}
/// Copy `source`'s file to `path`, mmap the primary section, load overflow into RAM.
/// Avoids iterating all n slots: the file copy is OS-level, overflow loading is O(n_overflow).
pub fn build_from(source: &PersistentCompactIntVec, path: &Path) -> io::Result<Self> {
fs::copy(source.path(), path)?;
let file = OpenOptions::new().read(true).write(true).open(path)?;
let mmap = unsafe { MmapMut::map_mut(&file)? };
let n = source.len();
let n = source.len();
let n_overflow = u64::from_le_bytes(mmap[16..24].try_into().unwrap()) as usize;
let data_offset = HEADER_SIZE + n;
let mut overflow = HashMap::with_capacity(n_overflow);
for i in 0..n_overflow {
let off = data_offset + i * OVERFLOW_ENTRY_SIZE;
let slot = u64::from_le_bytes(mmap[off..off + 8].try_into().unwrap()) as usize;
let value = u32::from_le_bytes(mmap[off + 8..off + 12].try_into().unwrap());
let (slot, value) = parse_overflow_entry(&mmap, data_offset, i);
overflow.insert(slot, value);
}
Ok(Self {
path: path.to_path_buf(),
mmap,
n,
overflow,
})
Ok(Self { path: path.to_path_buf(), mmap, n, overflow })
}
/// Get the value at the given slot, handling overflow if necessary.
pub fn get(&self, slot: usize) -> u32 {
match self.mmap[HEADER_SIZE + slot] {
255 => *self
.overflow
.get(&slot)
.expect("sentinel without overflow entry"),
v => v as u32,
255 => *self.overflow.get(&slot).expect("sentinel without overflow entry"),
v => v as u32,
}
}
@@ -83,61 +69,201 @@ impl PersistentCompactIntVecBuilder {
}
}
pub fn len(&self) -> usize {
self.n
pub fn len(&self) -> usize { self.n }
pub fn is_empty(&self) -> bool { self.n == 0 }
pub fn primary_bytes(&self) -> &[u8] { &self.mmap[HEADER_SIZE..HEADER_SIZE + self.n] }
pub fn primary_bytes_mut(&mut self) -> &mut [u8] { &mut self.mmap[HEADER_SIZE..HEADER_SIZE + self.n] }
pub fn clear_overflow(&mut self) { self.overflow.clear(); }
pub fn sum(&self) -> u64 {
byte_sum(&self.mmap[HEADER_SIZE..HEADER_SIZE + self.n], self.overflow.values().copied())
}
pub fn count_nonzero(&self) -> u64 {
byte_count_nonzero(&self.mmap[HEADER_SIZE..HEADER_SIZE + self.n])
}
pub fn is_empty(&self) -> bool {
self.n == 0
pub fn view(&self) -> IntSliceView<'_> {
// Builder overflow is a HashMap, not sorted raw bytes — convert on the fly
// by collecting into a sorted vec and storing in a thread-local buffer.
// For read-back during building, just call get(slot) directly.
// view() is primarily useful AFTER freeze (on PersistentCompactIntVec).
// Here we expose it via a zero-alloc path: primary only, no overflow raw.
// Callers that need overflow_entries during building use overflow_entries().
let primary = &self.mmap[HEADER_SIZE..HEADER_SIZE + self.n];
IntSliceView::new(primary, &[], 0, self.n)
}
pub fn min(&mut self, other: &PersistentCompactIntVec) {
assert_eq!(self.n, other.len(), "length mismatch");
for (slot, other_val) in other.iter().enumerate() {
if other_val < self.get(slot) {
self.set(slot, other_val);
pub fn overflow_entries(&self) -> impl Iterator<Item = (usize, u32)> + '_ {
self.overflow.iter().map(|(&k, &v)| (k, v))
}
pub fn inc(&mut self, slot: usize) {
let v = self.get(slot);
self.set(slot, v.saturating_add(1));
}
// ── Computation methods ───────────────────────────────────────────────────
/// Increment one counter per 1-bit of `col`. Safe for any group size.
pub fn inc_present(&mut self, col: BitSliceView<'_>) {
let n = self.n;
for (wi, &word) in col.words().iter().enumerate() {
if word == 0 { continue; }
let mut w = word;
while w != 0 {
let bit = w.trailing_zeros() as usize;
let slot = wi * 64 + bit;
if slot < n { self.inc(slot); }
w &= w - 1;
}
}
}
pub fn max(&mut self, other: &PersistentCompactIntVec) {
assert_eq!(self.n, other.len(), "length mismatch");
for (slot, other_val) in other.iter().enumerate() {
if other_val > self.get(slot) {
self.set(slot, other_val);
/// Increment one counter per 1-bit of `col`, using raw u8 arithmetic.
/// Caller guarantees no counter will reach 255 (group size < 255).
pub fn inc_present_fast(&mut self, col: BitSliceView<'_>) {
{
let primary = self.primary_bytes_mut();
let n = primary.len();
for (wi, &word) in col.words().iter().enumerate() {
if word == 0 { continue; }
let mut w = word;
while w != 0 {
let bit = w.trailing_zeros() as usize;
let s = wi * 64 + bit;
if s < n { primary[s] += 1; }
w &= w - 1;
}
}
}
debug_assert!(
!self.primary_bytes().contains(&255),
"sentinel 255 reached in inc_present_fast — group size must be < 255"
);
}
/// Two-pass: primary bytes then overflow. Increments `self[slot]` for each
/// slot where `pred(col[slot])` is true. Safe for any group size.
pub fn inc_predicate(&mut self, col: IntSliceView<'_>, pred: impl Fn(u32) -> bool) {
let n = col.len();
for slot in 0..n {
let b = col.primary_bytes()[slot];
if b < 255 && pred(b as u32) {
self.inc(slot);
}
}
for (slot, val) in col.overflow_entries() {
if pred(val) { self.inc(slot); }
}
}
/// Fast two-pass: raw u8 arithmetic. Caller guarantees no counter reaches 255.
pub fn inc_predicate_fast(&mut self, col: IntSliceView<'_>, pred: impl Fn(u32) -> bool) {
let n = col.len();
{
let primary = self.primary_bytes_mut();
for slot in 0..n {
let b = col.primary_bytes()[slot];
if b < 255 && pred(b as u32) {
primary[slot] += 1;
}
}
}
for (slot, val) in col.overflow_entries() {
if pred(val) { self.primary_bytes_mut()[slot] += 1; }
}
debug_assert!(
!self.primary_bytes().contains(&255),
"sentinel 255 reached in inc_predicate_fast — group size must be < 255"
);
}
pub fn add(&mut self, other: IntSliceView<'_>) {
let n = self.n;
for s in 0..n {
let sb = self.primary_bytes()[s];
let ob = other.primary_bytes()[s];
if sb < 255 && ob < 255 {
let sum = sb as u32 + ob as u32;
if sum < 255 { self.primary_bytes_mut()[s] = sum as u8; }
else { self.set(s, sum); }
} else {
let sv = self.get(s);
let ov = other.get(s);
self.set(s, sv + ov);
}
}
}
pub fn add(&mut self, other: &PersistentCompactIntVec) {
assert_eq!(self.n, other.len(), "length mismatch");
for (slot, other_val) in other.iter().enumerate() {
let cur = self.get(slot);
self.set(slot, cur.checked_add(other_val).expect("u32 overflow in add"));
pub fn min(&mut self, other: IntSliceView<'_>) {
let self_ov: Vec<(usize, u32)> = self.overflow_entries().collect();
let other_ov: HashMap<usize, u32> = other.overflow_entries().collect();
self.clear_overflow();
for (a, &b) in self.primary_bytes_mut().iter_mut().zip(other.primary_bytes()) {
if b < *a { *a = b; }
}
for (slot, self_val) in self_ov {
if let Some(&other_val) = other_ov.get(&slot) {
self.set(slot, self_val.min(other_val));
}
}
}
pub fn diff(&mut self, other: &PersistentCompactIntVec) {
assert_eq!(self.n, other.len(), "length mismatch");
for (slot, other_val) in other.iter().enumerate() {
self.set(slot, self.get(slot).saturating_sub(other_val));
pub fn max(&mut self, other: IntSliceView<'_>) {
for (slot, other_val) in other.overflow_entries() {
let sv = self.get(slot);
self.set(slot, sv.max(other_val));
}
for (a, &b) in self.primary_bytes_mut().iter_mut().zip(other.primary_bytes()) {
if b > *a { *a = b; }
}
}
pub fn diff(&mut self, other: IntSliceView<'_>) {
let n = self.n;
for s in 0..n {
let sb = self.primary_bytes()[s];
let ob = other.primary_bytes()[s];
if sb < 255 {
self.primary_bytes_mut()[s] = if ob < 255 { sb.saturating_sub(ob) } else { 0 };
} else {
let sv = self.get(s);
let ov = if ob < 255 { ob as u32 } else { other.get(s) };
self.set(s, sv.saturating_sub(ov));
}
}
}
pub fn mask_with(&mut self, mask: BitSliceView<'_>) {
let n = self.n;
for (wi, &word) in mask.words().iter().enumerate() {
if word == u64::MAX { continue; }
let mut zeros = !word;
while zeros != 0 {
let bit = zeros.trailing_zeros() as usize;
let s = wi * 64 + bit;
if s < n {
let b = self.primary_bytes()[s];
if b != 0 { self.set(s, 0); }
}
zeros &= zeros - 1;
}
}
}
/// Flush the primary mmap, then write sorted overflow data + index and fix the header.
pub fn close(self) -> io::Result<()> {
self.mmap.flush()?;
let Self {
path,
mmap,
n,
overflow,
} = self;
let Self { path, mmap, n, overflow } = self;
drop(mmap);
let mut entries: Vec<(usize, u32)> = overflow.into_iter().collect();
entries.sort_unstable_by_key(|&(slot, _)| slot);
finalize_pciv(&path, n, &entries)
}
pub fn finish(self) -> io::Result<PersistentCompactIntVec> {
let path = self.path.clone();
self.close()?;
PersistentCompactIntVec::open(&path)
}
}
src/obicompactvec/src/colgroup.rs
+137
View File
@@ -0,0 +1,137 @@
use std::io;
use crate::tempbitvec::{TempBitVec, TempBitVecBuilder};
use crate::tempintvec::TempCompactIntVec;
// ── ColGroup ──────────────────────────────────────────────────────────────────
/// A named subset of columns, identified by their indices within the matrix.
///
/// Defined once at the index level; the same indices are valid across all
/// partitions and layers because the column structure (samples / genomes) is
/// identical everywhere — only the row space (kmer slots) is partitioned.
pub struct ColGroup {
pub name: String,
pub indices: Vec<usize>,
}
impl ColGroup {
pub fn new(name: impl Into<String>, indices: Vec<usize>) -> Self {
Self { name: name.into(), indices }
}
}
// ── MatrixGroupOps ────────────────────────────────────────────────────────────
/// Per-matrix group aggregations.
///
/// `partial_group_presence_count`, `partial_group_sum`, `partial_group_any`,
/// `partial_group_min`, `partial_group_max` are the primitives; each impl must
/// provide all five.
///
/// `partial_group_all` and `partial_group_none` have default implementations
/// derived from `partial_group_presence_count` and should rarely need overriding.
pub trait MatrixGroupOps {
/// Per-slot count of group columns whose value ≥ `threshold`.
fn partial_group_presence_count(&self, g: &ColGroup, threshold: u32) -> io::Result<TempCompactIntVec>;
/// Per-slot sum of values across all group columns.
fn partial_group_sum(&self, g: &ColGroup) -> io::Result<TempCompactIntVec>;
/// Per-slot OR: 1 if any group column has value ≥ `threshold`.
fn partial_group_any(&self, g: &ColGroup, threshold: u32) -> io::Result<TempBitVec>;
/// Per-slot min value across all group columns (0 if group is empty).
fn partial_group_min(&self, g: &ColGroup) -> io::Result<TempCompactIntVec>;
/// Per-slot max value across all group columns (0 if group is empty).
fn partial_group_max(&self, g: &ColGroup) -> io::Result<TempCompactIntVec>;
/// Per-slot AND: 1 if ALL group columns have value ≥ `threshold`.
fn partial_group_all(&self, g: &ColGroup, threshold: u32) -> io::Result<TempBitVec> {
let counts = self.partial_group_presence_count(g, threshold)?;
let n = counts.len();
let n_required = g.indices.len() as u32;
let mut b = TempBitVecBuilder::new(n)?;
b.or_where(counts.view(), |v| v >= n_required);
b.freeze()
}
/// Per-slot NOR: 1 if NO group column has value ≥ `threshold`.
fn partial_group_none(&self, g: &ColGroup, threshold: u32) -> io::Result<TempBitVec> {
let counts = self.partial_group_presence_count(g, threshold)?;
let n = counts.len();
let mut b = TempBitVecBuilder::new(n)?;
b.or_where(counts.view(), |v| v == 0);
b.freeze()
}
}
// ── FilterMask — expression tree for column-based slot filters ────────────────
/// A composable filter expression that can be evaluated against a matrix
/// using only column operations (no MPHF lookup per kmer).
///
/// `threshold` semantics follow [`MatrixGroupOps::partial_group_presence_count`]:
/// a slot contributes to the count when its value is **≥ threshold**.
/// To match the row-level filter (`value > t`), callers should pass `t + 1`.
#[derive(Debug, Clone)]
pub enum FilterMask {
/// Slot passes if count of columns in `indices` with value ≥ `threshold` is ≥ `min_count`.
PresenceGeq { indices: Vec<usize>, threshold: u32, min_count: usize },
/// Slot passes if count of columns in `indices` with value ≥ `threshold` is ≤ `max_count`.
PresenceLeq { indices: Vec<usize>, threshold: u32, max_count: usize },
/// Slot passes if sum of values across `indices` columns is ≥ `min_sum`.
SumGeq { indices: Vec<usize>, min_sum: u32 },
/// Slot passes if sum of values across `indices` columns is ≤ `max_sum`.
SumLeq { indices: Vec<usize>, max_sum: u32 },
/// Slot passes if it passes all sub-expressions. Empty `And` is always true.
And(Vec<FilterMask>),
}
/// Evaluate a [`FilterMask`] against `mat`, returning a per-slot `TempBitVec`
/// where bit=1 means the slot passes the filter.
pub fn eval_filter_mask(expr: &FilterMask, mat: &dyn MatrixGroupOps, n: usize) -> io::Result<TempBitVec> {
match expr {
FilterMask::PresenceGeq { indices, threshold, min_count } => {
let g = ColGroup::new("", indices.clone());
let counts = mat.partial_group_presence_count(&g, *threshold)?;
let mut b = TempBitVecBuilder::new(n)?;
let mc = *min_count as u32;
b.or_where(counts.view(), |v| v >= mc);
b.freeze()
}
FilterMask::PresenceLeq { indices, threshold, max_count } => {
let g = ColGroup::new("", indices.clone());
let counts = mat.partial_group_presence_count(&g, *threshold)?;
let mut b = TempBitVecBuilder::new(n)?;
let mc = *max_count as u32;
b.or_where(counts.view(), |v| v <= mc);
b.freeze()
}
FilterMask::SumGeq { indices, min_sum } => {
let g = ColGroup::new("", indices.clone());
let sums = mat.partial_group_sum(&g)?;
let mut b = TempBitVecBuilder::new(n)?;
let ms = *min_sum;
b.or_where(sums.view(), |v| v >= ms);
b.freeze()
}
FilterMask::SumLeq { indices, max_sum } => {
let g = ColGroup::new("", indices.clone());
let sums = mat.partial_group_sum(&g)?;
let mut b = TempBitVecBuilder::new(n)?;
let ms = *max_sum;
b.or_where(sums.view(), |v| v <= ms);
b.freeze()
}
FilterMask::And(parts) => {
let mut b = TempBitVecBuilder::new_ones(n)?;
for part in parts {
let m = eval_filter_mask(part, mat, n)?;
b.and(m.view());
}
b.freeze()
}
}
}
src/obicompactvec/src/format.rs
+38
View File
@@ -13,6 +13,44 @@ pub const OVERFLOW_ENTRY_SIZE: usize = 12;
// Index entry: slot(u64) + pos(u64) = 16 bytes.
pub const INDEX_ENTRY_SIZE: usize = 16;
/// Sum all values in a compact-int primary byte slice, correcting for overflow sentinels.
///
/// `primary` is the raw `&[u8]` where 255 is a sentinel for large values.
/// `overflow` yields the true values (≥ 255) for each sentinel, in any order.
#[inline]
pub(crate) fn byte_sum(primary: &[u8], overflow: impl Iterator<Item = u32>) -> u64 {
let raw: u64 = primary.iter().map(|&b| b as u64).sum();
let (n, ov) = overflow.fold((0u64, 0u64), |(n, s), v| (n + 1, s + v as u64));
raw - 255 * n + ov
}
/// Count non-zero values in a compact-int primary byte slice.
///
/// Overflow sentinels (255) are always non-zero by construction, so a single
/// `b != 0` test is sufficient — no overflow map lookup needed.
#[inline]
pub(crate) fn byte_count_nonzero(primary: &[u8]) -> u64 {
primary.iter().filter(|&&b| b != 0).count() as u64
}
/// Parse a single overflow entry `(slot, value)` from a byte slice.
#[inline]
pub fn parse_overflow_entry(data: &[u8], base: usize, i: usize) -> (usize, u32) {
let off = base + i * OVERFLOW_ENTRY_SIZE;
let slot = u64::from_le_bytes(data[off..off+8].try_into().unwrap()) as usize;
let value = u32::from_le_bytes(data[off+8..off+12].try_into().unwrap());
(slot, value)
}
/// Parse a single sparse-index entry `(slot, pos)` from a byte slice.
#[inline]
pub fn parse_index_entry(data: &[u8], base: usize, i: usize) -> (usize, usize) {
let off = base + i * INDEX_ENTRY_SIZE;
let slot = u64::from_le_bytes(data[off..off+8].try_into().unwrap()) as usize;
let pos = u64::from_le_bytes(data[off+8..off+16].try_into().unwrap()) as usize;
(slot, pos)
}
// Sparse index target: ≤ 32 KB in L1 cache (16 B per entry → 2048 entries).
pub const L1_INDEX_ENTRIES: usize = 2048;
src/obicompactvec/src/intmatrix.rs
+144 -220
View File
@@ -1,4 +1,3 @@
use std::cmp::Ordering;
use std::fs::{self, File};
use std::io::{self, BufWriter, Write as _};
use std::path::{Path, PathBuf};
@@ -7,10 +6,15 @@ use memmap2::Mmap;
use ndarray::{Array1, Array2};
use rayon::prelude::*;
use crate::bitmatrix::{pairwise_matrix, pairwise2_matrix};
use crate::builder::PersistentCompactIntVecBuilder;
use crate::format::{HEADER_SIZE, INDEX_ENTRY_SIZE, OVERFLOW_ENTRY_SIZE};
use crate::colgroup::{ColGroup, MatrixGroupOps};
use crate::format::{HEADER_SIZE, OVERFLOW_ENTRY_SIZE};
use crate::meta::MatrixMeta;
use crate::reader::PersistentCompactIntVec;
use crate::tempbitvec::{TempBitVec, TempBitVecBuilder};
use crate::tempintvec::{TempCompactIntVec, TempCompactIntVecBuilder};
use crate::views::IntSliceView;
fn col_path(dir: &Path, col: usize) -> PathBuf {
dir.join(format!("col_{col:06}.pciv"))
@@ -41,9 +45,7 @@ impl ColumnarCompactIntMatrix {
}
pub(crate) fn fill_row(&self, slot: usize, buf: &mut [u32]) {
for (c, col) in self.cols.iter().enumerate() {
buf[c] = col.get(slot);
}
for (c, col) in self.cols.iter().enumerate() { buf[c] = col.get(slot); }
}
pub(crate) fn sum(&self) -> Array1<u64> {
@@ -63,49 +65,26 @@ impl ColumnarCompactIntMatrix {
}
pub(crate) fn partial_bray_dist_matrix(&self) -> Array2<u64> {
self.pairwise_u64(|i, j| self.col(i).partial_bray_dist(self.col(j)))
pairwise_matrix(self.n_cols(), |i, j| self.col(i).partial_bray_dist(self.col(j)))
}
pub(crate) fn partial_euclidean_dist_matrix(&self) -> Array2<f64> {
self.pairwise(|i, j| self.col(i).partial_euclidean_dist(self.col(j)))
pairwise_matrix(self.n_cols(), |i, j| self.col(i).partial_euclidean_dist(self.col(j)))
}
pub(crate) fn partial_threshold_jaccard_dist_matrix(
&self, threshold: u32,
) -> (Array2<u64>, Array2<u64>) {
let n = self.n_cols();
let pairs = upper_pairs(n);
let results: Vec<(usize, usize, u64, u64)> = pairs
.into_par_iter()
.map(|(i, j)| {
let (inter, union) =
self.col(i).partial_threshold_jaccard_dist(self.col(j), threshold);
(i, j, inter, union)
})
.collect();
let mut inter_m = Array2::zeros((n, n));
let mut union_m = Array2::zeros((n, n));
for (i, j, inter, union) in results {
inter_m[[i, j]] = inter; inter_m[[j, i]] = inter;
union_m[[i, j]] = union; union_m[[j, i]] = union;
}
(inter_m, union_m)
pub(crate) fn partial_threshold_jaccard_dist_matrix(&self, threshold: u32) -> (Array2<u64>, Array2<u64>) {
pairwise2_matrix(self.n_cols(), |i, j| self.col(i).partial_threshold_jaccard_dist(self.col(j), threshold))
}
pub(crate) fn partial_relfreq_bray_dist_matrix(&self, col_sums: &Array1<u64>) -> Array2<f64> {
self.pairwise(|i, j| {
pairwise_matrix(self.n_cols(), |i, j| {
self.col(i).partial_relfreq_bray_dist(self.col(j), col_sums[i] as f64, col_sums[j] as f64)
})
}
pub(crate) fn partial_relfreq_euclidean_dist_matrix(&self, col_sums: &Array1<u64>) -> Array2<f64> {
self.pairwise(|i, j| {
pairwise_matrix(self.n_cols(), |i, j| {
self.col(i).partial_relfreq_euclidean_dist(self.col(j), col_sums[i] as f64, col_sums[j] as f64)
})
}
pub(crate) fn partial_hellinger_euclidean_dist_matrix(&self, col_sums: &Array1<u64>) -> Array2<f64> {
self.pairwise(|i, j| {
pairwise_matrix(self.n_cols(), |i, j| {
self.col(i).partial_hellinger_euclidean_dist(self.col(j), col_sums[i] as f64, col_sums[j] as f64)
})
}
@@ -118,20 +97,6 @@ impl ColumnarCompactIntMatrix {
meta.n_cols += 1;
meta.save(dir)
}
fn pairwise(&self, f: impl Fn(usize, usize) -> f64 + Sync) -> Array2<f64> {
let n = self.n_cols();
let results: Vec<(usize, usize, f64)> = upper_pairs(n)
.into_par_iter().map(|(i, j)| (i, j, f(i, j))).collect();
fill_symmetric(n, results.into_iter().map(|(i, j, v)| (i, j, v, v)))
}
fn pairwise_u64(&self, f: impl Fn(usize, usize) -> u64 + Sync) -> Array2<u64> {
let n = self.n_cols();
let results: Vec<(usize, usize, u64)> = upper_pairs(n)
.into_par_iter().map(|(i, j)| (i, j, f(i, j))).collect();
fill_symmetric(n, results.into_iter().map(|(i, j, v)| (i, j, v, v)))
}
}
// ── PackedCompactIntMatrix ────────────────────────────────────────────────────
@@ -139,13 +104,10 @@ impl ColumnarCompactIntMatrix {
const PCMX_MAGIC: [u8; 4] = *b"PCMX";
const PCMX_HEADER: usize = 24; // magic(4) + pad(4) + n_rows(8) + n_cols(8)
/// Per-column metadata pre-parsed from the embedded PCIV header.
struct ColInfo {
primary_start: usize, // absolute mmap offset to primary array
data_offset: usize, // absolute mmap offset to overflow array
primary_start: usize,
data_offset: usize,
n_overflow: usize,
step: usize,
index: Vec<(usize, usize)>,
}
pub struct PackedCompactIntMatrix {
@@ -171,61 +133,31 @@ impl PackedCompactIntMatrix {
for c in 0..n_cols {
let off_pos = PCMX_HEADER + c * 8;
let col_base = u64::from_le_bytes(mmap[off_pos..off_pos+8].try_into().unwrap()) as usize;
// Parse embedded PCIV header at col_base
let n_ov = u64::from_le_bytes(mmap[col_base+16..col_base+24].try_into().unwrap()) as usize;
let n_idx = u64::from_le_bytes(mmap[col_base+24..col_base+32].try_into().unwrap()) as usize;
let step = u64::from_le_bytes(mmap[col_base+32..col_base+40].try_into().unwrap()) as usize;
let n_pciv = u64::from_le_bytes(mmap[col_base+8..col_base+16].try_into().unwrap()) as usize;
let n_ov = u64::from_le_bytes(mmap[col_base+16..col_base+24].try_into().unwrap()) as usize;
let n_pciv = u64::from_le_bytes(mmap[col_base+8..col_base+16].try_into().unwrap()) as usize;
let primary_start = col_base + HEADER_SIZE;
let data_offset = primary_start + n_pciv;
let index_offset = data_offset + n_ov * OVERFLOW_ENTRY_SIZE;
let mut index = Vec::with_capacity(n_idx);
for i in 0..n_idx {
let ioff = index_offset + i * INDEX_ENTRY_SIZE;
let slot = u64::from_le_bytes(mmap[ioff..ioff+8].try_into().unwrap()) as usize;
let pos = u64::from_le_bytes(mmap[ioff+8..ioff+16].try_into().unwrap()) as usize;
index.push((slot, pos));
}
columns.push(ColInfo { primary_start, data_offset, n_overflow: n_ov, step, index });
columns.push(ColInfo { primary_start, data_offset, n_overflow: n_ov });
}
Ok(Self { mmap, n_rows, n_cols, columns })
}
#[inline]
pub(crate) fn get(&self, col: usize, slot: usize) -> u32 {
let ci = &self.columns[col];
let v = self.mmap[ci.primary_start + slot];
if v < 255 { return v as u32; }
self.overflow_get(ci, slot)
pub(crate) fn col_view(&self, c: usize) -> IntSliceView<'_> {
let ci = &self.columns[c];
let primary = &self.mmap[ci.primary_start..ci.primary_start + self.n_rows];
let overflow_raw = &self.mmap[ci.data_offset..ci.data_offset + ci.n_overflow * OVERFLOW_ENTRY_SIZE];
IntSliceView::new(primary, overflow_raw, ci.n_overflow, self.n_rows)
}
fn overflow_get(&self, ci: &ColInfo, slot: usize) -> u32 {
let (pos_start, pos_end) = if ci.step == 0 {
(0, ci.n_overflow)
} else {
let i = ci.index.partition_point(|&(s, _)| s <= slot).saturating_sub(1);
let start = ci.index[i].1;
let end = if i + 1 < ci.index.len() { ci.index[i+1].1 } else { ci.n_overflow };
(start, end)
};
let mut lo = pos_start;
let mut hi = pos_end;
while lo < hi {
let mid = lo + (hi - lo) / 2;
let off = ci.data_offset + mid * OVERFLOW_ENTRY_SIZE;
let stored = u64::from_le_bytes(self.mmap[off..off+8].try_into().unwrap()) as usize;
match stored.cmp(&slot) {
Ordering::Equal => return u32::from_le_bytes(self.mmap[off+8..off+12].try_into().unwrap()),
Ordering::Less => lo = mid + 1,
Ordering::Greater => hi = mid,
}
}
panic!("slot {slot} marked overflow but not found")
pub(crate) fn col_persist(&self, c: usize, path: &Path) -> io::Result<PersistentCompactIntVecBuilder> {
let view = self.col_view(c);
let overflow: std::collections::HashMap<usize, u32> = view.overflow_entries().collect();
PersistentCompactIntVecBuilder::from_raw_primary(view.primary_bytes(), overflow, path)
}
#[inline]
pub(crate) fn get(&self, col: usize, slot: usize) -> u32 { self.col_view(col).get(slot) }
pub(crate) fn fill_row(&self, slot: usize, buf: &mut [u32]) {
for c in 0..self.n_cols { buf[c] = self.get(c, slot); }
}
@@ -236,152 +168,85 @@ impl PackedCompactIntMatrix {
pub(crate) fn sum(&self) -> Array1<u64> {
Array1::from_vec(
(0..self.n_cols).into_par_iter()
.map(|c| (0..self.n_rows).map(|s| self.get(c, s) as u64).sum())
.collect()
(0..self.n_cols).into_par_iter().map(|c| self.col_view(c).sum()).collect()
)
}
pub(crate) fn count_nonzero(&self) -> Array1<u64> {
Array1::from_vec(
(0..self.n_cols).into_par_iter()
.map(|c| (0..self.n_rows).filter(|&s| self.get(c, s) > 0).count() as u64)
.collect()
(0..self.n_cols).into_par_iter().map(|c| self.col_view(c).count_nonzero()).collect()
)
}
// ── Pair primitives ───────────────────────────────────────────────────────
fn pair_partial_bray(&self, i: usize, j: usize) -> u64 {
(0..self.n_rows).map(|s| self.get(i, s).min(self.get(j, s)) as u64).sum()
self.col_view(i).iter().zip(self.col_view(j).iter()).map(|(a, b)| a.min(b) as u64).sum()
}
fn pair_partial_euclidean(&self, i: usize, j: usize) -> f64 {
(0..self.n_rows).map(|s| {
let d = self.get(i, s) as f64 - self.get(j, s) as f64;
d * d
}).sum()
self.col_view(i).iter().zip(self.col_view(j).iter())
.map(|(a, b)| { let d = a as f64 - b as f64; d * d }).sum()
}
fn pair_partial_threshold_jaccard(&self, i: usize, j: usize, t: u32) -> (u64, u64) {
let (mut inter, mut union) = (0u64, 0u64);
for s in 0..self.n_rows {
let a = self.get(i, s) >= t;
let b = self.get(j, s) >= t;
if a && b { inter += 1; }
if a || b { union += 1; }
}
(inter, union)
self.col_view(i).iter().zip(self.col_view(j).iter())
.fold((0u64, 0u64), |(inter, uni), (a, b)| {
let ap = a >= t; let bp = b >= t;
(inter + (ap & bp) as u64, uni + (ap | bp) as u64)
})
}
fn pair_partial_relfreq_bray(&self, i: usize, j: usize, si: f64, sj: f64) -> f64 {
if si == 0.0 || sj == 0.0 { return 0.0; }
(0..self.n_rows).map(|s| {
(self.get(i, s) as f64 / si).min(self.get(j, s) as f64 / sj)
}).sum()
self.col_view(i).iter().zip(self.col_view(j).iter())
.map(|(a, b)| (a as f64 / si).min(b as f64 / sj)).sum()
}
fn pair_partial_relfreq_euclidean(&self, i: usize, j: usize, si: f64, sj: f64) -> f64 {
if si == 0.0 || sj == 0.0 { return 0.0; }
(0..self.n_rows).map(|s| {
let d = self.get(i, s) as f64 / si - self.get(j, s) as f64 / sj;
d * d
}).sum()
self.col_view(i).iter().zip(self.col_view(j).iter())
.map(|(a, b)| { let d = a as f64 / si - b as f64 / sj; d * d }).sum()
}
fn pair_partial_hellinger(&self, i: usize, j: usize, si: f64, sj: f64) -> f64 {
if si == 0.0 || sj == 0.0 { return 0.0; }
(0..self.n_rows).map(|s| {
let d = (self.get(i, s) as f64 / si).sqrt() - (self.get(j, s) as f64 / sj).sqrt();
d * d
}).sum()
}
// ── Matrix methods ────────────────────────────────────────────────────────
fn pairwise<T>(&self, f: impl Fn(usize, usize) -> T + Sync) -> Array2<T>
where T: Clone + Default + Send {
let n = self.n_cols;
let results: Vec<(usize, usize, T)> = upper_pairs(n)
.into_par_iter().map(|(i, j)| (i, j, f(i, j))).collect();
fill_symmetric(n, results.into_iter().map(|(i, j, v)| { let w = v.clone(); (i, j, v, w) }))
}
fn pairwise_u64(&self, f: impl Fn(usize, usize) -> u64 + Sync) -> Array2<u64> {
let n = self.n_cols;
let results: Vec<(usize, usize, u64)> = upper_pairs(n)
.into_par_iter().map(|(i, j)| (i, j, f(i, j))).collect();
fill_symmetric(n, results.into_iter().map(|(i, j, v)| (i, j, v, v)))
self.col_view(i).iter().zip(self.col_view(j).iter())
.map(|(a, b)| { let d = (a as f64 / si).sqrt() - (b as f64 / sj).sqrt(); d * d }).sum()
}
pub(crate) fn partial_bray_dist_matrix(&self) -> Array2<u64> {
self.pairwise_u64(|i, j| self.pair_partial_bray(i, j))
pairwise_matrix(self.n_cols, |i, j| self.pair_partial_bray(i, j))
}
pub(crate) fn partial_euclidean_dist_matrix(&self) -> Array2<f64> {
self.pairwise(|i, j| self.pair_partial_euclidean(i, j))
pairwise_matrix(self.n_cols, |i, j| self.pair_partial_euclidean(i, j))
}
pub(crate) fn partial_threshold_jaccard_dist_matrix(&self, t: u32) -> (Array2<u64>, Array2<u64>) {
let n = self.n_cols;
let results: Vec<(usize, usize, u64, u64)> = upper_pairs(n)
.into_par_iter()
.map(|(i, j)| { let (inter, union) = self.pair_partial_threshold_jaccard(i, j, t); (i, j, inter, union) })
.collect();
let mut inter_m = Array2::zeros((n, n));
let mut union_m = Array2::zeros((n, n));
for (i, j, inter, union) in results {
inter_m[[i, j]] = inter; inter_m[[j, i]] = inter;
union_m[[i, j]] = union; union_m[[j, i]] = union;
}
(inter_m, union_m)
pairwise2_matrix(self.n_cols, |i, j| self.pair_partial_threshold_jaccard(i, j, t))
}
pub(crate) fn partial_relfreq_bray_dist_matrix(&self, col_sums: &Array1<u64>) -> Array2<f64> {
self.pairwise(|i, j| self.pair_partial_relfreq_bray(i, j, col_sums[i] as f64, col_sums[j] as f64))
pairwise_matrix(self.n_cols, |i, j| self.pair_partial_relfreq_bray(i, j, col_sums[i] as f64, col_sums[j] as f64))
}
pub(crate) fn partial_relfreq_euclidean_dist_matrix(&self, col_sums: &Array1<u64>) -> Array2<f64> {
self.pairwise(|i, j| self.pair_partial_relfreq_euclidean(i, j, col_sums[i] as f64, col_sums[j] as f64))
pairwise_matrix(self.n_cols, |i, j| self.pair_partial_relfreq_euclidean(i, j, col_sums[i] as f64, col_sums[j] as f64))
}
pub(crate) fn partial_hellinger_euclidean_dist_matrix(&self, col_sums: &Array1<u64>) -> Array2<f64> {
self.pairwise(|i, j| self.pair_partial_hellinger(i, j, col_sums[i] as f64, col_sums[j] as f64))
pairwise_matrix(self.n_cols, |i, j| self.pair_partial_hellinger(i, j, col_sums[i] as f64, col_sums[j] as f64))
}
}
/// Build `counts/matrix.pcmx` from existing `col_*.pciv` files.
pub fn pack_compact_int_matrix(dir: &Path) -> io::Result<()> {
let packed_path = dir.join("matrix.pcmx");
if packed_path.exists() {
// Matrix complete; remove any leftover column files from a killed cleanup.
if let Ok(meta) = MatrixMeta::load(dir) {
for c in 0..meta.n_cols { let _ = fs::remove_file(col_path(dir, c)); }
let _ = fs::remove_file(dir.join("meta.json"));
}
return Ok(());
}
let meta = MatrixMeta::load(dir)?;
let meta = MatrixMeta::load(dir)?;
let n_cols = meta.n_cols;
// Compute offsets from file sizes — no column data loaded into RAM.
let col_sizes: Vec<u64> = (0..n_cols)
.map(|c| fs::metadata(col_path(dir, c)).map(|m| m.len()))
.collect::<io::Result<_>>()?;
let header_size = (PCMX_HEADER + n_cols * 8) as u64;
let mut col_offset = header_size;
let mut offsets = Vec::with_capacity(n_cols);
for &size in &col_sizes {
offsets.push(col_offset);
col_offset += size;
}
// Write to a temp file; rename atomically so a killed process never leaves
// a truncated matrix.pcmx that would be mistaken for a complete file.
for &size in &col_sizes { offsets.push(col_offset); col_offset += size; }
let tmp_path = dir.join("matrix.pcmx.tmp");
let mut out = BufWriter::new(File::create(&tmp_path)?);
out.write_all(&PCMX_MAGIC)?;
@@ -389,13 +254,10 @@ pub fn pack_compact_int_matrix(dir: &Path) -> io::Result<()> {
out.write_all(&(meta.n as u64).to_le_bytes())?;
out.write_all(&(n_cols as u64).to_le_bytes())?;
for &off in &offsets { out.write_all(&off.to_le_bytes())?; }
for c in 0..n_cols {
io::copy(&mut File::open(col_path(dir, c))?, &mut out)?;
}
for c in 0..n_cols { io::copy(&mut File::open(col_path(dir, c))?, &mut out)?; }
out.flush()?;
drop(out);
fs::rename(&tmp_path, &packed_path)?;
for c in 0..n_cols { fs::remove_file(col_path(dir, c))?; }
fs::remove_file(dir.join("meta.json"))?;
Ok(())
@@ -409,18 +271,14 @@ pub enum PersistentCompactIntMatrix {
}
impl PersistentCompactIntMatrix {
/// Open from `layer_dir`, auto-detecting Packed or Columnar.
pub fn open(layer_dir: &Path) -> io::Result<Self> {
let counts_dir = layer_dir.join("counts");
if counts_dir.join("matrix.pcmx").exists() {
return Ok(Self::Packed(PackedCompactIntMatrix::open(&counts_dir.join("matrix.pcmx"))?));
}
if MatrixMeta::load(&counts_dir).is_ok() {
return Ok(Self::Columnar(ColumnarCompactIntMatrix::open(&counts_dir)?));
}
Err(io::Error::new(
io::ErrorKind::NotFound,
format!("no count matrix found in {} — run 'obikmer upgrade'", layer_dir.display()),
@@ -430,7 +288,6 @@ impl PersistentCompactIntMatrix {
pub fn n(&self) -> usize {
match self { Self::Columnar(m) => m.n(), Self::Packed(m) => m.n_rows }
}
pub fn n_cols(&self) -> usize {
match self { Self::Columnar(m) => m.n_cols(), Self::Packed(m) => m.n_cols }
}
@@ -442,22 +299,32 @@ impl PersistentCompactIntMatrix {
}
}
pub fn col_view(&self, c: usize) -> IntSliceView<'_> {
match self {
Self::Columnar(m) => m.col(c).view(),
Self::Packed(m) => m.col_view(c),
}
}
pub fn col_persist(&self, c: usize, path: &Path) -> io::Result<PersistentCompactIntVecBuilder> {
match self {
Self::Columnar(m) => PersistentCompactIntVecBuilder::build_from(m.col(c), path),
Self::Packed(m) => m.col_persist(c, path),
}
}
pub fn row(&self, slot: usize) -> Box<[u32]> {
match self { Self::Columnar(m) => m.row(slot), Self::Packed(m) => m.row(slot) }
}
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) }
}
pub fn sum(&self) -> Array1<u64> {
match self { Self::Columnar(m) => m.sum(), Self::Packed(m) => m.sum() }
}
pub fn count_nonzero(&self) -> Array1<u64> {
match self { Self::Columnar(m) => m.count_nonzero(), Self::Packed(m) => m.count_nonzero() }
}
pub fn partial_bray_dist_matrix(&self) -> Array2<u64> {
match self { Self::Columnar(m) => m.partial_bray_dist_matrix(), Self::Packed(m) => m.partial_bray_dist_matrix() }
}
@@ -476,7 +343,6 @@ impl PersistentCompactIntMatrix {
pub fn partial_hellinger_euclidean_dist_matrix(&self, col_sums: &Array1<u64>) -> Array2<f64> {
match self { Self::Columnar(m) => m.partial_hellinger_euclidean_dist_matrix(col_sums), Self::Packed(m) => m.partial_hellinger_euclidean_dist_matrix(col_sums) }
}
pub fn append_column(dir: &Path, value_of: impl Fn(usize) -> u32) -> io::Result<()> {
ColumnarCompactIntMatrix::append_column(dir, value_of)
}
@@ -492,12 +358,12 @@ impl ColumnWeights for PersistentCompactIntMatrix {
}
impl CountPartials for PersistentCompactIntMatrix {
fn partial_bray(&self) -> Array2<u64> { self.partial_bray_dist_matrix() }
fn partial_euclidean(&self) -> Array2<f64> { self.partial_euclidean_dist_matrix() }
fn partial_bray(&self) -> Array2<u64> { self.partial_bray_dist_matrix() }
fn partial_euclidean(&self) -> Array2<f64> { self.partial_euclidean_dist_matrix() }
fn partial_threshold_jaccard(&self, t: u32) -> (Array2<u64>, Array2<u64>) { self.partial_threshold_jaccard_dist_matrix(t) }
fn partial_relfreq_bray(&self, g: &Array1<u64>) -> Array2<f64> { self.partial_relfreq_bray_dist_matrix(g) }
fn partial_relfreq_euclidean(&self, g: &Array1<u64>) -> Array2<f64> { self.partial_relfreq_euclidean_dist_matrix(g) }
fn partial_hellinger(&self, g: &Array1<u64>) -> Array2<f64> { self.partial_hellinger_euclidean_dist_matrix(g) }
fn partial_relfreq_bray(&self, g: &Array1<u64>) -> Array2<f64> { self.partial_relfreq_bray_dist_matrix(g) }
fn partial_relfreq_euclidean(&self, g: &Array1<u64>) -> Array2<f64> { self.partial_relfreq_euclidean_dist_matrix(g) }
fn partial_hellinger(&self, g: &Array1<u64>) -> Array2<f64> { self.partial_hellinger_euclidean_dist_matrix(g) }
}
// ── Builder ───────────────────────────────────────────────────────────────────
@@ -513,30 +379,88 @@ impl PersistentCompactIntMatrixBuilder {
fs::create_dir_all(dir)?;
Ok(Self { dir: dir.to_path_buf(), n, n_cols: 0 })
}
pub fn n(&self) -> usize { self.n }
pub fn n_cols(&self) -> usize { self.n_cols }
pub fn add_col(&mut self) -> io::Result<PersistentCompactIntVecBuilder> {
let path = col_path(&self.dir, self.n_cols);
self.n_cols += 1;
PersistentCompactIntVecBuilder::new(self.n, &path)
}
pub fn add_col_from(&mut self, src: &TempCompactIntVec) -> io::Result<()> {
src.make_persistent(&col_path(&self.dir, self.n_cols))?;
self.n_cols += 1;
Ok(())
}
pub fn add_col_from_bit(&mut self, src: &TempBitVec) -> io::Result<()> {
let path = col_path(&self.dir, self.n_cols);
self.n_cols += 1;
let mut b = PersistentCompactIntVecBuilder::new(self.n, &path)?;
b.inc_present(src.view());
b.close()
}
pub fn close(self) -> io::Result<()> {
MatrixMeta { n: self.n, n_cols: self.n_cols }.save(&self.dir)
}
}
// ── Helpers ───────────────────────────────────────────────────────────────────
// ── MatrixGroupOps ────────────────────────────────────────────────────────────
fn upper_pairs(n: usize) -> Vec<(usize, usize)> {
(0..n).flat_map(|i| (i + 1..n).map(move |j| (i, j))).collect()
}
impl MatrixGroupOps for PersistentCompactIntMatrix {
fn partial_group_presence_count(&self, g: &ColGroup, threshold: u32) -> io::Result<TempCompactIntVec> {
let n = self.n();
if g.indices.len() < 255 {
let mut builder = TempCompactIntVecBuilder::new(n)?;
for &c in &g.indices {
builder.inc_predicate_fast(self.col_view(c), |v| v >= threshold);
}
builder.freeze()
} else {
let mut result = TempCompactIntVecBuilder::new(n)?;
for chunk in g.indices.chunks(254) {
let mut chunk_b = TempCompactIntVecBuilder::new(n)?;
for &c in chunk {
chunk_b.inc_predicate_fast(self.col_view(c), |v| v >= threshold);
}
let frozen = chunk_b.freeze()?;
result.add(frozen.view());
}
result.freeze()
}
}
fn fill_symmetric<T>(n: usize, vals: impl Iterator<Item = (usize, usize, T, T)>) -> Array2<T>
where T: Clone + Default {
let mut m = Array2::from_elem((n, n), T::default());
for (i, j, vij, vji) in vals { m[[i, j]] = vij; m[[j, i]] = vji; }
m
fn partial_group_sum(&self, g: &ColGroup) -> io::Result<TempCompactIntVec> {
let n = self.n();
let mut result = TempCompactIntVecBuilder::new(n)?;
for &c in &g.indices { result.add(self.col_view(c)); }
result.freeze()
}
fn partial_group_any(&self, g: &ColGroup, threshold: u32) -> io::Result<TempBitVec> {
let n = self.n();
let mut result = TempBitVecBuilder::new(n)?;
for &c in &g.indices {
result.or_where(self.col_view(c), |v| v >= threshold);
}
result.freeze()
}
fn partial_group_min(&self, g: &ColGroup) -> io::Result<TempCompactIntVec> {
let n = self.n();
let mut result = TempCompactIntVecBuilder::new(n)?;
if let Some((&first, rest)) = g.indices.split_first() {
result.add(self.col_view(first));
for &c in rest { result.min(self.col_view(c)); }
}
result.freeze()
}
fn partial_group_max(&self, g: &ColGroup) -> io::Result<TempCompactIntVec> {
let n = self.n();
let mut result = TempCompactIntVecBuilder::new(n)?;
for &c in &g.indices { result.max(self.col_view(c)); }
result.freeze()
}
}
src/obicompactvec/src/layer_meta.rs
+1 -6
View File
@@ -23,11 +23,6 @@ impl LayerMeta {
}
fn parse(s: &str) -> Option<Self> {
let key = "\"n\":";
let pos = s.find(key)? + key.len();
let rest = s[pos..].trim_start();
let end = rest.find(|c: char| !c.is_ascii_digit()).unwrap_or(rest.len());
let n = rest[..end].parse().ok()?;
Some(Self { n })
Some(Self { n: crate::meta::field(s, "n")? })
}
}
src/obicompactvec/src/lib.rs
+9 -1
View File
@@ -1,20 +1,28 @@
mod bitvec;
mod bitmatrix;
mod builder;
mod colgroup;
mod format;
mod intmatrix;
mod layer_meta;
mod meta;
mod reader;
mod tempbitvec;
mod tempintvec;
mod views;
pub mod traits;
pub use bitvec::{BitIter, PersistentBitVec, PersistentBitVecBuilder};
pub use bitmatrix::{PersistentBitMatrix, PersistentBitMatrixBuilder, pack_bit_matrix};
pub use builder::PersistentCompactIntVecBuilder;
pub use colgroup::{ColGroup, FilterMask, MatrixGroupOps, eval_filter_mask};
pub use intmatrix::{PersistentCompactIntMatrix, PersistentCompactIntMatrixBuilder, pack_compact_int_matrix};
pub use layer_meta::LayerMeta;
pub use reader::PersistentCompactIntVec;
pub use reader::{PersistentCompactIntVec, Iter as CompactIntVecIter};
pub use tempbitvec::{TempBitVec, TempBitVecBuilder};
pub use tempintvec::{TempCompactIntVec, TempCompactIntVecBuilder};
pub use traits::{BitPartials, ColumnWeights, CountPartials};
pub use views::{BitSliceView, BitSliceIter, IntSliceView, IntSliceViewIter};
#[cfg(test)]
#[path = "tests/mod.rs"]
src/obicompactvec/src/meta.rs
+1 -1
View File
@@ -23,7 +23,7 @@ fn parse(s: &str) -> Option<MatrixMeta> {
Some(MatrixMeta { n: field(s, "n")?, n_cols: field(s, "n_cols")? })
}
fn field(s: &str, name: &str) -> Option<usize> {
pub(crate) fn field(s: &str, name: &str) -> Option<usize> {
let key = format!("\"{}\":", name);
let pos = s.find(&key)? + key.len();
let rest = s[pos..].trim_start();
src/obicompactvec/src/reader.rs
+70 -211
View File
@@ -4,7 +4,8 @@ use std::path::{Path, PathBuf};
use memmap2::Mmap;
use crate::format::{HEADER_SIZE, INDEX_ENTRY_SIZE, MAGIC, OVERFLOW_ENTRY_SIZE};
use crate::format::{byte_count_nonzero, byte_sum, HEADER_SIZE, MAGIC, OVERFLOW_ENTRY_SIZE, parse_index_entry};
use crate::views::IntSliceView;
pub struct PersistentCompactIntVec {
mmap: Mmap,
@@ -18,100 +19,60 @@ pub struct PersistentCompactIntVec {
}
impl PersistentCompactIntVec {
/// Opens a persistent compact int vector from the given path.
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,
"PCIV file too short",
));
return Err(io::Error::new(io::ErrorKind::InvalidData, "PCIV file too short"));
}
if &mmap[0..4] != &MAGIC {
return Err(io::Error::new(io::ErrorKind::InvalidData, "bad PCIV 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;
let n_overflow = u64::from_le_bytes(mmap[16..24].try_into().unwrap()) as usize;
let n_index = u64::from_le_bytes(mmap[24..32].try_into().unwrap()) as usize;
let step = u64::from_le_bytes(mmap[32..40].try_into().unwrap()) as usize;
let n_index = u64::from_le_bytes(mmap[24..32].try_into().unwrap()) as usize;
let step = u64::from_le_bytes(mmap[32..40].try_into().unwrap()) as usize;
let primary_offset = HEADER_SIZE;
let data_offset = primary_offset + n;
let index_offset = data_offset + n_overflow * OVERFLOW_ENTRY_SIZE;
let data_offset = primary_offset + n;
let index_offset = data_offset + n_overflow * OVERFLOW_ENTRY_SIZE;
let mut index = Vec::with_capacity(n_index);
for i in 0..n_index {
let off = index_offset + i * INDEX_ENTRY_SIZE;
let slot = u64::from_le_bytes(mmap[off..off + 8].try_into().unwrap()) as usize;
let pos = u64::from_le_bytes(mmap[off + 8..off + 16].try_into().unwrap()) as usize;
index.push((slot, pos));
index.push(parse_index_entry(&mmap, index_offset, i));
}
Ok(Self {
mmap,
n,
n_overflow,
step,
index,
primary_offset,
data_offset,
path: path.to_path_buf(),
})
Ok(Self { mmap, n, n_overflow, step, index, primary_offset, data_offset, path: path.to_path_buf() })
}
/// Returns the path of the compact int vector file.
pub fn path(&self) -> &Path {
&self.path
}
pub fn path(&self) -> &Path { &self.path }
pub fn len(&self) -> usize { self.n }
pub fn is_empty(&self) -> bool { self.n == 0 }
/// Returns the length of the compact int vector.
pub fn len(&self) -> usize {
self.n
}
/// Returns whether the compact int vector is empty.
pub fn is_empty(&self) -> bool {
self.n == 0
}
/// Returns the value at the given slot.
pub fn get(&self, slot: usize) -> u32 {
match self.mmap[self.primary_offset + slot] {
255 => self.overflow_get(slot),
v => v as u32,
v => v as u32,
}
}
/// Returns the value at the given slot from the overflow region.
fn overflow_get(&self, slot: usize) -> u32 {
let pos_start;
let pos_end;
if self.step == 0 {
pos_start = 0;
pos_end = self.n_overflow;
let (pos_start, pos_end) = if self.step == 0 {
(0, self.n_overflow)
} else {
let i = self
.index
.partition_point(|&(s, _)| s <= slot)
.saturating_sub(1);
pos_start = self.index[i].1;
pos_end = if i + 1 < self.index.len() {
self.index[i + 1].1
} else {
self.n_overflow
};
}
let i = self.index.partition_point(|&(s, _)| s <= slot).saturating_sub(1);
let start = self.index[i].1;
let end = if i + 1 < self.index.len() { self.index[i + 1].1 } else { self.n_overflow };
(start, end)
};
let mut lo = pos_start;
let mut hi = pos_end;
while lo < hi {
let mid = lo + (hi - lo) / 2;
match self.data_slot(mid).cmp(&slot) {
std::cmp::Ordering::Equal => return self.data_value(mid),
std::cmp::Ordering::Less => lo = mid + 1,
std::cmp::Ordering::Equal => return self.data_value(mid),
std::cmp::Ordering::Less => lo = mid + 1,
std::cmp::Ordering::Greater => hi = mid,
}
}
@@ -119,144 +80,91 @@ impl PersistentCompactIntVec {
}
#[inline]
/// Returns the slot at the given index in the overflow region.
fn data_slot(&self, i: usize) -> usize {
let off = self.data_offset + i * OVERFLOW_ENTRY_SIZE;
u64::from_le_bytes(self.mmap[off..off + 8].try_into().unwrap()) as usize
}
#[inline]
/// Returns the value at the given index in the overflow region.
fn data_value(&self, i: usize) -> u32 {
let off = self.data_offset + i * OVERFLOW_ENTRY_SIZE + 8;
u32::from_le_bytes(self.mmap[off..off + 4].try_into().unwrap())
}
#[inline]
pub fn sum(&self) -> u64 {
self.iter().map(|v| v as u64).sum()
let primary = &self.mmap[self.primary_offset..self.primary_offset + self.n];
byte_sum(primary, (0..self.n_overflow).map(|i| self.data_value(i)))
}
#[inline]
pub fn count_nonzero(&self) -> u64 {
self.iter().filter(|&v| v > 0).count() as u64
let primary = &self.mmap[self.primary_offset..self.primary_offset + self.n];
byte_count_nonzero(primary)
}
#[inline]
/// Returns the Bray-Curtis distance between two compact int vectors.
/// Lightweight zero-copy view — primary and overflow point into the mmap.
pub fn view(&self) -> IntSliceView<'_> {
let primary = &self.mmap[self.primary_offset..self.primary_offset + self.n];
let overflow_raw = &self.mmap[self.data_offset..self.data_offset + self.n_overflow * OVERFLOW_ENTRY_SIZE];
IntSliceView::new(primary, overflow_raw, self.n_overflow, self.n)
}
pub fn iter(&self) -> Iter<'_> {
Iter { pciv: self, slot: 0, overflow_pos: 0 }
}
// ── Distance methods ──────────────────────────────────────────────────────
pub fn bray_dist(&self, other: &PersistentCompactIntVec) -> f64 {
let sum_min = self.partial_bray_dist(other);
let denom = self.sum() + other.sum();
if denom == 0 {
return 0.0;
}
1.0 - 2.0 * sum_min as f64 / denom as f64
if denom == 0 { 0.0 } else { 1.0 - 2.0 * sum_min as f64 / denom as f64 }
}
/// Returns `Σ_slot min(self[slot], other[slot])` — the additive numerator of Bray-Curtis.
/// The denominator `sum_a + sum_b` is obtained from `self.sum() + other.sum()`.
pub fn partial_bray_dist(&self, other: &PersistentCompactIntVec) -> u64 {
assert_eq!(self.n, other.len(), "length mismatch");
self.iter()
.zip(other.iter())
.map(|(a, b)| a.min(b) as u64)
.sum()
self.iter().zip(other.iter()).map(|(a, b)| a.min(b) as u64).sum()
}
/// Returns the relative frequency Bray-Curtis distance between two compact int vectors.
///
/// This is a variant of [`bray_dist`] that uses relative frequencies instead of raw counts.
pub fn relfreq_bray_dist(&self, other: &PersistentCompactIntVec) -> f64 {
assert_eq!(self.n, other.len(), "length mismatch");
let sum_a = self.sum() as f64;
let sum_b = other.sum() as f64;
if sum_a == 0.0 && sum_b == 0.0 {
return 0.0;
}
let sum_min = self.partial_relfreq_bray_dist(other, sum_a, sum_b);
1.0 - sum_min
let sa = self.sum() as f64;
let sb = other.sum() as f64;
if sa == 0.0 && sb == 0.0 { return 0.0; }
1.0 - self.partial_relfreq_bray_dist(other, sa, sb)
}
/// Returns the partial relative frequency Bray-Curtis distance between two compact int vectors.
///
/// This is used internally by [`relfreq_bray_dist`] and to easily compute the relative frequency
/// Bray-Curtis distance over a set of vector pairs.
///
/// Arguments:
/// - `other`: the other compact int vector to compare with
/// - `sum_a`: the sum of the first vector's counts
/// - `sum_b`: the sum of the second vector's counts
///
/// Returns the sum of the minimum relative frequencies at each index.
pub fn partial_relfreq_bray_dist(
&self,
other: &PersistentCompactIntVec,
sum_a: f64,
sum_b: f64,
) -> f64 {
pub fn partial_relfreq_bray_dist(&self, other: &PersistentCompactIntVec, sum_a: f64, sum_b: f64) -> f64 {
assert_eq!(self.n, other.len(), "length mismatch");
let sum_min: f64 = self
.iter()
.zip(other.iter())
self.iter().zip(other.iter())
.map(|(a, b)| {
let pa = if sum_a > 0.0 { a as f64 / sum_a } else { 0.0 };
let pb = if sum_b > 0.0 { b as f64 / sum_b } else { 0.0 };
pa.min(pb)
})
.sum();
sum_min
.sum()
}
/// Returns the euclidean distance between two compact int vectors.
pub fn euclidean_dist(&self, other: &PersistentCompactIntVec) -> f64 {
self.partial_euclidean_dist(other).sqrt()
}
/// Returns the partial euclidean distance between two compact int vectors.
///
/// This is used internally by [`euclidean_dist`] and to easily compute the euclidean distance
/// over a set of vector pairs.
///
/// The result is the sum of the squared differences between corresponding elements of the two
/// vectors.
pub fn partial_euclidean_dist(&self, other: &PersistentCompactIntVec) -> f64 {
assert_eq!(self.n, other.len(), "length mismatch");
self.iter()
.zip(other.iter())
.map(|(a, b)| {
let d = a as f64 - b as f64;
d * d
})
self.iter().zip(other.iter())
.map(|(a, b)| { let d = a as f64 - b as f64; d * d })
.sum()
}
/// Returns the relative frequency euclidean distance between two compact int vectors.
///
/// This is a variant of [`euclidean_dist`] that uses relative frequencies instead of raw counts.
pub fn relfreq_euclidean_dist(&self, other: &PersistentCompactIntVec) -> f64 {
assert_eq!(self.n, other.len(), "length mismatch");
let sum_a = self.sum() as f64;
let sum_b = other.sum() as f64;
if sum_a == 0.0 && sum_b == 0.0 {
return 0.0;
}
self.partial_relfreq_euclidean_dist(other, sum_a, sum_b)
.sqrt()
let sa = self.sum() as f64;
let sb = other.sum() as f64;
if sa == 0.0 && sb == 0.0 { return 0.0; }
self.partial_relfreq_euclidean_dist(other, sa, sb).sqrt()
}
/// Returns the partial relative frequency euclidean distance between two compact int vectors.
///
/// This is used internally by [`relfreq_euclidean_dist`] and to easily compute the relative frequency
/// euclidean distance over a set of vector pairs.
pub fn partial_relfreq_euclidean_dist(
&self,
other: &PersistentCompactIntVec,
sum_a: f64,
sum_b: f64,
) -> f64 {
pub fn partial_relfreq_euclidean_dist(&self, other: &PersistentCompactIntVec, sum_a: f64, sum_b: f64) -> f64 {
assert_eq!(self.n, other.len(), "length mismatch");
self.iter()
.zip(other.iter())
self.iter().zip(other.iter())
.map(|(a, b)| {
let pa = if sum_a > 0.0 { a as f64 / sum_a } else { 0.0 };
let pb = if sum_b > 0.0 { b as f64 / sum_b } else { 0.0 };
@@ -266,46 +174,19 @@ impl PersistentCompactIntVec {
.sum()
}
/// Returns the Euclidean distance between two compact int vectors using the Hellinger transform.
///
/// The Hellinger transform is applied to the raw counts of each vector, and the result is
/// the Euclidean distance between the transformed vectors. The Hellinger transform is defined
/// as the square root of the relative frequencies.
pub fn hellinger_euclidean_dist(&self, other: &PersistentCompactIntVec) -> f64 {
assert_eq!(self.n, other.len(), "length mismatch");
let sum_a = self.sum() as f64;
let sum_b = other.sum() as f64;
if sum_a == 0.0 && sum_b == 0.0 {
return 0.0;
}
self.partial_hellinger_euclidean_dist(other, sum_a, sum_b)
.sqrt()
let sa = self.sum() as f64;
let sb = other.sum() as f64;
if sa == 0.0 && sb == 0.0 { return 0.0; }
self.partial_hellinger_euclidean_dist(other, sa, sb).sqrt()
}
/// Returns the partial Hellinger Euclidean distance between two compact int vectors.
///
/// This is used internally by [`hellinger_euclidean_dist`] and to easily compute the Hellinger
/// Euclidean distance over a set of vector pairs.
pub fn partial_hellinger_euclidean_dist(
&self,
other: &PersistentCompactIntVec,
sum_a: f64,
sum_b: f64,
) -> f64 {
pub fn partial_hellinger_euclidean_dist(&self, other: &PersistentCompactIntVec, sum_a: f64, sum_b: f64) -> f64 {
assert_eq!(self.n, other.len(), "length mismatch");
self.iter()
.zip(other.iter())
self.iter().zip(other.iter())
.map(|(a, b)| {
let pa = if sum_a > 0.0 {
(a as f64 / sum_a).sqrt()
} else {
0.0
};
let pb = if sum_b > 0.0 {
(b as f64 / sum_b).sqrt()
} else {
0.0
};
let pa = if sum_a > 0.0 { (a as f64 / sum_a).sqrt() } else { 0.0 };
let pb = if sum_b > 0.0 { (b as f64 / sum_b).sqrt() } else { 0.0 };
let d = pa - pb;
d * d
})
@@ -317,22 +198,13 @@ impl PersistentCompactIntVec {
}
pub fn threshold_jaccard_dist(&self, other: &PersistentCompactIntVec, threshold: u32) -> f64 {
assert_eq!(self.n, other.len(), "length mismatch");
let (intersection, union) = self.partial_threshold_jaccard_dist(other, threshold);
if union == 0 {
return 0.0;
}
1.0 - intersection as f64 / union as f64
if union == 0 { 0.0 } else { 1.0 - intersection as f64 / union as f64 }
}
pub fn partial_threshold_jaccard_dist(
&self,
other: &PersistentCompactIntVec,
threshold: u32,
) -> (u64, u64) {
pub fn partial_threshold_jaccard_dist(&self, other: &PersistentCompactIntVec, threshold: u32) -> (u64, u64) {
assert_eq!(self.n, other.len(), "length mismatch");
self.iter()
.zip(other.iter())
self.iter().zip(other.iter())
.fold((0u64, 0u64), |(inter, uni), (a, b)| {
let ap = a >= threshold;
let bp = b >= threshold;
@@ -343,23 +215,12 @@ impl PersistentCompactIntVec {
pub fn jaccard_dist(&self, other: &PersistentCompactIntVec) -> f64 {
self.threshold_jaccard_dist(other, 1)
}
pub fn iter(&self) -> Iter<'_> {
Iter {
pciv: self,
slot: 0,
overflow_pos: 0,
}
}
}
impl<'a> IntoIterator for &'a PersistentCompactIntVec {
type Item = u32;
type IntoIter = Iter<'a>;
fn into_iter(self) -> Iter<'a> {
self.iter()
}
fn into_iter(self) -> Iter<'a> { self.iter() }
}
pub struct Iter<'a> {
@@ -374,9 +235,7 @@ impl Iterator for Iter<'_> {
type Item = u32;
fn next(&mut self) -> Option<u32> {
if self.slot >= self.pciv.n {
return None;
}
if self.slot >= self.pciv.n { return None; }
let v = self.pciv.mmap[self.pciv.primary_offset + self.slot];
self.slot += 1;
if v < 255 {
src/obicompactvec/src/tempbitvec.rs
+111
View File
@@ -0,0 +1,111 @@
use std::io;
use std::path::Path;
use tempfile::TempDir;
use crate::bitvec::{PersistentBitVec, PersistentBitVecBuilder};
use crate::views::{BitSliceIter, BitSliceView, IntSliceView};
// ── TempBitVec — frozen read-only, auto-deleted on drop ──────────────────────
pub struct TempBitVec {
vec: PersistentBitVec,
// Dropped after `vec` (field order), so the mmap is released before the
// temp directory is deleted.
_temp: TempDir,
}
impl TempBitVec {
pub fn make_persistent(&self, path: &Path) -> io::Result<PersistentBitVec> {
std::fs::copy(self.vec.path(), path)?;
PersistentBitVec::open(path)
}
pub fn len(&self) -> usize {
self.vec.len()
}
pub fn is_empty(&self) -> bool {
self.vec.is_empty()
}
pub fn get(&self, slot: usize) -> bool {
self.vec.get(slot)
}
pub fn count_ones(&self) -> u64 {
self.vec.count_ones()
}
pub fn view(&self) -> BitSliceView<'_> {
self.vec.view()
}
pub fn iter(&self) -> BitSliceIter<'_> {
self.view().iter()
}
}
// ── TempBitVecBuilder — mutable, becomes TempBitVec on freeze ────────────────
pub struct TempBitVecBuilder {
builder: PersistentBitVecBuilder,
temp: TempDir,
}
impl TempBitVecBuilder {
pub fn new(n: usize) -> io::Result<Self> {
let temp = TempDir::new()?;
let path = temp.path().join("data.pbiv");
let builder = PersistentBitVecBuilder::new(n, &path)?;
Ok(Self { builder, temp })
}
pub fn new_ones(n: usize) -> io::Result<Self> {
let temp = TempDir::new()?;
let path = temp.path().join("data.pbiv");
let builder = PersistentBitVecBuilder::new_ones(n, &path)?;
Ok(Self { builder, temp })
}
pub fn freeze(self) -> io::Result<TempBitVec> {
let Self { builder, temp } = self;
let vec = builder.finish()?;
Ok(TempBitVec { vec, _temp: temp })
}
pub fn set(&mut self, slot: usize, value: bool) {
self.builder.set(slot, value);
}
pub fn view(&self) -> BitSliceView<'_> {
self.builder.view()
}
pub fn or(&mut self, other: BitSliceView<'_>) {
self.builder.or(other);
}
pub fn and(&mut self, other: BitSliceView<'_>) {
self.builder.and(other);
}
pub fn xor(&mut self, other: BitSliceView<'_>) {
self.builder.xor(other);
}
pub fn not(&mut self) {
self.builder.not();
}
pub fn copy_from(&mut self, src: BitSliceView<'_>) {
self.builder.copy_from(src);
}
pub fn or_where(&mut self, col: IntSliceView<'_>, pred: impl Fn(u32) -> bool) {
self.builder.or_where(col, pred);
}
pub fn and_where(&mut self, col: IntSliceView<'_>, pred: impl Fn(u32) -> bool) {
self.builder.and_where(col, pred);
}
pub fn xor_where(&mut self, col: IntSliceView<'_>, pred: impl Fn(u32) -> bool) {
self.builder.xor_where(col, pred);
}
}
src/obicompactvec/src/tempintvec.rs
+89
View File
@@ -0,0 +1,89 @@
use std::io;
use std::path::Path;
use tempfile::TempDir;
use crate::builder::PersistentCompactIntVecBuilder;
use crate::reader::PersistentCompactIntVec;
use crate::views::{BitSliceView, IntSliceView};
// ── TempCompactIntVec — frozen read-only, auto-deleted on drop ────────────────
pub struct TempCompactIntVec {
vec: PersistentCompactIntVec,
// Dropped after `vec` (field order), so the mmap is released before the
// temp directory is deleted.
_temp: TempDir,
}
impl TempCompactIntVec {
pub fn make_persistent(&self, path: &Path) -> io::Result<PersistentCompactIntVec> {
std::fs::copy(self.vec.path(), path)?;
PersistentCompactIntVec::open(path)
}
pub fn len(&self) -> usize { self.vec.len() }
pub fn is_empty(&self) -> bool { self.vec.is_empty() }
pub fn get(&self, slot: usize) -> u32 { self.vec.get(slot) }
pub fn sum(&self) -> u64 { self.vec.sum() }
pub fn view(&self) -> IntSliceView<'_> { self.vec.view() }
pub fn iter(&self) -> crate::reader::Iter<'_> { self.vec.iter() }
}
// ── TempCompactIntVecBuilder — mutable, becomes TempCompactIntVec on freeze ──
pub struct TempCompactIntVecBuilder {
builder: PersistentCompactIntVecBuilder,
temp: TempDir,
}
impl TempCompactIntVecBuilder {
pub fn new(n: usize) -> io::Result<Self> {
let temp = TempDir::new()?;
let path = temp.path().join("data.pciv");
let builder = PersistentCompactIntVecBuilder::new(n, &path)?;
Ok(Self { builder, temp })
}
pub fn freeze(self) -> io::Result<TempCompactIntVec> {
let Self { builder, temp } = self;
let vec = builder.finish()?;
Ok(TempCompactIntVec { vec, _temp: temp })
}
pub fn n(&self) -> usize { self.builder.len() }
pub fn set(&mut self, slot: usize, value: u32) { self.builder.set(slot, value); }
pub fn get(&self, slot: usize) -> u32 { self.builder.get(slot) }
pub fn primary_bytes(&self) -> &[u8] { self.builder.primary_bytes() }
pub fn primary_bytes_mut(&mut self) -> &mut [u8] { self.builder.primary_bytes_mut() }
pub fn inc_present(&mut self, col: BitSliceView<'_>) {
self.builder.inc_present(col);
}
pub fn inc_present_fast(&mut self, col: BitSliceView<'_>) {
self.builder.inc_present_fast(col);
}
pub fn inc_predicate(&mut self, col: IntSliceView<'_>, pred: impl Fn(u32) -> bool) {
self.builder.inc_predicate(col, pred);
}
pub fn inc_predicate_fast(&mut self, col: IntSliceView<'_>, pred: impl Fn(u32) -> bool) {
self.builder.inc_predicate_fast(col, pred);
}
pub fn add(&mut self, other: IntSliceView<'_>) {
self.builder.add(other);
}
pub fn mask_with(&mut self, mask: BitSliceView<'_>) {
self.builder.mask_with(mask);
}
pub fn min(&mut self, other: IntSliceView<'_>) { self.builder.min(other); }
pub fn max(&mut self, other: IntSliceView<'_>) { self.builder.max(other); }
pub fn diff(&mut self, other: IntSliceView<'_>) { self.builder.diff(other); }
}
src/obicompactvec/src/tests/bitmatrix.rs
+55 -1
View File
@@ -1,6 +1,6 @@
use tempfile::tempdir;
use crate::{PersistentBitMatrix, PersistentBitMatrixBuilder};
use crate::{pack_bit_matrix, PersistentBitMatrix, PersistentBitMatrixBuilder};
use crate::traits::BitPartials;
fn make_matrix(cols: &[&[bool]]) -> (tempfile::TempDir, PersistentBitMatrix) {
@@ -203,3 +203,57 @@ fn partial_hamming_matches_hamming() {
let full = m.hamming_dist_matrix();
assert_eq!(partial, full);
}
// ── col_view on Packed ────────────────────────────────────────────────────────
#[test]
fn col_view_packed_values() {
let (dir, _) = make_matrix(&[
&[true, false, true, true],
&[false, true, false, true],
]);
pack_bit_matrix(&dir.path().join("presence")).unwrap();
let m = PersistentBitMatrix::open(dir.path()).unwrap();
// col 0: [T, F, T, T]
let v0 = m.col_view(0);
assert_eq!(v0.len(), 4);
assert_eq!(v0.get(0), true);
assert_eq!(v0.get(1), false);
assert_eq!(v0.get(2), true);
assert_eq!(v0.get(3), true);
assert_eq!(v0.count_ones(), 3);
// col 1: [F, T, F, T]
let v1 = m.col_view(1);
assert_eq!(v1.get(0), false);
assert_eq!(v1.get(1), true);
assert_eq!(v1.get(2), false);
assert_eq!(v1.get(3), true);
assert_eq!(v1.count_ones(), 2);
}
#[test]
fn col_view_packed_matches_columnar() {
let data: &[&[bool]] = &[
&[true, false, true, false, true, true, false, true],
&[false, false, true, true, false, true, true, false],
&[true, true, true, false, false, false, true, true],
];
let (dir_col, m_col) = make_matrix(data);
let (dir_pack, _) = make_matrix(data);
pack_bit_matrix(&dir_pack.path().join("presence")).unwrap();
let m_pack = PersistentBitMatrix::open(dir_pack.path()).unwrap();
for c in 0..data.len() {
let col_ref = m_col.col(c);
let col_view = m_pack.col_view(c);
assert_eq!(col_view.len(), col_ref.len(), "col={c} len");
for s in 0..col_ref.len() {
assert_eq!(col_view.get(s), col_ref.get(s), "col={c} slot={s}");
}
assert_eq!(col_view.count_ones(), col_ref.count_ones(), "col={c} count_ones");
assert_eq!(col_view.words(), col_ref.words(), "col={c} words");
}
drop(dir_col);
}
src/obicompactvec/src/tests/bitvec.rs
+3 -3
View File
@@ -77,7 +77,7 @@ fn op_and() {
let dir = tempdir().unwrap();
let path = dir.path().join("out.pbiv");
let mut b = PersistentBitVecBuilder::build_from(&ra, &path).unwrap();
b.and(&rb);
b.and(rb.view());
b.close().unwrap();
let r = PersistentBitVec::open(&path).unwrap();
assert_eq!(r.iter().collect::<Vec<_>>(), vec![true, false, false, false]);
@@ -90,7 +90,7 @@ fn op_or() {
let dir = tempdir().unwrap();
let path = dir.path().join("out.pbiv");
let mut b = PersistentBitVecBuilder::build_from(&ra, &path).unwrap();
b.or(&rb);
b.or(rb.view());
b.close().unwrap();
let r = PersistentBitVec::open(&path).unwrap();
assert_eq!(r.iter().collect::<Vec<_>>(), vec![true, true, true, false]);
@@ -103,7 +103,7 @@ fn op_xor() {
let dir = tempdir().unwrap();
let path = dir.path().join("out.pbiv");
let mut b = PersistentBitVecBuilder::build_from(&ra, &path).unwrap();
b.xor(&rb);
b.xor(rb.view());
b.close().unwrap();
let r = PersistentBitVec::open(&path).unwrap();
assert_eq!(r.iter().collect::<Vec<_>>(), vec![false, true, true, false]);
src/obicompactvec/src/tests/colgroup.rs
+223
View File
@@ -0,0 +1,223 @@
use tempfile::tempdir;
use crate::{
ColGroup, MatrixGroupOps,
PersistentBitMatrix, PersistentBitMatrixBuilder,
PersistentCompactIntMatrix, PersistentCompactIntMatrixBuilder,
};
use crate::{PersistentBitVecBuilder, PersistentCompactIntVec, PersistentCompactIntVecBuilder};
// ── helpers ───────────────────────────────────────────────────────────────────
fn make_int_matrix(cols: &[&[u32]]) -> (tempfile::TempDir, PersistentCompactIntMatrix) {
let n = cols.first().map_or(0, |c| c.len());
let dir = tempdir().unwrap();
let mut b = PersistentCompactIntMatrixBuilder::new(n, &dir.path().join("counts")).unwrap();
for &col in cols {
let mut cb = b.add_col().unwrap();
for (slot, &v) in col.iter().enumerate() { cb.set(slot, v); }
cb.close().unwrap();
}
b.close().unwrap();
let m = PersistentCompactIntMatrix::open(dir.path()).unwrap();
(dir, m)
}
fn make_bit_matrix(cols: &[&[bool]]) -> (tempfile::TempDir, PersistentBitMatrix) {
let n = cols.first().map_or(0, |c| c.len());
let dir = tempdir().unwrap();
let presence = dir.path().join("presence");
let mut b = PersistentBitMatrixBuilder::new(n, &presence).unwrap();
for &col in cols {
let mut cb = b.add_col().unwrap();
for (slot, &v) in col.iter().enumerate() { cb.set(slot, v); }
cb.close().unwrap();
}
b.close().unwrap();
let m = PersistentBitMatrix::open(dir.path()).unwrap();
(dir, m)
}
// ── IntMatrix: partial_group_sum ──────────────────────────────────────────────
#[test]
fn int_partial_group_sum_basic() {
// col0=[1,2,3], col1=[10,20,30], col2=[100,0,5]
// group {0,2}: sum = [101, 2, 8]
let (_d, m) = make_int_matrix(&[&[1, 2, 3], &[10, 20, 30], &[100, 0, 5]]);
let g = ColGroup::new("g", vec![0, 2]);
let result = m.partial_group_sum(&g).unwrap();
assert_eq!(result.get(0), 101);
assert_eq!(result.get(1), 2);
assert_eq!(result.get(2), 8);
}
#[test]
fn int_partial_group_sum_with_overflow() {
// col0=[300,0], col1=[200,400]: group {0,1}: sum=[500, 400]
let (_d, m) = make_int_matrix(&[&[300, 0], &[200, 400]]);
let g = ColGroup::new("g", vec![0, 1]);
let result = m.partial_group_sum(&g).unwrap();
assert_eq!(result.get(0), 500);
assert_eq!(result.get(1), 400);
assert_eq!(result.sum(), 900);
}
// ── IntMatrix: partial_group_presence_count ───────────────────────────────────
#[test]
fn int_partial_group_presence_count() {
// col0=[5,1,0,3], col1=[2,0,4,3], col2=[0,3,1,0]
// threshold=2: col0: [T,F,F,T], col1: [T,F,T,T], col2: [F,T,F,F]
// group {0,1,2}: counts = [2, 1, 1, 2]
let (_d, m) = make_int_matrix(&[&[5, 1, 0, 3], &[2, 0, 4, 3], &[0, 3, 1, 0]]);
let g = ColGroup::new("g", vec![0, 1, 2]);
let result = m.partial_group_presence_count(&g, 2).unwrap();
assert_eq!(result.get(0), 2);
assert_eq!(result.get(1), 1);
assert_eq!(result.get(2), 1);
assert_eq!(result.get(3), 2);
}
#[test]
fn int_partial_group_presence_count_with_overflow() {
// col0=[300,0,10], col1=[0,400,10], col2=[1,1,10]
// threshold=5: col0: [T,F,T], col1: [F,T,T], col2: [F,F,T]
// group {0,1,2}: counts = [1, 1, 3]
let (_d, m) = make_int_matrix(&[&[300, 0, 10], &[0, 400, 10], &[1, 1, 10]]);
let g = ColGroup::new("g", vec![0, 1, 2]);
let result = m.partial_group_presence_count(&g, 5).unwrap();
assert_eq!(result.get(0), 1);
assert_eq!(result.get(1), 1);
assert_eq!(result.get(2), 3);
}
// ── IntMatrix: partial_group_any ──────────────────────────────────────────────
#[test]
fn int_partial_group_any() {
// col0=[0,3,0,1], col1=[2,0,0,0], col2=[0,0,5,0]
// threshold=2: col0: [F,T,F,F], col1: [T,F,F,F], col2: [F,F,T,F]
// group {0,1,2}: any = [T, T, T, F]
let (_d, m) = make_int_matrix(&[&[0, 3, 0, 1], &[2, 0, 0, 0], &[0, 0, 5, 0]]);
let g = ColGroup::new("g", vec![0, 1, 2]);
let result = m.partial_group_any(&g, 2).unwrap();
assert_eq!(result.get(0), true);
assert_eq!(result.get(1), true);
assert_eq!(result.get(2), true);
assert_eq!(result.get(3), false);
}
// ── IntMatrix: mask_with ──────────────────────────────────────────────────────
#[test]
fn mask_with_zeros_selected_slots() {
// count vec [10, 20, 30, 40], mask [T, F, T, F] → [10, 0, 30, 0]
let dir = tempdir().unwrap();
let mut v = PersistentCompactIntVecBuilder::new(4, &dir.path().join("v.pciv")).unwrap();
v.set(0, 10); v.set(1, 20); v.set(2, 30); v.set(3, 40);
let mut mask = PersistentBitVecBuilder::new(4, &dir.path().join("m.pbiv")).unwrap();
mask.set(0, true); mask.set(2, true);
v.mask_with(mask.view());
v.close().unwrap();
let r = PersistentCompactIntVec::open(&dir.path().join("v.pciv")).unwrap();
assert_eq!(r.get(0), 10);
assert_eq!(r.get(1), 0);
assert_eq!(r.get(2), 30);
assert_eq!(r.get(3), 0);
}
#[test]
fn mask_with_overflow_slot_zeroed() {
// overflow slot (value 500) masked out → removed from overflow, primary=0
let dir = tempdir().unwrap();
let mut v = PersistentCompactIntVecBuilder::new(3, &dir.path().join("v.pciv")).unwrap();
v.set(0, 10); v.set(1, 500); v.set(2, 5);
let mut mask = PersistentBitVecBuilder::new(3, &dir.path().join("m.pbiv")).unwrap();
mask.set(0, true); mask.set(2, true); // slot 1 masked out
v.mask_with(mask.view());
v.close().unwrap();
let r = PersistentCompactIntVec::open(&dir.path().join("v.pciv")).unwrap();
assert_eq!(r.get(0), 10);
assert_eq!(r.get(1), 0);
assert_eq!(r.get(2), 5);
let ov: Vec<_> = r.view().overflow_entries().collect();
assert!(ov.is_empty(), "overflow entry for masked-out slot should be gone");
}
#[test]
fn mask_with_all_ones_is_noop() {
let dir = tempdir().unwrap();
let mut v = PersistentCompactIntVecBuilder::new(4, &dir.path().join("v.pciv")).unwrap();
v.set(0, 300); v.set(1, 1); v.set(2, 0); v.set(3, 42);
let mask = PersistentBitVecBuilder::new_ones(4, &dir.path().join("m.pbiv")).unwrap();
v.mask_with(mask.view());
v.close().unwrap();
let r = PersistentCompactIntVec::open(&dir.path().join("v.pciv")).unwrap();
assert_eq!(r.get(0), 300);
assert_eq!(r.get(1), 1);
assert_eq!(r.get(2), 0);
assert_eq!(r.get(3), 42);
}
// ── BitMatrix: partial_group_presence_count ───────────────────────────────────
#[test]
fn bit_partial_group_presence_count() {
// col0=[T,F,T,F], col1=[T,T,F,F], col2=[F,T,T,F]
// group {0,1,2}: counts = [2, 2, 2, 0]
let (_d, m) = make_bit_matrix(&[
&[true, false, true, false],
&[true, true, false, false],
&[false,true, true, false],
]);
let g = ColGroup::new("g", vec![0, 1, 2]);
let result = m.partial_group_presence_count(&g, 1).unwrap();
assert_eq!(result.get(0), 2);
assert_eq!(result.get(1), 2);
assert_eq!(result.get(2), 2);
assert_eq!(result.get(3), 0);
}
// ── BitMatrix: partial_group_any ──────────────────────────────────────────────
#[test]
fn bit_partial_group_any() {
// col0=[T,F,F], col1=[F,F,T], group {0,1}: any = [T, F, T]
let (_d, m) = make_bit_matrix(&[
&[true, false, false],
&[false, false, true],
]);
let g = ColGroup::new("g", vec![0, 1]);
let result = m.partial_group_any(&g, 1).unwrap();
assert_eq!(result.get(0), true);
assert_eq!(result.get(1), false);
assert_eq!(result.get(2), true);
}
// ── Composition: partial results are additive ─────────────────────────────────
#[test]
fn int_presence_count_additive_across_split() {
// Simulate two partitions (different kmer ranges) whose counts should add.
// Global data for col0: [5,1,0,3,2], col1: [2,0,4,3,1] — threshold=2
// Split: partition A = slots 0..2, partition B = slots 2..5
let data_a: &[&[u32]] = &[&[5, 1], &[2, 0]];
let data_b: &[&[u32]] = &[&[0, 3, 2], &[4, 3, 1]];
let (_da, ma) = make_int_matrix(data_a);
let (_db, mb) = make_int_matrix(data_b);
let g = ColGroup::new("g", vec![0, 1]);
let pa = ma.partial_group_presence_count(&g, 2).unwrap();
let pb = mb.partial_group_presence_count(&g, 2).unwrap();
// Concatenate by adding (disjoint kmer ranges — here we just verify
// individual results match the expected per-partition counts).
// partition A: col0=[5≥2,1<2]=[T,F], col1=[2≥2,0<2]=[T,F] → [2, 0]
assert_eq!(pa.get(0), 2);
assert_eq!(pa.get(1), 0);
// partition B: col0=[0<2,3≥2,2≥2]=[F,T,T], col1=[4≥2,3≥2,1<2]=[T,T,F] → [1, 2, 1]
assert_eq!(pb.get(0), 1);
assert_eq!(pb.get(1), 2);
assert_eq!(pb.get(2), 1);
}
src/obicompactvec/src/tests/intmatrix.rs
+56 -1
View File
@@ -1,6 +1,6 @@
use tempfile::tempdir;
use crate::{PersistentCompactIntMatrix, PersistentCompactIntMatrixBuilder};
use crate::{pack_compact_int_matrix, PersistentCompactIntMatrix, PersistentCompactIntMatrixBuilder};
use crate::traits::CountPartials;
fn make_matrix(cols: &[&[u32]]) -> (tempfile::TempDir, PersistentCompactIntMatrix) {
@@ -243,6 +243,61 @@ fn partial_hellinger_matches_full() {
}
}
#[test]
fn col_view_packed_values() {
// Build Columnar with overflow values (≥ 255), pack, reopen as Packed, exercise col_view().
let (dir, _col) = make_matrix(&[&[10, 300, 500], &[200, 50, 1000]]);
pack_compact_int_matrix(&dir.path().join("counts")).unwrap();
let m = PersistentCompactIntMatrix::open(dir.path()).unwrap();
// col 0: [10, 300, 500] — two overflow slots
let v0 = m.col_view(0);
assert_eq!(v0.get(0), 10);
assert_eq!(v0.get(1), 300);
assert_eq!(v0.get(2), 500);
assert_eq!(v0.sum(), 810);
assert_eq!(v0.count_nonzero(), 3);
let mut ov0: Vec<(usize, u32)> = v0.overflow_entries().collect();
ov0.sort_unstable_by_key(|&(s, _)| s);
assert_eq!(ov0, vec![(1, 300), (2, 500)]);
// col 1: [200, 50, 1000] — one overflow slot
let v1 = m.col_view(1);
assert_eq!(v1.get(0), 200);
assert_eq!(v1.get(1), 50);
assert_eq!(v1.get(2), 1000);
let mut ov1: Vec<(usize, u32)> = v1.overflow_entries().collect();
ov1.sort_unstable_by_key(|&(s, _)| s);
assert_eq!(ov1, vec![(2, 1000)]);
}
#[test]
fn col_view_packed_matches_columnar() {
// Same data, compare col_view() on Packed against col() on Columnar slot-by-slot.
let data: &[&[u32]] = &[&[0, 255, 1, 300, 128], &[500, 3, 0, 700, 42]];
let (dir_col, m_col) = make_matrix(data);
// Re-build in a separate dir so we can pack without touching m_col's files.
let (dir_pack, _) = make_matrix(data);
pack_compact_int_matrix(&dir_pack.path().join("counts")).unwrap();
let m_pack = PersistentCompactIntMatrix::open(dir_pack.path()).unwrap();
for c in 0..data.len() {
let col_ref = m_col.col(c);
let col_view = m_pack.col_view(c);
assert_eq!(col_view.len(), col_ref.len());
for s in 0..col_ref.len() {
assert_eq!(col_view.get(s), col_ref.get(s), "col={c} slot={s}");
}
assert_eq!(col_view.sum(), col_ref.sum(), "col={c} sum");
let mut ov_view: Vec<(usize, u32)> = col_view.overflow_entries().collect();
let mut ov_ref: Vec<(usize, u32)> = col_ref.view().overflow_entries().collect();
ov_view.sort_unstable_by_key(|&(s, _)| s);
ov_ref.sort_unstable_by_key(|&(s, _)| s);
assert_eq!(ov_view, ov_ref, "col={c} overflow_entries");
}
drop(dir_col);
}
#[test]
fn partial_relfreq_bray_additive_across_split() {
// Split rows [1,2,3,4,5] between two matrices; partial sums should add up.
src/obicompactvec/src/tests/mod.rs
+5 -4
View File
@@ -1,5 +1,6 @@
mod bitmatrix;
mod bitvec;
mod colgroup;
mod intmatrix;
use tempfile::tempdir;
@@ -169,7 +170,7 @@ fn combine_min() {
let dir = tempdir().unwrap();
let path = dir.path().join("out.pciv");
let mut b = PersistentCompactIntVecBuilder::build_from(&ra, &path).unwrap();
b.min(&rb);
b.min(rb.view());
b.close().unwrap();
let r = PersistentCompactIntVec::open(&path).unwrap();
assert_eq!(r.iter().collect::<Vec<_>>(), vec![10, 100, 0, 800]);
@@ -182,7 +183,7 @@ fn combine_max() {
let dir = tempdir().unwrap();
let path = dir.path().join("out.pciv");
let mut b = PersistentCompactIntVecBuilder::build_from(&ra, &path).unwrap();
b.max(&rb);
b.max(rb.view());
b.close().unwrap();
let r = PersistentCompactIntVec::open(&path).unwrap();
assert_eq!(r.iter().collect::<Vec<_>>(), vec![20, 300, 500, 1000]);
@@ -195,7 +196,7 @@ fn combine_add() {
let dir = tempdir().unwrap();
let path = dir.path().join("out.pciv");
let mut b = PersistentCompactIntVecBuilder::build_from(&ra, &path).unwrap();
b.add(&rb);
b.add(rb.view());
b.close().unwrap();
let r = PersistentCompactIntVec::open(&path).unwrap();
assert_eq!(r.iter().collect::<Vec<_>>(), vec![30, 300, 5, 101]);
@@ -220,7 +221,7 @@ fn combine_diff() {
let dir = tempdir().unwrap();
let path = dir.path().join("out.pciv");
let mut b = PersistentCompactIntVecBuilder::build_from(&ra, &path).unwrap();
b.diff(&rb);
b.diff(rb.view());
b.close().unwrap();
let r = PersistentCompactIntVec::open(&path).unwrap();
assert_eq!(r.iter().collect::<Vec<_>>(), vec![10, 700, 0, 0]);
src/obicompactvec/src/traits.rs
+1 -1
View File
@@ -1,6 +1,6 @@
use ndarray::{Array1, Array2};
/// Column-level weight statistic — total count or presence count per column.
// ── Column-level weight statistic — total count or presence count per column.
/// Additive across layers and partitions; used as denominator in normalised distances.
///
/// `partial_kmer_counts` returns the number of **distinct k-mers** present per
src/obicompactvec/src/views.rs
+278
View File
@@ -0,0 +1,278 @@
use crate::format::{byte_count_nonzero, byte_sum, parse_overflow_entry};
// ── BitSliceView ──────────────────────────────────────────────────────────────
/// Lightweight, copy-able read-only view over a u64 word array.
/// Bit `i` is in `words[i >> 6]` at position `i & 63`. Padding bits are zero.
#[derive(Clone, Copy)]
pub struct BitSliceView<'a> {
pub(crate) words: &'a [u64],
pub(crate) n: usize,
}
impl<'a> BitSliceView<'a> {
#[inline]
pub fn new(words: &'a [u64], n: usize) -> Self { Self { words, n } }
pub fn len(&self) -> usize { self.n }
pub fn is_empty(&self) -> bool { self.n == 0 }
pub fn words(&self) -> &'a [u64] { self.words }
#[inline]
pub fn get(&self, slot: usize) -> bool {
(self.words[slot >> 6] >> (slot & 63)) & 1 != 0
}
pub fn count_ones(&self) -> u64 {
self.words.iter().map(|w| w.count_ones() as u64).sum()
}
pub fn count_zeros(&self) -> u64 { self.n as u64 - self.count_ones() }
pub fn iter(&self) -> BitSliceIter<'a> {
BitSliceIter { words: self.words, slot: 0, n: self.n }
}
pub fn partial_jaccard_dist(self, other: BitSliceView<'_>) -> (u64, u64) {
assert_eq!(self.n, other.n, "BitSliceView length mismatch");
self.words.iter().zip(other.words)
.fold((0u64, 0u64), |(i, u), (&a, &b)| {
(i + (a & b).count_ones() as u64, u + (a | b).count_ones() as u64)
})
}
pub fn jaccard_dist(self, other: BitSliceView<'_>) -> f64 {
let (inter, union) = self.partial_jaccard_dist(other);
if union == 0 { 0.0 } else { 1.0 - inter as f64 / union as f64 }
}
pub fn hamming_dist(self, other: BitSliceView<'_>) -> u64 {
assert_eq!(self.n, other.n, "BitSliceView length mismatch");
self.words.iter().zip(other.words)
.map(|(&a, &b)| (a ^ b).count_ones() as u64)
.sum()
}
}
// ── BitSliceIter ──────────────────────────────────────────────────────────────
pub struct BitSliceIter<'a> {
words: &'a [u64],
slot: usize,
n: usize,
}
impl Iterator for BitSliceIter<'_> {
type Item = bool;
fn next(&mut self) -> Option<bool> {
if self.slot >= self.n { return None; }
let v = (self.words[self.slot >> 6] >> (self.slot & 63)) & 1 != 0;
self.slot += 1;
Some(v)
}
fn size_hint(&self) -> (usize, Option<usize>) {
let rem = self.n - self.slot;
(rem, Some(rem))
}
}
impl ExactSizeIterator for BitSliceIter<'_> {}
// ── IntSliceView ──────────────────────────────────────────────────────────────
/// Lightweight, copy-able read-only view over a compact-int primary array plus
/// its sorted raw overflow bytes. Zero-copy: all data lives in the caller's mmap.
#[derive(Clone, Copy)]
pub struct IntSliceView<'a> {
pub(crate) primary: &'a [u8],
pub(crate) overflow_raw: &'a [u8], // n_overflow × OVERFLOW_ENTRY_SIZE bytes, sorted by slot
pub(crate) n_overflow: usize,
pub(crate) n: usize,
}
impl<'a> IntSliceView<'a> {
#[inline]
pub fn new(primary: &'a [u8], overflow_raw: &'a [u8], n_overflow: usize, n: usize) -> Self {
Self { primary, overflow_raw, n_overflow, n }
}
pub fn len(&self) -> usize { self.n }
pub fn is_empty(&self) -> bool { self.n == 0 }
pub fn primary_bytes(&self) -> &'a [u8] { self.primary }
pub fn n_overflow(&self) -> usize { self.n_overflow }
pub fn overflow_entries(&self) -> impl Iterator<Item = (usize, u32)> + 'a {
let raw = self.overflow_raw;
let n_ov = self.n_overflow;
(0..n_ov).map(move |i| parse_overflow_entry(raw, 0, i))
}
/// O(log n_overflow) via binary search (overflow is always sorted by slot).
pub fn get(&self, slot: usize) -> u32 {
let b = self.primary[slot];
if b < 255 { return b as u32; }
let mut lo = 0usize;
let mut hi = self.n_overflow;
while lo < hi {
let mid = lo + (hi - lo) / 2;
let (s, v) = parse_overflow_entry(self.overflow_raw, 0, mid);
match s.cmp(&slot) {
std::cmp::Ordering::Equal => return v,
std::cmp::Ordering::Less => lo = mid + 1,
std::cmp::Ordering::Greater => hi = mid,
}
}
panic!("slot {slot} marked overflow but not found")
}
/// Sequential merge scan: yields all n values in slot order.
pub fn iter(&self) -> IntSliceViewIter<'a> {
IntSliceViewIter {
primary: self.primary,
overflow_raw: self.overflow_raw,
slot: 0,
overflow_pos: 0,
n: self.n,
}
}
pub fn sum(&self) -> u64 {
byte_sum(self.primary, self.overflow_entries().map(|(_, v)| v))
}
pub fn count_nonzero(&self) -> u64 {
byte_count_nonzero(self.primary)
}
// ── Distance methods ──────────────────────────────────────────────────────
pub fn partial_bray_dist(self, other: IntSliceView<'_>) -> u64 {
assert_eq!(self.n, other.n, "length mismatch");
self.iter().zip(other.iter()).map(|(a, b)| a.min(b) as u64).sum()
}
pub fn bray_dist(self, other: IntSliceView<'_>) -> f64 {
let sum_min = self.partial_bray_dist(other);
let denom = self.sum() + other.sum();
if denom == 0 { 0.0 } else { 1.0 - 2.0 * sum_min as f64 / denom as f64 }
}
pub fn partial_relfreq_bray_dist(self, other: IntSliceView<'_>, sa: f64, sb: f64) -> f64 {
assert_eq!(self.n, other.n, "length mismatch");
self.iter().zip(other.iter())
.map(|(a, b)| {
let pa = if sa > 0.0 { a as f64 / sa } else { 0.0 };
let pb = if sb > 0.0 { b as f64 / sb } else { 0.0 };
pa.min(pb)
})
.sum()
}
pub fn relfreq_bray_dist(self, other: IntSliceView<'_>) -> f64 {
let sa = self.sum() as f64;
let sb = other.sum() as f64;
if sa == 0.0 && sb == 0.0 { return 0.0; }
1.0 - self.partial_relfreq_bray_dist(other, sa, sb)
}
pub fn partial_euclidean_dist(self, other: IntSliceView<'_>) -> f64 {
assert_eq!(self.n, other.n, "length mismatch");
self.iter().zip(other.iter())
.map(|(a, b)| { let d = a as f64 - b as f64; d * d })
.sum()
}
pub fn euclidean_dist(self, other: IntSliceView<'_>) -> f64 {
self.partial_euclidean_dist(other).sqrt()
}
pub fn partial_relfreq_euclidean_dist(self, other: IntSliceView<'_>, sa: f64, sb: f64) -> f64 {
assert_eq!(self.n, other.n, "length mismatch");
self.iter().zip(other.iter())
.map(|(a, b)| {
let pa = if sa > 0.0 { a as f64 / sa } else { 0.0 };
let pb = if sb > 0.0 { b as f64 / sb } else { 0.0 };
let d = pa - pb;
d * d
})
.sum()
}
pub fn relfreq_euclidean_dist(self, other: IntSliceView<'_>) -> f64 {
let sa = self.sum() as f64;
let sb = other.sum() as f64;
if sa == 0.0 && sb == 0.0 { return 0.0; }
self.partial_relfreq_euclidean_dist(other, sa, sb).sqrt()
}
pub fn partial_hellinger_euclidean_dist(self, other: IntSliceView<'_>, sa: f64, sb: f64) -> f64 {
assert_eq!(self.n, other.n, "length mismatch");
self.iter().zip(other.iter())
.map(|(a, b)| {
let pa = if sa > 0.0 { (a as f64 / sa).sqrt() } else { 0.0 };
let pb = if sb > 0.0 { (b as f64 / sb).sqrt() } else { 0.0 };
let d = pa - pb;
d * d
})
.sum()
}
pub fn hellinger_euclidean_dist(self, other: IntSliceView<'_>) -> f64 {
let sa = self.sum() as f64;
let sb = other.sum() as f64;
if sa == 0.0 && sb == 0.0 { return 0.0; }
self.partial_hellinger_euclidean_dist(other, sa, sb).sqrt()
}
pub fn hellinger_dist(self, other: IntSliceView<'_>) -> f64 {
self.hellinger_euclidean_dist(other) / std::f64::consts::SQRT_2
}
pub fn partial_threshold_jaccard_dist(self, other: IntSliceView<'_>, threshold: u32) -> (u64, u64) {
assert_eq!(self.n, other.n, "length mismatch");
self.iter().zip(other.iter())
.fold((0u64, 0u64), |(inter, uni), (a, b)| {
let ap = a >= threshold;
let bp = b >= threshold;
(inter + (ap & bp) as u64, uni + (ap | bp) as u64)
})
}
pub fn threshold_jaccard_dist(self, other: IntSliceView<'_>, threshold: u32) -> f64 {
let (inter, union) = self.partial_threshold_jaccard_dist(other, threshold);
if union == 0 { 0.0 } else { 1.0 - inter as f64 / union as f64 }
}
pub fn jaccard_dist(self, other: IntSliceView<'_>) -> f64 {
self.threshold_jaccard_dist(other, 1)
}
}
// ── IntSliceViewIter ──────────────────────────────────────────────────────────
pub struct IntSliceViewIter<'a> {
primary: &'a [u8],
overflow_raw: &'a [u8],
slot: usize,
overflow_pos: usize,
n: usize,
}
impl Iterator for IntSliceViewIter<'_> {
type Item = u32;
fn next(&mut self) -> Option<u32> {
if self.slot >= self.n { return None; }
let v = self.primary[self.slot];
self.slot += 1;
if v < 255 {
Some(v as u32)
} else {
let (_, val) = parse_overflow_entry(self.overflow_raw, 0, self.overflow_pos);
self.overflow_pos += 1;
Some(val)
}
}
fn size_hint(&self) -> (usize, Option<usize>) {
let rem = self.n - self.slot;
(rem, Some(rem))
}
}
impl ExactSizeIterator for IntSliceViewIter<'_> {}
src/obidebruinj/src/debruijn.rs
+1
View File
@@ -3,6 +3,7 @@ use crossbeam_channel;
use hashbrown::HashMap;
use obikseq::k;
use obikseq::{CanonicalKmer, Sequence, Unitig};
#[cfg(not(any(test, feature = "test-utils")))]
use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
use std::cell::RefCell;
use std::fmt;
src/obikindex/src/index.rs
+5 -2
View File
@@ -204,6 +204,7 @@ impl KmerIndex {
let n = self.n_partitions();
let order: Vec<usize> = (0..n).collect();
let pb = progress_bar("pack", n as u64, "partitions");
crate::numa::PartitionRunner::new().run(
&order,
|i| -> OKIResult<()> {
@@ -220,8 +221,10 @@ impl KmerIndex {
}
Ok(())
},
|_, _, _| {},
)
|_, _, _| { pb.inc(1); },
)?;
pb.finish_and_clear();
Ok(())
}
/// Write a `layer_meta.json` in any layer directory that is missing one.
src/obikindex/src/merge.rs
+3 -1
View File
@@ -11,7 +11,7 @@ use obilayeredmap::IndexMode;
use crate::error::{OKIError, OKIResult};
use crate::index::KmerIndex;
use crate::meta::{GenomeInfo, IndexMeta};
use crate::state::IndexState;
use crate::state::{IndexState, SENTINEL_INDEXED};
pub use obikpartitionner::MergeMode;
@@ -263,6 +263,8 @@ impl KmerIndex {
rep.push(t.stop());
}
fs::File::create(output.join(SENTINEL_INDEXED)).map_err(OKIError::Io)?;
KmerIndex::open(output)
}
}
src/obikindex/src/rebuild.rs
+2
View File
@@ -98,7 +98,9 @@ impl KmerIndex {
fs::File::create(output.join(SENTINEL_INDEXED))?;
let idx = KmerIndex::open(output)?;
let t_pack = Stage::start("pack");
idx.pack_matrices()?;
rep.push(t_pack.stop());
Ok(idx)
}
}
src/obikindex/src/select.rs
+9 -7
View File
@@ -3,7 +3,7 @@ use std::io;
use std::path::Path;
use obikpartitionner::{KmerPartition, OutputCol, PARTITIONS_SUBDIR};
use obisys::{Stage, progress_bar};
use obisys::{Reporter, Stage, progress_bar};
use tracing::info;
use crate::error::{OKIError, OKIResult};
@@ -25,6 +25,7 @@ impl KmerIndex {
threshold: u32,
output_presence: bool,
force: bool,
rep: &mut Reporter,
) -> OKIResult<Self> {
let output = output.as_ref();
@@ -80,13 +81,14 @@ impl KmerIndex {
).map_err(OKIError::Partition)?;
pb.finish_and_clear();
let _ = t.stop();
rep.push(t.stop());
fs::File::create(output.join(SENTINEL_INDEXED))?;
let idx = KmerIndex::open(output)?;
let t_pack = Stage::start("pack");
idx.pack_matrices()?;
rep.push(t_pack.stop());
Ok(idx)
}
@@ -98,6 +100,7 @@ impl KmerIndex {
specs: &[OutputCol],
threshold: u32,
output_presence: bool,
rep: &mut Reporter,
) -> OKIResult<()> {
if self.state() != IndexState::Indexed {
return Err(OKIError::NotIndexed(self.root_path.clone()));
@@ -106,7 +109,6 @@ impl KmerIndex {
let n_src_genomes = self.meta.genomes.len();
let n_partitions = self.partition.n_partitions();
// Open a second handle to the same path so we can borrow src and dst simultaneously.
let src_partition = KmerPartition::open_with_config(
&self.root_path,
self.meta.config.kmer_size,
@@ -132,17 +134,17 @@ impl KmerIndex {
).map_err(OKIError::Partition)?;
pb.finish_and_clear();
rep.push(t.stop());
let _ = t.stop();
// Update index.meta with new genome list and with_counts flag.
self.meta.config.with_counts = !output_presence;
self.meta.genomes = specs.iter()
.map(|s| GenomeInfo::new(s.label.clone()))
.collect();
self.meta.write(&self.root_path)?;
let t_pack = Stage::start("pack");
self.pack_matrices()?;
rep.push(t_pack.stop());
Ok(())
}
}
src/obikmer/Cargo.toml
+2 -1
View File
@@ -1,6 +1,6 @@
[package]
name = "obikmer"
version = "0.1.0"
version = "0.1.3"
edition = "2024"
[[bin]]
@@ -19,6 +19,7 @@ obikpartitionner = { path = "../obikpartitionner" }
obisys = { path = "../obisys" }
obiskio = { path = "../obiskio" }
obikindex = { path = "../obikindex" }
obitaxonomy = { path = "../obitaxonomy" }
obilayeredmap = { path = "../obilayeredmap" }
clap = { version = "4", features = ["derive"] }
serde_json = "1"
src/obikmer/src/cmd/predicate.rs
+8 -11
View File
@@ -3,6 +3,7 @@ use std::collections::HashMap;
use clap::Args;
use obikindex::GenomeInfo;
use obikpartitionner::{GroupQuorumFilter, KmerFilter};
use obitaxonomy::{TaxPath, TaxPattern};
// ── Operator ──────────────────────────────────────────────────────────────────
@@ -49,7 +50,6 @@ impl MetaPred {
if values.iter().any(|v| v.is_empty()) {
return Err(format!("empty value in predicate: {s}"));
}
Ok(Self { key, op, values })
}
@@ -70,18 +70,15 @@ impl MetaPred {
// ── Path matching ─────────────────────────────────────────────────────────────
/// True if `value` is equal to `pattern` or is a descendant of it in a `/`-separated hierarchy.
/// True if the stored taxonomy `value` matches `pattern`.
///
/// - Absolute pattern (`/a/b`): `value` must start with `/a/b` at a segment boundary.
/// - Bare segment (`b`): `value` must contain `b` as an exact segment anywhere.
/// `value` must be a valid `TaxPath` (starts with `taxonomy:/`).
/// `pattern` is a `TaxPattern` query (see `obitaxonomy::TaxPattern` for syntax).
/// Returns `false` if either fails to parse.
fn path_matches(value: &str, pattern: &str) -> bool {
if pattern.starts_with('/') {
value == pattern
|| (value.starts_with(pattern)
&& value[pattern.len()..].starts_with('/'))
} else {
value.split('/').any(|seg| seg == pattern)
}
let Ok(path) = TaxPath::parse(value) else { return false };
let Ok(pat) = TaxPattern::parse(pattern) else { return false };
pat.matches(&path)
}
// ── Three-value group evaluation ──────────────────────────────────────────────
src/obikmer/src/cmd/select.rs
+7 -2
View File
@@ -4,6 +4,7 @@ use std::path::PathBuf;
use clap::{Args, ValueEnum};
use obikindex::{GenomeInfo, KmerIndex};
use obikpartitionner::{AggOp, OutputCol};
use obisys::Reporter;
use tracing::info;
use super::predicate::matching_genome_indices;
@@ -229,20 +230,24 @@ pub fn run(args: SelectArgs) {
if output_presence { "presence" } else { "count" },
);
let mut rep = Reporter::new();
if args.in_place {
src.select_in_place(&specs, args.presence_threshold, output_presence)
src.select_in_place(&specs, args.presence_threshold, output_presence, &mut rep)
.unwrap_or_else(|e| {
eprintln!("select error: {e}");
std::process::exit(1);
});
rep.print();
info!("selected in-place → {}", args.source.display());
} else {
let output = args.output.unwrap();
KmerIndex::select(&output, &src, &specs, args.presence_threshold, output_presence, args.force)
KmerIndex::select(&output, &src, &specs, args.presence_threshold, output_presence, args.force, &mut rep)
.unwrap_or_else(|e| {
eprintln!("select error: {e}");
std::process::exit(1);
});
rep.print();
info!("selected index → {}", output.display());
}
}
src/obikmer/src/main.rs
+1 -1
View File
@@ -6,7 +6,7 @@ use clap::{Parser, Subcommand};
use tracing_subscriber::{EnvFilter, fmt};
#[derive(Parser)]
#[command(name = "obikmer", about = "DNA k-mer tools")]
#[command(name = "obikmer", about = "DNA k-mer tools", version)]
struct Cli {
#[command(subcommand)]
command: Commands,
src/obikpartitionner/src/filter.rs
+120
View File
@@ -1,9 +1,24 @@
use obicompactvec::FilterMask;
/// Trait for kmer row filters.
///
/// `row` contains raw per-genome counts (or 0/1 for presence/absence data).
/// `n_genomes` equals `row.len()`.
pub trait KmerFilter: Send + Sync {
fn passes(&self, row: &[u32], n_genomes: usize) -> bool;
/// Express this filter as a [`FilterMask`] column-operation expression.
///
/// Returns `Some(expr)` if the filter can be evaluated solely from matrix
/// column aggregates (no per-kmer row scan needed). Returns `None` if the
/// filter requires row-level inspection.
///
/// `threshold` semantics in the returned mask use `>= threshold`, matching
/// [`obicompactvec::MatrixGroupOps`]. Implementations must add 1 to any
/// row-level threshold that uses strict `>` comparison.
fn column_mask_expr(&self, _n_genomes: usize) -> Option<FilterMask> {
None
}
}
/// True when `row` passes every filter in `filters`.
@@ -29,6 +44,16 @@ impl KmerFilter for MinGenomeFraction {
let p = present_count(row, self.threshold);
p as f64 / n_genomes as f64 >= self.frac
}
fn column_mask_expr(&self, n_genomes: usize) -> Option<FilterMask> {
let t = self.threshold.checked_add(1)?;
let min_count = (self.frac * n_genomes as f64).ceil() as usize;
Some(FilterMask::PresenceGeq {
indices: (0..n_genomes).collect(),
threshold: t,
min_count,
})
}
}
/// At most `frac` fraction of genomes contain this kmer (count > `threshold`).
@@ -42,6 +67,16 @@ impl KmerFilter for MaxGenomeFraction {
let p = present_count(row, self.threshold);
p as f64 / n_genomes as f64 <= self.frac
}
fn column_mask_expr(&self, n_genomes: usize) -> Option<FilterMask> {
let t = self.threshold.checked_add(1)?;
let max_count = (self.frac * n_genomes as f64).floor() as usize;
Some(FilterMask::PresenceLeq {
indices: (0..n_genomes).collect(),
threshold: t,
max_count,
})
}
}
/// At least `count` genomes contain this kmer (count > `threshold`).
@@ -54,6 +89,15 @@ impl KmerFilter for MinGenomeCount {
fn passes(&self, row: &[u32], _n_genomes: usize) -> bool {
present_count(row, self.threshold) >= self.count
}
fn column_mask_expr(&self, n_genomes: usize) -> Option<FilterMask> {
let t = self.threshold.checked_add(1)?;
Some(FilterMask::PresenceGeq {
indices: (0..n_genomes).collect(),
threshold: t,
min_count: self.count,
})
}
}
/// At most `count` genomes contain this kmer (count > `threshold`).
@@ -66,6 +110,15 @@ impl KmerFilter for MaxGenomeCount {
fn passes(&self, row: &[u32], _n_genomes: usize) -> bool {
present_count(row, self.threshold) <= self.count
}
fn column_mask_expr(&self, n_genomes: usize) -> Option<FilterMask> {
let t = self.threshold.checked_add(1)?;
Some(FilterMask::PresenceLeq {
indices: (0..n_genomes).collect(),
threshold: t,
max_count: self.count,
})
}
}
// ── Total-count filters (count indexes only) ───────────────────────────────────
@@ -79,6 +132,13 @@ impl KmerFilter for MinTotalCount {
fn passes(&self, row: &[u32], _n_genomes: usize) -> bool {
row.iter().sum::<u32>() >= self.total
}
fn column_mask_expr(&self, n_genomes: usize) -> Option<FilterMask> {
Some(FilterMask::SumGeq {
indices: (0..n_genomes).collect(),
min_sum: self.total,
})
}
}
/// Sum of counts across all genomes <= `total`.
@@ -90,6 +150,13 @@ impl KmerFilter for MaxTotalCount {
fn passes(&self, row: &[u32], _n_genomes: usize) -> bool {
row.iter().sum::<u32>() <= self.total
}
fn column_mask_expr(&self, n_genomes: usize) -> Option<FilterMask> {
Some(FilterMask::SumLeq {
indices: (0..n_genomes).collect(),
max_sum: self.total,
})
}
}
// ── Group-based quorum filter ─────────────────────────────────────────────────
@@ -113,6 +180,37 @@ pub struct GroupQuorumFilter {
pub max_outgroup_frac: f64,
}
impl GroupQuorumFilter {
// Build PresenceGeq/PresenceLeq constraints for one group (ingroup or outgroup).
fn group_mask_parts(
indices: &[usize],
threshold: u32,
min_count: usize,
max_count: usize,
min_frac: f64,
max_frac: f64,
parts: &mut Vec<FilterMask>,
) {
let n = indices.len();
let geq = min_count.max((min_frac * n as f64).ceil() as usize);
if geq > 0 {
parts.push(FilterMask::PresenceGeq {
indices: indices.to_vec(),
threshold,
min_count: geq,
});
}
let leq = max_count.min((max_frac * n as f64).floor() as usize);
if leq < n {
parts.push(FilterMask::PresenceLeq {
indices: indices.to_vec(),
threshold,
max_count: leq,
});
}
}
}
impl KmerFilter for GroupQuorumFilter {
fn passes(&self, row: &[u32], _n_genomes: usize) -> bool {
if !self.ingroup_idx.is_empty() {
@@ -139,4 +237,26 @@ impl KmerFilter for GroupQuorumFilter {
}
true
}
fn column_mask_expr(&self, _n_genomes: usize) -> Option<FilterMask> {
let t = self.threshold.checked_add(1)?;
let mut parts: Vec<FilterMask> = Vec::new();
if !self.ingroup_idx.is_empty() {
Self::group_mask_parts(
&self.ingroup_idx, t,
self.min_count, self.max_count,
self.min_frac, self.max_frac,
&mut parts,
);
}
if !self.outgroup_idx.is_empty() {
Self::group_mask_parts(
&self.outgroup_idx, t,
self.min_outgroup_count, self.max_outgroup_count,
self.min_outgroup_frac, self.max_outgroup_frac,
&mut parts,
);
}
Some(FilterMask::And(parts))
}
}
src/obikpartitionner/src/merge_layer.rs
+36
View File
@@ -10,6 +10,7 @@ use obipipeline::{
};
use obicompactvec::{
MatrixGroupOps,
PersistentBitMatrix, PersistentBitMatrixBuilder, PersistentBitVecBuilder,
PersistentCompactIntMatrix, PersistentCompactIntMatrixBuilder, PersistentCompactIntVecBuilder,
};
@@ -78,6 +79,41 @@ impl SrcLayerData {
}
buf
}
pub(crate) fn n_slots(&self) -> usize {
match self {
SrcLayerData::Presence(_, mat) => mat.n(),
SrcLayerData::Count(_, mat) => mat.n(),
}
}
/// MPHF lookup: returns the slot index for `kmer` (kmer must be in the domain).
#[inline]
pub(crate) fn slot(&self, kmer: CanonicalKmer) -> usize {
match self {
SrcLayerData::Presence(mphf, _) => mphf.index(kmer),
SrcLayerData::Count(mphf, _) => mphf.index(kmer),
}
}
/// Row lookup by slot index, bypassing the MPHF.
#[inline]
pub(crate) fn fill_row_by_slot(&self, slot: usize, n_genomes: usize) -> Vec<u32> {
let mut buf = vec![0u32; n_genomes];
match self {
SrcLayerData::Presence(_, mat) => mat.fill_row(slot, &mut buf),
SrcLayerData::Count(_, mat) => mat.fill_row(slot, &mut buf),
}
buf
}
/// Call `f` with a reference to the underlying matrix as `&dyn MatrixGroupOps`.
pub(crate) fn with_matrix<R>(&self, f: impl FnOnce(&dyn MatrixGroupOps) -> R) -> R {
match self {
SrcLayerData::Presence(_, mat) => f(mat),
SrcLayerData::Count(_, mat) => f(mat),
}
}
}
// ── helpers ───────────────────────────────────────────────────────────────────
src/obikpartitionner/src/rebuild_layer.rs
+143 -51
View File
@@ -1,8 +1,9 @@
use std::path::Path;
use obicompactvec::{
PersistentBitMatrixBuilder, PersistentBitVecBuilder, PersistentCompactIntMatrixBuilder,
PersistentCompactIntVecBuilder,
FilterMask, eval_filter_mask,
PersistentBitMatrixBuilder, PersistentBitVecBuilder,
PersistentCompactIntMatrixBuilder, PersistentCompactIntVecBuilder,
};
use obidebruinj::GraphDeBruijn;
use obikseq::CanonicalKmer;
@@ -10,18 +11,135 @@ use obilayeredmap::meta::PartitionMeta;
use obilayeredmap::{IndexMode, MphfLayer};
use obiskio::{SKError, SKResult, UnitigFileReader};
use crate::common::{ColBuilder, col_path_bit, col_path_int, load_meta, olm_to_sk, write_matrix_meta};
use crate::filter::{KmerFilter, passes_all};
use crate::common::{load_meta, olm_to_sk};
use crate::filter::KmerFilter;
use crate::graph_pipeline::materialize_layer;
use crate::merge_layer::{MergeMode, SrcLayerData};
use crate::partition::KmerPartition;
const INDEX_SUBDIR: &str = "index";
/// Iterate all kmers in `src_index_dir` that pass `filters`, yielding `(kmer, row)`.
// ── Builders — pair matrix builder + column builders for one mode ─────────────
enum Builders {
Presence(PersistentBitMatrixBuilder, Vec<PersistentBitVecBuilder>),
Count(PersistentCompactIntMatrixBuilder, Vec<PersistentCompactIntVecBuilder>),
}
impl Builders {
fn new(mode: MergeMode, n: usize, dir: &Path, n_genomes: usize) -> SKResult<Self> {
match mode {
MergeMode::Presence => {
let mut mat = PersistentBitMatrixBuilder::new(n, dir).map_err(SKError::Io)?;
let mut cols = Vec::with_capacity(n_genomes);
for _ in 0..n_genomes { cols.push(mat.add_col().map_err(SKError::Io)?); }
Ok(Builders::Presence(mat, cols))
}
MergeMode::Count => {
let mut mat = PersistentCompactIntMatrixBuilder::new(n, dir).map_err(SKError::Io)?;
let mut cols = Vec::with_capacity(n_genomes);
for _ in 0..n_genomes { cols.push(mat.add_col().map_err(SKError::Io)?); }
Ok(Builders::Count(mat, cols))
}
}
}
fn set_val(&mut self, col: usize, slot: usize, value: u32) {
match self {
Builders::Presence(_, cols) => cols[col].set(slot, value > 0),
Builders::Count(_, cols) => cols[col].set(slot, value),
}
}
fn close(self) -> SKResult<()> {
match self {
Builders::Presence(mat, cols) => {
for b in cols { b.close().map_err(SKError::Io)?; }
mat.close().map_err(SKError::Io)
}
Builders::Count(mat, cols) => {
for b in cols { b.close().map_err(SKError::Io)?; }
mat.close().map_err(SKError::Io)
}
}
}
}
// ── try_compute_combined_mask ─────────────────────────────────────────────────
/// Build a per-slot `TempBitVec` mask from `filters` using column operations
/// on the source matrix — no per-kmer MPHF lookup or row read needed.
///
/// Uses [`SrcLayerData`] semantics: counts take priority over presence when
/// `mode = Count`; presence (or implicit all-ones) is used for `Presence`.
/// Returns `Some(mask)` when every filter in `filters` can express itself as
/// a [`FilterMask`] expression. Returns `None` when any filter requires
/// row-level inspection (fall back to `passes_all`).
fn try_compute_combined_mask(
filters: &[Box<dyn KmerFilter>],
src_data: &SrcLayerData,
n_genomes: usize,
) -> SKResult<Option<obicompactvec::TempBitVec>> {
if filters.is_empty() {
return Ok(None);
}
let mut exprs: Vec<FilterMask> = Vec::with_capacity(filters.len());
for f in filters {
match f.column_mask_expr(n_genomes) {
Some(expr) => exprs.push(expr),
None => return Ok(None),
}
}
let combined = FilterMask::And(exprs);
let n = src_data.n_slots();
let mask = src_data
.with_matrix(|mat| eval_filter_mask(&combined, mat, n))
.map_err(SKError::Io)?;
Ok(Some(mask))
}
// ── iter_src_kmers_masked (pass 1) ────────────────────────────────────────────
/// Iterate all passing kmers in `src_index_dir`, yielding only the kmer value.
///
/// When all filters can be expressed as column operations, a per-slot mask is
/// computed once per layer and used for O(1) slot-check per kmer instead of a
/// full row read. Falls back to row-level `passes_all` otherwise.
fn iter_src_kmers_masked(
src_index_dir: &Path,
mode: MergeMode,
n_genomes: usize,
filters: &[Box<dyn KmerFilter>],
mut cb: impl FnMut(CanonicalKmer),
) -> SKResult<()> {
let src_meta = load_meta(src_index_dir, "rebuild")?;
for l in 0..src_meta.n_layers {
let src_layer_dir = src_index_dir.join(format!("layer_{l}"));
let unitigs_path = src_layer_dir.join("unitigs.bin");
if !unitigs_path.exists() { continue; }
let src_data = SrcLayerData::open(&src_layer_dir, mode)?;
let mask = try_compute_combined_mask(filters, &src_data, n_genomes)?;
let reader = UnitigFileReader::open_sequential(&unitigs_path)?;
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
let slot = src_data.slot(kmer);
let passes = match &mask {
Some(m) => m.get(slot),
None => {
let row = src_data.fill_row_by_slot(slot, n_genomes);
filters.iter().all(|f| f.passes(&row, n_genomes))
}
};
if passes { cb(kmer); }
}
}
Ok(())
}
// ── iter_src_layers (pass 2) ──────────────────────────────────────────────────
/// Iterate all passing kmers in `src_index_dir`, yielding `(kmer, row)`.
///
/// When the slot mask is available, skips the row read for filtered-out slots.
fn iter_src_layers(
src_index_dir: &Path,
mode: MergeMode,
@@ -33,17 +151,23 @@ fn iter_src_layers(
for l in 0..src_meta.n_layers {
let src_layer_dir = src_index_dir.join(format!("layer_{l}"));
let unitigs_path = src_layer_dir.join("unitigs.bin");
if !unitigs_path.exists() {
continue;
}
if !unitigs_path.exists() { continue; }
let reader = UnitigFileReader::open_sequential(&unitigs_path)?;
let src_data = SrcLayerData::open(&src_layer_dir, mode)?;
let mask = try_compute_combined_mask(filters, &src_data, n_genomes)?;
let reader = UnitigFileReader::open_sequential(&unitigs_path)?;
for (kmer, _, _) in reader.iter_indexed_canonical_kmers() {
let row = src_data.lookup(kmer, n_genomes);
if passes_all(filters, &row, n_genomes) {
let slot = src_data.slot(kmer);
if let Some(ref m) = mask {
if !m.get(slot) { continue; }
let row = src_data.fill_row_by_slot(slot, n_genomes);
cb(kmer, row.into_boxed_slice());
} else {
let row = src_data.fill_row_by_slot(slot, n_genomes);
if filters.iter().all(|f| f.passes(&row, n_genomes)) {
cb(kmer, row.into_boxed_slice());
}
}
}
}
@@ -81,7 +205,7 @@ impl KmerPartition {
// ── Pass 1: collect filtered kmers into de Bruijn graph ───────────────
let mut g = GraphDeBruijn::new();
iter_src_layers(&src_index_dir, mode, n_genomes, filters, |kmer, _row| {
iter_src_kmers_masked(&src_index_dir, mode, n_genomes, filters, |kmer| {
g.push(kmer);
})?;
@@ -100,54 +224,22 @@ impl KmerPartition {
// ── Prepare matrix builders (one column per genome) ───────────────────
let data_dir = match mode {
MergeMode::Presence => dst_layer_dir.join("presence"),
MergeMode::Count => dst_layer_dir.join("counts"),
MergeMode::Count => dst_layer_dir.join("counts"),
};
std::fs::create_dir_all(&data_dir)?;
let mut builders: Vec<ColBuilder> = match mode {
MergeMode::Presence => {
PersistentBitMatrixBuilder::new(n_new, &data_dir)
.map_err(SKError::Io)?
.close()
.map_err(SKError::Io)?;
(0..n_genomes)
.map(|g| -> SKResult<ColBuilder> {
let b = PersistentBitVecBuilder::new(n_new, &col_path_bit(&data_dir, g))?;
Ok(ColBuilder::Bit(b))
})
.collect::<SKResult<_>>()?
}
MergeMode::Count => {
PersistentCompactIntMatrixBuilder::new(n_new, &data_dir)
.map_err(SKError::Io)?
.close()
.map_err(SKError::Io)?;
(0..n_genomes)
.map(|g| -> SKResult<ColBuilder> {
let b = PersistentCompactIntVecBuilder::new(
n_new,
&col_path_int(&data_dir, g),
)?;
Ok(ColBuilder::Int(b))
})
.collect::<SKResult<_>>()?
}
};
let mut builders = Builders::new(mode, n_new, &data_dir, n_genomes)?;
// ── Pass 2: fill builders ─────────────────────────────────────────────
iter_src_layers(&src_index_dir, mode, n_genomes, filters, |kmer, row| {
if let Some(slot) = dst_mphf.find(kmer) {
for (col, &value) in row.iter().enumerate() {
builders[col].set_val(slot, value);
builders.set_val(col, slot, value);
}
}
})?;
// ── Close builders, write metadata ────────────────────────────────────
for b in builders {
b.close()?;
}
write_matrix_meta(&data_dir, n_new, n_genomes).map_err(SKError::Io)?;
// ── Close builders and write metadata ─────────────────────────────────
builders.close()?;
PartitionMeta {
n_layers: 1,
src/obikpartitionner/src/select_layer.rs
+64 -106
View File
@@ -3,8 +3,9 @@ use std::io;
use std::path::{Path, PathBuf};
use obicompactvec::{
PersistentBitMatrix, PersistentBitMatrixBuilder, PersistentBitVecBuilder,
PersistentCompactIntMatrix, PersistentCompactIntMatrixBuilder, PersistentCompactIntVecBuilder,
ColGroup, MatrixGroupOps,
PersistentBitMatrix, PersistentBitMatrixBuilder,
PersistentCompactIntMatrix, PersistentCompactIntMatrixBuilder,
};
use obilayeredmap::meta::PartitionMeta;
use obilayeredmap::OLMError;
@@ -40,52 +41,6 @@ pub struct OutputCol {
pub op: AggOp,
}
// ── Aggregation ───────────────────────────────────────────────────────────────
#[inline]
fn aggregate(op: AggOp, indices: &[usize], src_row: &[u32], threshold: u32) -> u32 {
match op {
AggOp::Any => {
if indices.iter().any(|&i| src_row[i] > threshold) { 1 } else { 0 }
}
AggOp::All => {
if indices.is_empty() { return 0; }
if indices.iter().all(|&i| src_row[i] > threshold) { 1 } else { 0 }
}
AggOp::None => {
if indices.iter().all(|&i| src_row[i] <= threshold) { 1 } else { 0 }
}
AggOp::Sum => {
indices.iter().map(|&i| src_row[i]).fold(0u32, |a, b| a.saturating_add(b))
}
AggOp::Min => indices.iter().map(|&i| src_row[i]).min().unwrap_or(0),
AggOp::Max => indices.iter().map(|&i| src_row[i]).max().unwrap_or(0),
}
}
// ── ColBuilder ────────────────────────────────────────────────────────────────
enum ColBuilder {
Bit(PersistentBitVecBuilder),
Int(PersistentCompactIntVecBuilder),
}
impl ColBuilder {
fn set_val(&mut self, slot: usize, value: u32) {
match self {
ColBuilder::Bit(b) => b.set(slot, value > 0),
ColBuilder::Int(b) => b.set(slot, value),
}
}
fn close(self) -> SKResult<()> {
match self {
ColBuilder::Bit(b) => b.close().map_err(SKError::Io),
ColBuilder::Int(b) => b.close().map_err(SKError::Io),
}
}
}
// ── Helpers ───────────────────────────────────────────────────────────────────
fn olm_to_sk(e: OLMError) -> SKError {
@@ -95,21 +50,6 @@ fn olm_to_sk(e: OLMError) -> SKError {
}
}
fn col_path_bit(dir: &Path, col: usize) -> PathBuf {
dir.join(format!("col_{col:06}.pbiv"))
}
fn col_path_int(dir: &Path, col: usize) -> PathBuf {
dir.join(format!("col_{col:06}.pciv"))
}
fn write_matrix_meta(dir: &Path, n: usize, n_cols: usize) -> io::Result<()> {
fs::write(
dir.join("meta.json"),
format!("{{\"n\":{n},\"n_cols\":{n_cols}}}\n"),
)
}
/// Copy all plain files (not subdirectories) from `src_dir` to `dst_dir`.
fn copy_layer_files(src_dir: &Path, dst_dir: &Path) -> io::Result<()> {
for entry in fs::read_dir(src_dir)? {
@@ -125,30 +65,64 @@ fn copy_layer_files(src_dir: &Path, dst_dir: &Path) -> io::Result<()> {
// ── fill_builders ─────────────────────────────────────────────────────────────
fn fill_builders(
builders: &mut [ColBuilder],
specs: &[OutputCol],
n: usize,
n_src: usize,
src_layer_dir: &Path,
src_is_count: bool,
threshold: u32,
output_presence: bool,
mut dst_bit: Option<&mut PersistentBitMatrixBuilder>,
mut dst_int: Option<&mut PersistentCompactIntMatrixBuilder>,
) -> SKResult<()> {
let mut src_buf = vec![0u32; n_src];
if src_is_count {
let mat = PersistentCompactIntMatrix::open(src_layer_dir).map_err(SKError::Io)?;
for slot in 0..n {
mat.fill_row(slot, &mut src_buf);
for (col, spec) in specs.iter().enumerate() {
builders[col].set_val(slot, aggregate(spec.op, &spec.indices, &src_buf, threshold));
for spec in specs {
let g = ColGroup::new(&spec.label, spec.indices.clone());
if output_presence {
let b = dst_bit.as_deref_mut().unwrap();
match spec.op {
AggOp::Any => b.add_col_from (&mat.partial_group_any (&g, threshold).map_err(SKError::Io)?),
AggOp::All => b.add_col_from (&mat.partial_group_all (&g, threshold).map_err(SKError::Io)?),
AggOp::None => b.add_col_from (&mat.partial_group_none(&g, threshold).map_err(SKError::Io)?),
AggOp::Sum => b.add_col_from_int(&mat.partial_group_sum (&g).map_err(SKError::Io)?),
AggOp::Min => b.add_col_from_int(&mat.partial_group_min (&g).map_err(SKError::Io)?),
AggOp::Max => b.add_col_from_int(&mat.partial_group_max (&g).map_err(SKError::Io)?),
}.map_err(SKError::Io)?;
} else {
let b = dst_int.as_deref_mut().unwrap();
match spec.op {
AggOp::Sum => b.add_col_from (&mat.partial_group_sum (&g).map_err(SKError::Io)?),
AggOp::Min => b.add_col_from (&mat.partial_group_min (&g).map_err(SKError::Io)?),
AggOp::Max => b.add_col_from (&mat.partial_group_max (&g).map_err(SKError::Io)?),
AggOp::Any => b.add_col_from_bit(&mat.partial_group_any (&g, threshold).map_err(SKError::Io)?),
AggOp::All => b.add_col_from_bit(&mat.partial_group_all (&g, threshold).map_err(SKError::Io)?),
AggOp::None => b.add_col_from_bit(&mat.partial_group_none(&g, threshold).map_err(SKError::Io)?),
}.map_err(SKError::Io)?;
}
}
} else {
let mat = PersistentBitMatrix::open(src_layer_dir).map_err(SKError::Io)?;
for slot in 0..n {
mat.fill_row(slot, &mut src_buf);
for (col, spec) in specs.iter().enumerate() {
builders[col].set_val(slot, aggregate(spec.op, &spec.indices, &src_buf, threshold));
for spec in specs {
let g = ColGroup::new(&spec.label, spec.indices.clone());
if output_presence {
let b = dst_bit.as_deref_mut().unwrap();
match spec.op {
AggOp::Any => b.add_col_from (&mat.partial_group_any (&g, 1).map_err(SKError::Io)?),
AggOp::All => b.add_col_from (&mat.partial_group_all (&g, 1).map_err(SKError::Io)?),
AggOp::None => b.add_col_from (&mat.partial_group_none(&g, 1).map_err(SKError::Io)?),
AggOp::Sum => b.add_col_from_int(&mat.partial_group_sum (&g).map_err(SKError::Io)?),
AggOp::Min => b.add_col_from_int(&mat.partial_group_min (&g).map_err(SKError::Io)?),
AggOp::Max => b.add_col_from_int(&mat.partial_group_max (&g).map_err(SKError::Io)?),
}.map_err(SKError::Io)?;
} else {
let b = dst_int.as_deref_mut().unwrap();
match spec.op {
AggOp::Sum => b.add_col_from (&mat.partial_group_sum (&g).map_err(SKError::Io)?),
AggOp::Min => b.add_col_from (&mat.partial_group_min (&g).map_err(SKError::Io)?),
AggOp::Max => b.add_col_from (&mat.partial_group_max (&g).map_err(SKError::Io)?),
AggOp::Any => b.add_col_from_bit(&mat.partial_group_any (&g, 1).map_err(SKError::Io)?),
AggOp::All => b.add_col_from_bit(&mat.partial_group_all (&g, 1).map_err(SKError::Io)?),
AggOp::None => b.add_col_from_bit(&mat.partial_group_none(&g, 1).map_err(SKError::Io)?),
}.map_err(SKError::Io)?;
}
}
}
@@ -168,7 +142,7 @@ impl KmerPartition {
src: &KmerPartition,
i: usize,
specs: &[OutputCol],
n_src_genomes: usize,
_n_src_genomes: usize,
threshold: u32,
output_presence: bool,
in_place: bool,
@@ -188,7 +162,6 @@ impl KmerPartition {
fs::create_dir_all(&dst_index_dir)?;
}
let n_out = specs.len();
let data_subdir = if output_presence { "presence" } else { "counts" };
for l in 0..src_meta.n_layers {
@@ -201,7 +174,7 @@ impl KmerPartition {
let presence_dir = src_layer_dir.join("presence");
let src_is_count = counts_dir.exists() && !presence_dir.exists();
// Determine number of slots from the source matrix.
// Determine number of slots and detect implicit layers.
let n = if counts_dir.exists() {
PersistentCompactIntMatrix::open(&src_layer_dir).map_err(SKError::Io)?.n()
} else if presence_dir.exists() {
@@ -216,7 +189,7 @@ impl KmerPartition {
};
// Choose the output data directory (temp name for in-place).
let (dst_data_dir, final_data_dir) = if in_place {
let (dst_data_dir, final_data_dir): (PathBuf, PathBuf) = if in_place {
let tmp = dst_layer_dir.join(format!("{data_subdir}_new"));
let perm = dst_layer_dir.join(data_subdir);
(tmp, perm)
@@ -231,37 +204,22 @@ impl KmerPartition {
}
fs::create_dir_all(&dst_data_dir)?;
// Initialise packed-format skeleton.
if output_presence {
PersistentBitMatrixBuilder::new(n, &dst_data_dir)
.map_err(SKError::Io)?.close().map_err(SKError::Io)?;
let (mut dst_bit, mut dst_int) = if output_presence {
(Some(PersistentBitMatrixBuilder::new(n, &dst_data_dir).map_err(SKError::Io)?), None)
} else {
PersistentCompactIntMatrixBuilder::new(n, &dst_data_dir)
.map_err(SKError::Io)?.close().map_err(SKError::Io)?;
}
// Create column builders.
let mut builders: Vec<ColBuilder> = (0..n_out)
.map(|col| -> SKResult<ColBuilder> {
if output_presence {
Ok(ColBuilder::Bit(PersistentBitVecBuilder::new(
n, &col_path_bit(&dst_data_dir, col),
)?))
} else {
Ok(ColBuilder::Int(PersistentCompactIntVecBuilder::new(
n, &col_path_int(&dst_data_dir, col),
)?))
}
})
.collect::<SKResult<_>>()?;
(None, Some(PersistentCompactIntMatrixBuilder::new(n, &dst_data_dir).map_err(SKError::Io)?))
};
fill_builders(
&mut builders, specs, n, n_src_genomes,
&src_layer_dir, src_is_count, threshold,
specs, &src_layer_dir, src_is_count, threshold, output_presence,
dst_bit.as_mut(), dst_int.as_mut(),
)?;
for b in builders { b.close()?; }
write_matrix_meta(&dst_data_dir, n, n_out).map_err(SKError::Io)?;
if output_presence {
dst_bit.unwrap().close().map_err(SKError::Io)?;
} else {
dst_int.unwrap().close().map_err(SKError::Io)?;
}
// In-place: swap old data dir for new.
if in_place {
src/obilayeredmap/src/layer.rs
+1 -5
View File
@@ -106,11 +106,7 @@ impl Layer<()> {
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)?;
let mut col = mb.add_col().map_err(OLMError::Io)?;
for slot in 0..n_kmers {
col.set(slot, true);
}
col.close().map_err(OLMError::Io)?;
mb.add_col_ones().map_err(OLMError::Io)?.close().map_err(OLMError::Io)?;
mb.close().map_err(OLMError::Io)
}
}
src/obitaxonomy/Cargo.toml
+6
View File
@@ -0,0 +1,6 @@
[package]
name = "obitaxonomy"
version = "0.1.0"
edition = "2024"
[dependencies]
src/obitaxonomy/src/error.rs
+38
View File
@@ -0,0 +1,38 @@
use std::fmt;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum TaxError {
/// Stored value does not start with the `taxonomy:/` prefix.
MissingPrefix,
/// Stored path contains no segments after the prefix.
EmptyPath,
/// Query pattern contains no segments (after stripping anchors).
EmptyPattern,
/// A segment has an empty name (e.g. consecutive `/`).
EmptySegmentName,
/// A segment has a trailing `@` with no rank name.
EmptyRankName { segment: String },
/// A segment contains more than one `@`.
AmbiguousRank { segment: String },
}
impl fmt::Display for TaxError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
TaxError::MissingPrefix =>
write!(f, "taxonomy path must start with \"taxonomy:/\""),
TaxError::EmptyPath =>
write!(f, "taxonomy path has no segments"),
TaxError::EmptyPattern =>
write!(f, "taxonomy query pattern has no segments"),
TaxError::EmptySegmentName =>
write!(f, "segment has an empty name"),
TaxError::EmptyRankName { segment } =>
write!(f, "segment has '@' with no rank name: {segment:?}"),
TaxError::AmbiguousRank { segment } =>
write!(f, "segment contains more than one '@': {segment:?}"),
}
}
}
impl std::error::Error for TaxError {}
src/obitaxonomy/src/lib.rs
+11
View File
@@ -0,0 +1,11 @@
mod error;
mod segment;
mod segment_pattern;
mod path;
mod pattern;
pub use error::TaxError;
pub use segment::TaxSegment;
pub use segment_pattern::SegmentPattern;
pub use path::{TaxPath, PREFIX};
pub use pattern::TaxPattern;
src/obitaxonomy/src/path.rs
+82
View File
@@ -0,0 +1,82 @@
use std::fmt;
use std::str::FromStr;
use crate::error::TaxError;
use crate::segment::TaxSegment;
/// The prefix that marks a metadata value as a taxonomy path.
pub const PREFIX: &str = "taxonomy:/";
/// A rooted, `/`-separated taxonomy path with optional per-segment rank annotations.
///
/// Stored form: `taxonomy:/seg1@rank1/seg2/seg3@rank3`
/// The leading `taxonomy:/` is the discriminator; the remainder is one or more
/// `/`-separated segments, each of the form `name` or `name@rank`.
///
/// `@` is reserved and may not appear in segment names or rank names.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TaxPath {
segments: Vec<TaxSegment>,
}
impl TaxPath {
pub fn parse(s: &str) -> Result<Self, TaxError> {
let tail = s.strip_prefix(PREFIX).ok_or(TaxError::MissingPrefix)?;
if tail.is_empty() {
return Err(TaxError::EmptyPath);
}
let segments = tail.split('/')
.map(TaxSegment::parse)
.collect::<Result<Vec<_>, _>>()?;
Ok(Self { segments })
}
/// True if `self` is an ancestor of — or equal to — `other`.
///
/// Comparison is by segment name only; rank annotations are ignored.
/// `self` must be a prefix of `other` at segment granularity.
pub fn is_ancestor_of(&self, other: &TaxPath) -> bool {
self.segments.len() <= other.segments.len()
&& self.segments.iter().zip(other.segments.iter())
.all(|(a, b)| a.name() == b.name())
}
/// Returns the name of the first segment whose rank equals `rank`, if any.
pub fn name_at_rank(&self, rank: &str) -> Option<&str> {
self.segments.iter()
.find(|s| s.rank() == Some(rank))
.map(|s| s.name())
}
/// True if any segment has the given rank.
pub fn has_rank(&self, rank: &str) -> bool {
self.segments.iter().any(|s| s.rank() == Some(rank))
}
/// True if the path contains a segment with both the given rank and name.
pub fn matches_rank(&self, rank: &str, name: &str) -> bool {
self.segments.iter().any(|s| s.rank() == Some(rank) && s.name() == name)
}
pub fn segments(&self) -> &[TaxSegment] { &self.segments }
pub fn depth(&self) -> usize { self.segments.len() }
pub fn is_empty(&self) -> bool { self.segments.is_empty() }
}
impl fmt::Display for TaxPath {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{}", PREFIX)?;
let mut first = true;
for seg in &self.segments {
if !first { write!(f, "/")?; }
write!(f, "{seg}")?;
first = false;
}
Ok(())
}
}
impl FromStr for TaxPath {
type Err = TaxError;
fn from_str(s: &str) -> Result<Self, Self::Err> { Self::parse(s) }
}
src/obitaxonomy/src/pattern.rs
+72
View File
@@ -0,0 +1,72 @@
use crate::error::TaxError;
use crate::path::TaxPath;
use crate::segment::TaxSegment;
use crate::segment_pattern::SegmentPattern;
/// A query pattern for matching against stored `TaxPath` values.
///
/// Syntax:
///
/// | Form | Semantics |
/// |----------|-----------|
/// | `A/B` | A then B as a contiguous sub-path, anywhere in the value |
/// | `/A/B` | value starts with A then B (start-anchored) |
/// | `A/B$` | value ends with A then B (end-anchored) |
/// | `/A/B$` | value is exactly A then B (fully anchored) |
/// | `A@x/B` | A with rank `x`, followed by B with any rank |
///
/// A segment pattern without `@` matches any segment with that name regardless
/// of its stored rank.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TaxPattern {
start_anchored: bool,
end_anchored: bool,
segments: Vec<SegmentPattern>,
}
impl TaxPattern {
pub fn parse(s: &str) -> Result<Self, TaxError> {
let s = s.trim();
let start_anchored = s.starts_with('/');
let s = if start_anchored { &s[1..] } else { s };
let end_anchored = s.ends_with('$');
let s = if end_anchored { &s[..s.len() - 1] } else { s };
if s.is_empty() {
return Err(TaxError::EmptyPattern);
}
let segments = s.split('/')
.map(SegmentPattern::parse)
.collect::<Result<Vec<_>, _>>()?;
Ok(Self { start_anchored, end_anchored, segments })
}
/// True if this pattern matches `path` according to the anchor flags.
///
/// The pattern must match a contiguous run of segments in the path.
/// Start/end anchors restrict where that run may begin or end.
pub fn matches(&self, path: &TaxPath) -> bool {
let n = self.segments.len();
let m = path.depth();
if n > m { return false; }
let segs = path.segments();
match (self.start_anchored, self.end_anchored) {
(true, true) => n == m && self.window_matches(segs, 0),
(true, false) => self.window_matches(segs, 0),
(false, true) => self.window_matches(segs, m - n),
(false, false) => (0..=(m - n)).any(|i| self.window_matches(segs, i)),
}
}
fn window_matches(&self, segs: &[TaxSegment], start: usize) -> bool {
self.segments.iter()
.zip(segs[start..start + self.segments.len()].iter())
.all(|(pat, seg)| pat.matches(seg))
}
}
src/obitaxonomy/src/segment.rs
+49
View File
@@ -0,0 +1,49 @@
use std::fmt;
use crate::error::TaxError;
/// A single node in a taxonomy path: a name and an optional rank.
///
/// Neither `name` nor `rank` may contain `@` (reserved separator).
/// Serialised form: `name` or `name@rank`.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TaxSegment {
name: String,
rank: Option<String>,
}
impl TaxSegment {
pub fn parse(raw: &str) -> Result<Self, TaxError> {
let parts: Vec<&str> = raw.splitn(3, '@').collect();
let (name_raw, rank_raw) = match parts.as_slice() {
[name] => (*name, None),
[name, rank] => (*name, Some(*rank)),
_ => return Err(TaxError::AmbiguousRank { segment: raw.to_string() }),
};
if name_raw.is_empty() {
return Err(TaxError::EmptySegmentName);
}
let rank = match rank_raw {
None => None,
Some("") => return Err(TaxError::EmptyRankName { segment: raw.to_string() }),
Some(r) => Some(r.to_string()),
};
Ok(Self { name: name_raw.to_string(), rank })
}
pub fn name(&self) -> &str { &self.name }
pub fn rank(&self) -> Option<&str> { self.rank.as_deref() }
}
impl fmt::Display for TaxSegment {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match &self.rank {
None => write!(f, "{}", self.name),
Some(r) => write!(f, "{}@{}", self.name, r),
}
}
}
src/obitaxonomy/src/segment_pattern.rs
+41
View File
@@ -0,0 +1,41 @@
use crate::error::TaxError;
use crate::segment::TaxSegment;
/// A single segment in a query pattern: a required name and an optional rank filter.
///
/// If `rank` is `None`, the pattern matches any segment with the given name,
/// regardless of its stored rank. If `rank` is `Some(r)`, both name and rank
/// must match exactly.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SegmentPattern {
name: String,
rank: Option<String>,
}
impl SegmentPattern {
pub fn parse(raw: &str) -> Result<Self, TaxError> {
let parts: Vec<&str> = raw.splitn(3, '@').collect();
let (name_raw, rank_raw) = match parts.as_slice() {
[name] => (*name, None),
[name, rank] => (*name, Some(*rank)),
_ => return Err(TaxError::AmbiguousRank { segment: raw.to_string() }),
};
if name_raw.is_empty() {
return Err(TaxError::EmptySegmentName);
}
let rank = match rank_raw {
None => None,
Some("") => return Err(TaxError::EmptyRankName { segment: raw.to_string() }),
Some(r) => Some(r.to_string()),
};
Ok(Self { name: name_raw.to_string(), rank })
}
/// True if this pattern matches `seg`.
/// Name must match exactly. If a rank is specified in the pattern, the
/// segment's rank must match; otherwise any rank (or no rank) is accepted.
pub fn matches(&self, seg: &TaxSegment) -> bool {
self.name == seg.name()
&& self.rank.as_deref().map_or(true, |r| seg.rank() == Some(r))
}
}