From 640f29725b9acc389a3809914b356b55edf8c618 Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Thu, 20 Aug 2026 06:12:28 +0200 Subject: [PATCH] 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. --- src/obicompactvec/src/bitmatrix/sparse.rs | 35 ++++++++---- src/obicompactvec/src/tests/sparse.rs | 67 ++++++++++++++++++++++- 2 files changed, 89 insertions(+), 13 deletions(-) diff --git a/src/obicompactvec/src/bitmatrix/sparse.rs b/src/obicompactvec/src/bitmatrix/sparse.rs index 409ef6fb..51df6a84 100644 --- a/src/obicompactvec/src/bitmatrix/sparse.rs +++ b/src/obicompactvec/src/bitmatrix/sparse.rs @@ -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(()) } diff --git a/src/obicompactvec/src/tests/sparse.rs b/src/obicompactvec/src/tests/sparse.rs index 59c1dacf..e464535e 100644 --- a/src/obicompactvec/src/tests/sparse.rs +++ b/src/obicompactvec/src/tests/sparse.rs @@ -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> = (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> = (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}"); + } +}