Eric Coissac
ac38aa759b
Renames `get_batch`, `fill_batch`, and `fill_batch_sorted` to `collect_slots_values`, `fill_slots_values`, and `fill_slots_values_sorted` to align with updated `common_traits 0.13` APIs. Introduces optimized batch retrieval that sorts input indices for sequential mmap access before reordering outputs to match the original query order. Updates test suites to reflect the new method signatures without altering validation logic or coverage.
462 lines
16 KiB
Rust
462 lines
16 KiB
Rust
use tempfile::tempdir;
|
|
|
|
use crate::{pack_compact_int_matrix, PersistentCompactIntMatrix, PersistentCompactIntMatrixBuilder, PersistentCompactIntVec, PersistentCompactIntVecBuilder};
|
|
use crate::traits::CountPartials;
|
|
|
|
fn make_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)
|
|
}
|
|
|
|
#[test]
|
|
fn single_col_roundtrip() {
|
|
let dir = tempdir().unwrap();
|
|
let mut b = PersistentCompactIntMatrixBuilder::new(4, &dir.path().join("counts")).unwrap();
|
|
let mut col = b.add_col().unwrap();
|
|
col.set(0, 10);
|
|
col.set(1, 200);
|
|
col.set(2, 300);
|
|
col.set(3, 1000);
|
|
col.close().unwrap();
|
|
b.close().unwrap();
|
|
|
|
let m = PersistentCompactIntMatrix::open(dir.path()).unwrap();
|
|
assert_eq!(m.n_cols(), 1);
|
|
assert_eq!(m.n(), 4);
|
|
assert_eq!(&*m.row(0), &[10u32]);
|
|
assert_eq!(&*m.row(1), &[200u32]);
|
|
assert_eq!(&*m.row(2), &[300u32]);
|
|
assert_eq!(&*m.row(3), &[1000u32]);
|
|
}
|
|
|
|
#[test]
|
|
fn two_cols_roundtrip() {
|
|
let dir = tempdir().unwrap();
|
|
let mut b = PersistentCompactIntMatrixBuilder::new(3, &dir.path().join("counts")).unwrap();
|
|
let mut col0 = b.add_col().unwrap();
|
|
col0.set(0, 1); col0.set(1, 2); col0.set(2, 3);
|
|
col0.close().unwrap();
|
|
let mut col1 = b.add_col().unwrap();
|
|
col1.set(0, 10); col1.set(1, 20); col1.set(2, 30);
|
|
col1.close().unwrap();
|
|
b.close().unwrap();
|
|
|
|
let m = PersistentCompactIntMatrix::open(dir.path()).unwrap();
|
|
assert_eq!(m.n_cols(), 2);
|
|
assert_eq!(&*m.row(0), &[1u32, 10]);
|
|
assert_eq!(&*m.row(1), &[2u32, 20]);
|
|
assert_eq!(&*m.row(2), &[3u32, 30]);
|
|
}
|
|
|
|
#[test]
|
|
fn resume_continues_n_cols_and_appends_columns() {
|
|
let dir = tempdir().unwrap();
|
|
let counts = dir.path().join("counts");
|
|
|
|
let mut b = PersistentCompactIntMatrixBuilder::new(3, &counts).unwrap();
|
|
let mut col0 = b.add_col().unwrap();
|
|
col0.set(0, 1); col0.set(1, 2); col0.set(2, 3);
|
|
col0.close().unwrap();
|
|
b.close().unwrap();
|
|
|
|
let mut resumed = PersistentCompactIntMatrixBuilder::resume(&counts).unwrap();
|
|
assert_eq!(resumed.n(), 3);
|
|
assert_eq!(resumed.n_cols(), 1);
|
|
let mut col1 = resumed.add_col().unwrap();
|
|
col1.set(0, 10); col1.set(1, 20); col1.set(2, 30);
|
|
col1.close().unwrap();
|
|
resumed.close().unwrap();
|
|
|
|
let m = PersistentCompactIntMatrix::open(dir.path()).unwrap();
|
|
assert_eq!(m.n_cols(), 2);
|
|
assert_eq!(&*m.row(0), &[1u32, 10]);
|
|
assert_eq!(&*m.row(1), &[2u32, 20]);
|
|
assert_eq!(&*m.row(2), &[3u32, 30]);
|
|
}
|
|
|
|
#[test]
|
|
fn resume_twice_keeps_appending() {
|
|
let dir = tempdir().unwrap();
|
|
let counts = dir.path().join("counts");
|
|
|
|
PersistentCompactIntMatrixBuilder::new(2, &counts).unwrap().close().unwrap();
|
|
|
|
for v in [1u32, 2, 3] {
|
|
let mut b = PersistentCompactIntMatrixBuilder::resume(&counts).unwrap();
|
|
let mut col = b.add_col().unwrap();
|
|
col.set(0, v); col.set(1, v * 10);
|
|
col.close().unwrap();
|
|
b.close().unwrap();
|
|
}
|
|
|
|
let m = PersistentCompactIntMatrix::open(dir.path()).unwrap();
|
|
assert_eq!(m.n_cols(), 3);
|
|
assert_eq!(&*m.row(0), &[1u32, 2, 3]);
|
|
assert_eq!(&*m.row(1), &[10u32, 20, 30]);
|
|
}
|
|
|
|
#[test]
|
|
fn col_accessor() {
|
|
let dir = tempdir().unwrap();
|
|
let mut b = PersistentCompactIntMatrixBuilder::new(2, &dir.path().join("counts")).unwrap();
|
|
let mut col0 = b.add_col().unwrap();
|
|
col0.set(0, 5); col0.set(1, 7);
|
|
col0.close().unwrap();
|
|
b.close().unwrap();
|
|
|
|
let m = PersistentCompactIntMatrix::open(dir.path()).unwrap();
|
|
assert_eq!(m.col(0).get(0), 5);
|
|
assert_eq!(m.col(0).get(1), 7);
|
|
}
|
|
|
|
#[test]
|
|
fn zero_cols_roundtrip() {
|
|
let dir = tempdir().unwrap();
|
|
let b = PersistentCompactIntMatrixBuilder::new(10, &dir.path().join("counts")).unwrap();
|
|
b.close().unwrap();
|
|
|
|
let m = PersistentCompactIntMatrix::open(dir.path()).unwrap();
|
|
assert_eq!(m.n_cols(), 0);
|
|
assert_eq!(m.n(), 10);
|
|
}
|
|
|
|
// ── Distance matrix tests ─────────────────────────────────────────────────────
|
|
|
|
#[test]
|
|
fn bray_dist_matrix_symmetry_and_diagonal() {
|
|
// col0=[1,0,1], col1=[1,1,0], col2=[0,1,1]
|
|
let (_d, m) = make_matrix(&[&[1, 0, 1], &[1, 1, 0], &[0, 1, 1]]);
|
|
let dm = m.bray_dist_matrix();
|
|
let n = m.n_cols();
|
|
for i in 0..n { assert_eq!(dm[[i, i]], 0.0, "diagonal"); }
|
|
for i in 0..n {
|
|
for j in 0..n {
|
|
assert!((dm[[i, j]] - dm[[j, i]]).abs() < 1e-12, "symmetry");
|
|
}
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn bray_dist_matrix_values_match_pairwise() {
|
|
let (_d, m) = make_matrix(&[&[1, 0, 1], &[1, 1, 0], &[0, 1, 1]]);
|
|
let dm = m.bray_dist_matrix();
|
|
for i in 0..m.n_cols() {
|
|
for j in 0..m.n_cols() {
|
|
let expected = m.col(i).bray_dist(m.col(j));
|
|
assert!((dm[[i, j]] - expected).abs() < 1e-12, "[{i},{j}]");
|
|
}
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn jaccard_dist_matrix_values_match_pairwise() {
|
|
let (_d, m) = make_matrix(&[&[1, 0, 2], &[0, 1, 1], &[1, 1, 0]]);
|
|
let dm = m.jaccard_dist_matrix();
|
|
for i in 0..m.n_cols() {
|
|
for j in 0..m.n_cols() {
|
|
let expected = m.col(i).jaccard_dist(m.col(j));
|
|
assert!((dm[[i, j]] - expected).abs() < 1e-12, "[{i},{j}]");
|
|
}
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn partial_bray_dist_matrix_consistent() {
|
|
let (_d, m) = make_matrix(&[&[1, 0, 1], &[1, 1, 0], &[0, 1, 1]]);
|
|
let sum_min = m.partial_bray_dist_matrix();
|
|
let col_sums = m.sum();
|
|
let n = m.n_cols();
|
|
|
|
// symmetry
|
|
for i in 0..n {
|
|
for j in 0..n {
|
|
assert_eq!(sum_min[[i, j]], sum_min[[j, i]]);
|
|
}
|
|
}
|
|
|
|
// reconstruct distance from partials and compare to direct method
|
|
for i in 0..n {
|
|
for j in i + 1..n {
|
|
let denom = col_sums[i] + col_sums[j];
|
|
let dist = if denom == 0 { 0.0 } else { 1.0 - 2.0 * sum_min[[i, j]] as f64 / denom as f64 };
|
|
let expected = m.col(i).bray_dist(m.col(j));
|
|
assert!((dist - expected).abs() < 1e-12, "[{i},{j}]");
|
|
}
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn partial_euclidean_dist_matrix_consistent() {
|
|
let (_d, m) = make_matrix(&[&[3, 0], &[0, 4], &[1, 1]]);
|
|
let sq = m.partial_euclidean_dist_matrix();
|
|
let n = m.n_cols();
|
|
for i in 0..n {
|
|
for j in i + 1..n {
|
|
let expected = m.col(i).partial_euclidean_dist(m.col(j));
|
|
assert!((sq[[i, j]] - expected).abs() < 1e-12, "[{i},{j}]");
|
|
assert!((sq[[j, i]] - expected).abs() < 1e-12, "symmetry [{j},{i}]");
|
|
}
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn partial_threshold_jaccard_consistent() {
|
|
let (_d, m) = make_matrix(&[&[3, 0, 2], &[1, 2, 0], &[0, 3, 1]]);
|
|
let threshold = 2u32;
|
|
let (inter, union) = m.partial_threshold_jaccard_dist_matrix(threshold);
|
|
let n = m.n_cols();
|
|
for i in 0..n {
|
|
for j in i + 1..n {
|
|
let (ei, eu) = m.col(i).partial_threshold_jaccard_dist(m.col(j), threshold);
|
|
assert_eq!(inter[[i, j]], ei);
|
|
assert_eq!(union[[i, j]], eu);
|
|
assert_eq!(inter[[j, i]], ei, "symmetry inter");
|
|
assert_eq!(union[[j, i]], eu, "symmetry union");
|
|
}
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn single_col_distance_matrix_is_zero() {
|
|
let (_d, m) = make_matrix(&[&[1, 2, 3]]);
|
|
let dm = m.bray_dist_matrix();
|
|
assert_eq!(dm.shape(), &[1, 1]);
|
|
assert_eq!(dm[[0, 0]], 0.0);
|
|
}
|
|
|
|
#[test]
|
|
fn sum_returns_column_sums() {
|
|
let (_d, m) = make_matrix(&[&[1, 2, 3], &[10, 20, 30], &[0, 0, 5]]);
|
|
let s = m.sum();
|
|
assert_eq!(s[0], 6);
|
|
assert_eq!(s[1], 60);
|
|
assert_eq!(s[2], 5);
|
|
}
|
|
|
|
#[test]
|
|
fn partial_relfreq_bray_matches_full() {
|
|
let (_d, m) = make_matrix(&[&[1, 0, 2], &[0, 1, 1], &[1, 1, 0]]);
|
|
let col_sums = m.sum();
|
|
let partial = m.partial_relfreq_bray_dist_matrix(&col_sums);
|
|
let full = m.relfreq_bray_dist_matrix();
|
|
let n = m.n_cols();
|
|
// partial[i,j] = sum_min_relfreq; full[i,j] = 1 - sum_min_relfreq (off-diagonal only)
|
|
for i in 0..n {
|
|
for j in 0..n {
|
|
if i == j { continue; }
|
|
assert!((partial[[i, j]] - (1.0 - full[[i, j]])).abs() < 1e-12, "[{i},{j}]");
|
|
}
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn partial_relfreq_euclidean_matches_full() {
|
|
let (_d, m) = make_matrix(&[&[3, 0], &[0, 4], &[1, 1]]);
|
|
let col_sums = m.sum();
|
|
let partial = m.partial_relfreq_euclidean_dist_matrix(&col_sums);
|
|
let full = m.relfreq_euclidean_dist_matrix();
|
|
let n = m.n_cols();
|
|
for i in 0..n {
|
|
for j in 0..n {
|
|
// partial = squared euclidean; full = sqrt(partial)
|
|
assert!((partial[[i, j]].sqrt() - full[[i, j]]).abs() < 1e-12, "[{i},{j}]");
|
|
}
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn partial_hellinger_matches_full() {
|
|
let (_d, m) = make_matrix(&[&[3, 0], &[0, 4], &[1, 1]]);
|
|
let col_sums = m.sum();
|
|
let partial = m.partial_hellinger_euclidean_dist_matrix(&col_sums);
|
|
let full = m.hellinger_dist_matrix();
|
|
let n = m.n_cols();
|
|
for i in 0..n {
|
|
for j in 0..n {
|
|
// partial / sqrt(2) gives the Hellinger distance
|
|
assert!((partial[[i, j]].sqrt() / std::f64::consts::SQRT_2 - full[[i, j]]).abs() < 1e-12, "[{i},{j}]");
|
|
}
|
|
}
|
|
}
|
|
|
|
#[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);
|
|
}
|
|
|
|
/// `nonzero_iter` must agree with `row`/`sub_matrix` (already exercised
|
|
/// elsewhere): every nonzero cell among `slots`, no more, no less, values
|
|
/// matching. Compares against both `Columnar` and `Packed` — same data,
|
|
/// two on-disk layouts, one shared `nonzero_iter` implementation.
|
|
#[test]
|
|
fn nonzero_iter_matches_row() {
|
|
let data: &[&[u32]] = &[&[0, 5, 0, 3, 7], &[2, 0, 0, 4, 0], &[0, 0, 9, 0, 1]];
|
|
let (dir_col, m_col) = make_matrix(data);
|
|
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();
|
|
|
|
let slots = [4usize, 0, 3, 1];
|
|
let mut expected: Vec<(usize, usize, u32)> = Vec::new();
|
|
for (i, &slot) in slots.iter().enumerate() {
|
|
for c in 0..data.len() {
|
|
let v = data[c][slot];
|
|
if v != 0 {
|
|
expected.push((i, c, v));
|
|
}
|
|
}
|
|
}
|
|
expected.sort();
|
|
|
|
for (label, m) in [("columnar", &m_col), ("packed", &m_pack)] {
|
|
let mut got: Vec<(usize, usize, u32)> = m.nonzero_iter(&slots).collect();
|
|
got.sort();
|
|
assert_eq!(got, expected, "{label}");
|
|
}
|
|
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.
|
|
// col0=[1,2,3,4,5], col1=[5,4,3,2,1]
|
|
// global sums: col0=15, col1=15
|
|
// split: first=[1,2], second=[3,4,5]
|
|
let (_d1, m1) = make_matrix(&[&[1, 2], &[5, 4]]);
|
|
let (_d2, m2) = make_matrix(&[&[3, 4, 5], &[3, 2, 1]]);
|
|
let global_sums = ndarray::array![15u64, 15u64];
|
|
let p1 = m1.partial_relfreq_bray_dist_matrix(&global_sums);
|
|
let p2 = m2.partial_relfreq_bray_dist_matrix(&global_sums);
|
|
let combined = &p1 + &p2;
|
|
|
|
let (_d, m_full) = make_matrix(&[&[1, 2, 3, 4, 5], &[5, 4, 3, 2, 1]]);
|
|
let full_partial = m_full.partial_relfreq_bray_dist_matrix(&global_sums);
|
|
let n = m_full.n_cols();
|
|
for i in 0..n {
|
|
for j in 0..n {
|
|
assert!((combined[[i, j]] - full_partial[[i, j]]).abs() < 1e-12, "[{i},{j}]");
|
|
}
|
|
}
|
|
}
|
|
|
|
// ── collect_slots_values tests ────────────────────────────────────────────────────────────
|
|
|
|
fn make_pciv(counts: &[u32]) -> (tempfile::TempDir, PersistentCompactIntVec) {
|
|
let dir = tempdir().unwrap();
|
|
let path = dir.path().join("c.pciv");
|
|
let mut b = PersistentCompactIntVecBuilder::new(counts.len(), &path).unwrap();
|
|
for (i, &v) in counts.iter().enumerate() { b.set(i, v); }
|
|
b.close().unwrap();
|
|
let r = PersistentCompactIntVec::open(&path).unwrap();
|
|
(dir, r)
|
|
}
|
|
|
|
#[test]
|
|
fn pciv_collect_slots_values_in_order() {
|
|
let counts = vec![10u32, 255, 300, 1000];
|
|
let (_dir, v) = make_pciv(&counts);
|
|
let got = v.collect_slots_values(&[0, 1, 2, 3]);
|
|
assert_eq!(got, counts);
|
|
}
|
|
|
|
#[test]
|
|
fn pciv_collect_slots_values_out_of_order() {
|
|
let counts = vec![10u32, 255, 300, 1000];
|
|
let (_dir, v) = make_pciv(&counts);
|
|
let got = v.collect_slots_values(&[3, 0, 2, 1]);
|
|
assert_eq!(got, vec![1000, 10, 300, 255]);
|
|
}
|
|
|
|
#[test]
|
|
fn pciv_collect_slots_values_with_duplicates() {
|
|
let counts = vec![10u32, 255, 300];
|
|
let (_dir, v) = make_pciv(&counts);
|
|
let got = v.collect_slots_values(&[0, 2, 0, 1]);
|
|
assert_eq!(got, vec![10, 300, 10, 255]);
|
|
}
|
|
|
|
#[test]
|
|
fn pciv_collect_slots_values_empty() {
|
|
let (_dir, v) = make_pciv(&[10u32, 20]);
|
|
let got: Vec<u32> = v.collect_slots_values(&[]);
|
|
assert!(got.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn pciv_collect_slots_values_out_of_bounds_panics() {
|
|
let (_dir, v) = make_pciv(&[10u32, 20]);
|
|
let result = std::panic::catch_unwind(|| v.collect_slots_values(&[0, 2]));
|
|
assert!(result.is_err(), "collect_slots_values should panic on out-of-bounds slot");
|
|
}
|
|
|
|
// IntSliceView collect_slots_values (same logic, exercised through the view)
|
|
#[test]
|
|
fn intslice_view_collect_slots_values() {
|
|
let counts = vec![10u32, 255, 300, 1000];
|
|
let (_dir, v) = make_pciv(&counts);
|
|
let view = v.view();
|
|
assert_eq!(view.collect_slots_values(&[0, 1, 2]), vec![10, 255, 300]);
|
|
assert_eq!(view.collect_slots_values(&[3, 1]), vec![1000, 255]);
|
|
}
|