Optimize sparse bit matrix packing to avoid intermediate copies

Refactor `pack_sparse_bit_matrix` to conditionally open either a packed or columnar dense representation based on file existence. This eliminates the previous unconditional pre-conversion step that forced a full packed copy. Update cleanup logic and add tests to verify direct transposition, proper file handling, and idempotency for both formats.
This commit is contained in:
Eric Coissac
2026-08-20 14:38:20 +02:00
parent c2e0533fa9
commit 640f29725b
2 changed files with 89 additions and 13 deletions
+23 -12
View File
@@ -34,6 +34,8 @@ use crate::meta::field;
use crate::rankselect::{PersistentRankSelectBitVec, PersistentRankSelectBitVecBuilder};
use super::PersistentBitMatrix;
use super::col_path;
use super::columnar::ColumnarBitMatrix;
use super::packed::PackedBitMatrix;
fn is_multi_path(dir: &Path) -> PathBuf { dir.join("is_multi.prsb") }
@@ -485,24 +487,33 @@ impl BitPartials for PersistentSparseBitMatrix {
/// convention (old-format files removed only after the new format is
/// fully written, so a crash mid-conversion leaves the previous, still
/// valid format rather than a half-written one). Idempotent — does
/// nothing if `is_multi.prsb` already exists. Packs to dense
/// (`matrix.pbmx`) first via [`super::pack_bit_matrix`] if that hasn't
/// happened yet, since the dense→sparse transpose (`build_from_dense`)
/// needs random row access, which only the packed/columnar dense forms
/// give — not a further reason to keep the dense file around afterward.
/// nothing if `is_multi.prsb` already exists.
///
/// The dense→sparse transpose (`build_from_dense`) only needs `fill_row`,
/// which costs the same `O(n_cols)` per row on `Columnar` as on `Packed`
/// (see `ColumnarBitMatrix::fill_row`) — so this reads straight from
/// whichever dense form is already on disk. In particular, when this runs
/// right after a build (the common case: `presence/` is still `Columnar`,
/// `matrix.pbmx` doesn't exist yet), it never materialises a full packed
/// copy of the presence matrix just to read it back once and delete it.
pub fn pack_sparse_bit_matrix(dir: &Path) -> io::Result<()> {
if is_multi_path(dir).exists() {
return Ok(());
}
let packed_path = dir.join("matrix.pbmx");
if !packed_path.exists() {
super::pack_bit_matrix(dir)?;
if packed_path.exists() {
let dense = PersistentBitMatrix::Packed(PackedBitMatrix::open(&packed_path)?);
PersistentSparseBitMatrixBuilder::build_from_dense(&dense, dir)?.close()?;
drop(dense);
fs::remove_file(&packed_path)?;
} else {
let dense = PersistentBitMatrix::Columnar(ColumnarBitMatrix::open(dir)?);
let n_cols = dense.n_cols();
PersistentSparseBitMatrixBuilder::build_from_dense(&dense, dir)?.close()?;
drop(dense);
for c in 0..n_cols { let _ = fs::remove_file(col_path(dir, c)); }
let _ = fs::remove_file(dir.join("meta.json"));
}
let dense = PersistentBitMatrix::Packed(PackedBitMatrix::open(&packed_path)?);
PersistentSparseBitMatrixBuilder::build_from_dense(&dense, dir)?.close()?;
drop(dense);
fs::remove_file(&packed_path)?;
Ok(())
}
+66 -1
View File
@@ -1,6 +1,6 @@
use tempfile::tempdir;
use crate::{BinaryMatrix, PersistentBitMatrix, PersistentBitMatrixBuilder, PersistentSparseBitMatrix, PersistentSparseBitMatrixBuilder};
use crate::{pack_bit_matrix, pack_sparse_bit_matrix, BinaryMatrix, PersistentBitMatrix, PersistentBitMatrixBuilder, PersistentSparseBitMatrix, PersistentSparseBitMatrixBuilder};
/// Builds a dense `PersistentBitMatrix` from column-major `bool` data —
/// mirrors `tests/bitmatrix.rs`'s own `make_matrix` helper.
@@ -307,3 +307,68 @@ fn reopen_large_real_shaped_data_with_heavy_dedup() {
assert_eq!(&row_as_bool(&m, slot), expected, "row {slot}");
}
}
/// Common case: `pack_sparse_bit_matrix` runs right after a build, while
/// `presence/` is still `Columnar` (`matrix.pbmx` doesn't exist yet). Must
/// transpose straight from the per-column `.pbiv` files — never write a
/// full packed copy just to read it back and delete it.
#[test]
fn pack_sparse_bit_matrix_transposes_columnar_directly() {
let cols: &[&[bool]] = &[
&[true, false, true, false],
&[true, true, false, false],
&[false, true, true, true],
];
let (dir, dense) = make_dense(cols);
let presence_dir = dir.path().join("presence");
assert!(!presence_dir.join("matrix.pbmx").exists(), "still Columnar before packing");
let expected_rows: Vec<Box<[bool]>> = (0..dense.n()).map(|s| dense.row(s)).collect();
let n_cols = dense.n_cols();
drop(dense);
pack_sparse_bit_matrix(&presence_dir).unwrap();
assert!(presence_dir.join("is_multi.prsb").exists());
assert!(!presence_dir.join("meta.json").exists(), "columnar meta.json cleaned up");
assert!(!presence_dir.join("col_000000.pbiv").exists(), "columnar column files cleaned up");
assert!(!presence_dir.join("matrix.pbmx").exists(), "never materialised");
let sparse = PersistentSparseBitMatrix::open(&presence_dir).unwrap();
assert_eq!(sparse.n(), expected_rows.len());
assert_eq!(sparse.n_cols(), n_cols);
for (slot, expected) in expected_rows.iter().enumerate() {
assert_eq!(&*sparse.row(slot), &**expected, "slot {slot}");
}
// Idempotent — a second call is a no-op, doesn't error.
pack_sparse_bit_matrix(&presence_dir).unwrap();
}
/// Regression for the pre-existing path: if `matrix.pbmx` is already there
/// (e.g. `pack_bit_matrix` ran earlier for a non-sparse consumer), it's
/// still used as the transpose source and cleaned up afterward.
#[test]
fn pack_sparse_bit_matrix_from_already_packed_matrix() {
let cols: &[&[bool]] = &[
&[true, false, true],
&[false, true, true],
];
let (dir, dense) = make_dense(cols);
let presence_dir = dir.path().join("presence");
let expected_rows: Vec<Box<[bool]>> = (0..dense.n()).map(|s| dense.row(s)).collect();
let n_cols = dense.n_cols();
drop(dense);
pack_bit_matrix(&presence_dir).unwrap();
assert!(presence_dir.join("matrix.pbmx").exists());
pack_sparse_bit_matrix(&presence_dir).unwrap();
assert!(!presence_dir.join("matrix.pbmx").exists(), "packed intermediate cleaned up");
let sparse = PersistentSparseBitMatrix::open(&presence_dir).unwrap();
assert_eq!(sparse.n_cols(), n_cols);
for (slot, expected) in expected_rows.iter().enumerate() {
assert_eq!(&*sparse.row(slot), &**expected, "slot {slot}");
}
}