INT-09: Persistent BM25 index via .bm25 sidecar file
Eliminates the O(N × terms) rebuild on every HDF5Memory::open() call
for large corpora.
Changes:
bm25.rs — Add BM25Index::to_bytes() / from_bytes()
Compact binary format (magic "BM25" + version byte, then doc_lengths,
inverted posting lists, and idf cache, all length-prefixed LE u32/f32).
from_bytes() validates magic, version, and expected doc count so a
stale or corrupted sidecar falls back to a fresh build.
lib.rs — Wire sidecar into open() and flush()
- bm25_sidecar_path() free function returns <h5 path>.bm25
- open(): after WAL replay, tries to load the sidecar; uses it if
valid, otherwise leaves bm25_cache = None for lazy rebuild.
- flush(): if bm25_cache is Some, writes the sidecar alongside the
.h5 file. Failure is best-effort — a write error is silently
swallowed so it never disrupts the main flush path.
Tests (bm25.rs):
- sidecar_round_trip_preserves_search_results: verifies identical
doc_id and score (within 1e-5) before and after round-trip.
- sidecar_stale_doc_count_rejected: wrong expected_doc_count → None.
- sidecar_bad_magic_rejected: corrupted magic bytes → None.
- sidecar_empty_index_round_trip: zero-doc edge case.
Co-Authored-By: Claude Sonnet 4.6 <[email protected]>
This commit is contained in:
ClawHDF5 Planner
co-authored by
Claude Sonnet 4.6
Claude Sonnet 4.6
parent
fdc4572ab7
commit
ca8a3a4a2e
@@ -218,6 +218,171 @@ impl BM25Index {
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Sidecar serialization (BM25 persistence — INT-09)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Magic bytes for the `.bm25` sidecar format.
|
||||
const SIDECAR_MAGIC: [u8; 4] = [0x42, 0x4D, 0x32, 0x35]; // "BM25"
|
||||
/// Current sidecar format version.
|
||||
const SIDECAR_VERSION: u8 = 0x01;
|
||||
|
||||
impl BM25Index {
|
||||
/// Serialize the index into a compact binary format suitable for writing to
|
||||
/// the `.bm25` sidecar file.
|
||||
///
|
||||
/// Format:
|
||||
/// ```text
|
||||
/// [4] magic "BM25"
|
||||
/// [1] version byte
|
||||
/// [4] doc_lengths.len() as le u32 (= total chunk count, including tombstones)
|
||||
/// [4] num_docs as le u32
|
||||
/// [4] avg_dl as le f32
|
||||
/// [N*4] doc_lengths as le u32 each
|
||||
/// [4] inverted entry count as le u32
|
||||
/// per inverted entry:
|
||||
/// [4] token byte length as le u32
|
||||
/// [L] UTF-8 token bytes
|
||||
/// [4] posting count as le u32
|
||||
/// per posting: [4] doc_id le u32, [4] term_freq le u32
|
||||
/// [4] idf entry count as le u32
|
||||
/// per idf entry:
|
||||
/// [4] token byte length as le u32
|
||||
/// [L] UTF-8 token bytes
|
||||
/// [4] idf score as le f32
|
||||
/// ```
|
||||
pub fn to_bytes(&self) -> Vec<u8> {
|
||||
let mut buf = Vec::with_capacity(
|
||||
9 + self.doc_lengths.len() * 4 + self.inverted.len() * 16 + self.idf_cache.len() * 16,
|
||||
);
|
||||
|
||||
buf.extend_from_slice(&SIDECAR_MAGIC);
|
||||
buf.push(SIDECAR_VERSION);
|
||||
buf.extend_from_slice(&(self.doc_lengths.len() as u32).to_le_bytes());
|
||||
buf.extend_from_slice(&(self.num_docs as u32).to_le_bytes());
|
||||
buf.extend_from_slice(&self.avg_dl.to_le_bytes());
|
||||
|
||||
for &dl in &self.doc_lengths {
|
||||
buf.extend_from_slice(&dl.to_le_bytes());
|
||||
}
|
||||
|
||||
buf.extend_from_slice(&(self.inverted.len() as u32).to_le_bytes());
|
||||
for (token, postings) in &self.inverted {
|
||||
let tb = token.as_bytes();
|
||||
buf.extend_from_slice(&(tb.len() as u32).to_le_bytes());
|
||||
buf.extend_from_slice(tb);
|
||||
buf.extend_from_slice(&(postings.len() as u32).to_le_bytes());
|
||||
for &(doc_id, tf) in postings {
|
||||
buf.extend_from_slice(&(doc_id as u32).to_le_bytes());
|
||||
buf.extend_from_slice(&tf.to_le_bytes());
|
||||
}
|
||||
}
|
||||
|
||||
buf.extend_from_slice(&(self.idf_cache.len() as u32).to_le_bytes());
|
||||
for (token, &idf) in &self.idf_cache {
|
||||
let tb = token.as_bytes();
|
||||
buf.extend_from_slice(&(tb.len() as u32).to_le_bytes());
|
||||
buf.extend_from_slice(tb);
|
||||
buf.extend_from_slice(&idf.to_le_bytes());
|
||||
}
|
||||
|
||||
buf
|
||||
}
|
||||
|
||||
/// Deserialize an index from the bytes produced by [`to_bytes`].
|
||||
///
|
||||
/// Returns `None` if the bytes are malformed (bad magic, wrong version,
|
||||
/// truncated data, or non-UTF-8 tokens). The caller should fall back to
|
||||
/// [`BM25Index::build`] when `None` is returned.
|
||||
///
|
||||
/// `expected_doc_count` is the total number of chunks (including tombstones)
|
||||
/// currently in the cache. If it does not match the serialized
|
||||
/// `doc_lengths.len()`, the sidecar is stale and `None` is returned.
|
||||
pub fn from_bytes(data: &[u8], expected_doc_count: usize) -> Option<Self> {
|
||||
let mut pos = 0usize;
|
||||
|
||||
macro_rules! read_bytes {
|
||||
($n:expr) => {{
|
||||
let end = pos + $n;
|
||||
if end > data.len() {
|
||||
return None;
|
||||
}
|
||||
let slice = &data[pos..end];
|
||||
pos = end;
|
||||
slice
|
||||
}};
|
||||
}
|
||||
macro_rules! read_u32 {
|
||||
() => {{
|
||||
u32::from_le_bytes(read_bytes!(4).try_into().ok()?)
|
||||
}};
|
||||
}
|
||||
macro_rules! read_f32 {
|
||||
() => {{
|
||||
f32::from_le_bytes(read_bytes!(4).try_into().ok()?)
|
||||
}};
|
||||
}
|
||||
|
||||
// Magic + version
|
||||
let magic = read_bytes!(4);
|
||||
if magic != SIDECAR_MAGIC {
|
||||
return None;
|
||||
}
|
||||
let version = read_bytes!(1)[0];
|
||||
if version != SIDECAR_VERSION {
|
||||
return None;
|
||||
}
|
||||
|
||||
// doc_lengths
|
||||
let doc_count = read_u32!() as usize;
|
||||
if doc_count != expected_doc_count {
|
||||
return None; // stale sidecar
|
||||
}
|
||||
let num_docs = read_u32!() as usize;
|
||||
let avg_dl = read_f32!();
|
||||
let mut doc_lengths = Vec::with_capacity(doc_count);
|
||||
for _ in 0..doc_count {
|
||||
doc_lengths.push(read_u32!());
|
||||
}
|
||||
|
||||
// inverted index
|
||||
let inv_count = read_u32!() as usize;
|
||||
let mut inverted: HashMap<String, Vec<(usize, u32)>> = HashMap::with_capacity(inv_count);
|
||||
for _ in 0..inv_count {
|
||||
let tlen = read_u32!() as usize;
|
||||
let token = std::str::from_utf8(read_bytes!(tlen)).ok()?.to_string();
|
||||
let plen = read_u32!() as usize;
|
||||
let mut postings = Vec::with_capacity(plen);
|
||||
for _ in 0..plen {
|
||||
let doc_id = read_u32!() as usize;
|
||||
let tf = read_u32!();
|
||||
postings.push((doc_id, tf));
|
||||
}
|
||||
inverted.insert(token, postings);
|
||||
}
|
||||
|
||||
// idf cache
|
||||
let idf_count = read_u32!() as usize;
|
||||
let mut idf_cache: HashMap<String, f32> = HashMap::with_capacity(idf_count);
|
||||
for _ in 0..idf_count {
|
||||
let tlen = read_u32!() as usize;
|
||||
let token = std::str::from_utf8(read_bytes!(tlen)).ok()?.to_string();
|
||||
let idf = read_f32!();
|
||||
idf_cache.insert(token, idf);
|
||||
}
|
||||
|
||||
Some(Self {
|
||||
inverted,
|
||||
idf_cache,
|
||||
doc_lengths,
|
||||
avg_dl,
|
||||
num_docs,
|
||||
k1: DEFAULT_K1,
|
||||
b: DEFAULT_B,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Tokenize a string: lowercase, split on non-alphanumeric characters,
|
||||
/// filter empty tokens.
|
||||
fn tokenize(text: &str) -> Vec<String> {
|
||||
@@ -451,4 +616,74 @@ mod tests {
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// Sidecar serialization round-trip (INT-09)
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
#[test]
|
||||
fn sidecar_round_trip_preserves_search_results() {
|
||||
let docs = vec![
|
||||
"the quick brown fox jumps over the lazy dog".to_string(),
|
||||
"rust programming language systems programming".to_string(),
|
||||
"python scripting and data science".to_string(),
|
||||
];
|
||||
let tombstones = vec![0u8, 0, 0];
|
||||
let original = BM25Index::build(&docs, &tombstones);
|
||||
|
||||
// Serialize then deserialize.
|
||||
let bytes = original.to_bytes();
|
||||
let restored =
|
||||
BM25Index::from_bytes(&bytes, docs.len()).expect("round-trip must succeed");
|
||||
|
||||
// Both indexes must return identical results for the same query.
|
||||
let orig_results = original.search("rust programming", 10);
|
||||
let rest_results = restored.search("rust programming", 10);
|
||||
assert_eq!(
|
||||
orig_results.len(),
|
||||
rest_results.len(),
|
||||
"result count mismatch"
|
||||
);
|
||||
for (a, b) in orig_results.iter().zip(rest_results.iter()) {
|
||||
assert_eq!(a.0, b.0, "doc_id mismatch after round-trip");
|
||||
assert!(
|
||||
(a.1 - b.1).abs() < 1e-5,
|
||||
"score mismatch: {} vs {} for doc {}",
|
||||
a.1,
|
||||
b.1,
|
||||
a.0
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sidecar_stale_doc_count_rejected() {
|
||||
let docs = vec!["hello world".to_string()];
|
||||
let tombstones = vec![0u8];
|
||||
let idx = BM25Index::build(&docs, &tombstones);
|
||||
let bytes = idx.to_bytes();
|
||||
// Pass wrong expected_doc_count — should return None.
|
||||
assert!(BM25Index::from_bytes(&bytes, 999).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sidecar_bad_magic_rejected() {
|
||||
let docs = vec!["hello".to_string()];
|
||||
let tombstones = vec![0u8];
|
||||
let idx = BM25Index::build(&docs, &tombstones);
|
||||
let mut bytes = idx.to_bytes();
|
||||
// Corrupt the magic bytes.
|
||||
bytes[0] = 0xFF;
|
||||
assert!(BM25Index::from_bytes(&bytes, 1).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sidecar_empty_index_round_trip() {
|
||||
let docs: Vec<String> = vec![];
|
||||
let tombstones: Vec<u8> = vec![];
|
||||
let idx = BM25Index::build(&docs, &tombstones);
|
||||
let bytes = idx.to_bytes();
|
||||
let restored = BM25Index::from_bytes(&bytes, 0).expect("empty index must round-trip");
|
||||
assert_eq!(restored.search("anything", 5).len(), 0);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -60,6 +60,17 @@ pub fn cosine_similarity_prenorm(
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use cache::MemoryCache;
|
||||
|
||||
/// Returns the path to the BM25 sidecar file for an HDF5 memory file at `h5_path`.
|
||||
///
|
||||
/// The sidecar lives next to the `.h5` file with a `.bm25` extension appended
|
||||
/// (e.g. `memory.h5` → `memory.h5.bm25`). It is loaded on `open()` to skip the
|
||||
/// O(N × terms) rebuild when the cache is large, and written on every `flush()`.
|
||||
fn bm25_sidecar_path(h5_path: &Path) -> PathBuf {
|
||||
let mut p = h5_path.as_os_str().to_owned();
|
||||
p.push(".bm25");
|
||||
PathBuf::from(p)
|
||||
}
|
||||
#[cfg(feature = "hnsw")]
|
||||
use clawhdf5_ann::{DistanceMetric, HnswIndex};
|
||||
use ephemeral::{EphemeralConfig, EphemeralStore};
|
||||
@@ -290,6 +301,16 @@ impl HDF5Memory {
|
||||
None
|
||||
};
|
||||
|
||||
// Try to load the BM25 sidecar so the first hybrid_search after open()
|
||||
// skips the O(N × terms) rebuild. Fall back to None (lazy rebuild) if
|
||||
// the sidecar is absent, malformed, or has a mismatched doc count.
|
||||
let bm25_cache = {
|
||||
let sidecar_path = bm25_sidecar_path(&config.path);
|
||||
std::fs::read(&sidecar_path)
|
||||
.ok()
|
||||
.and_then(|b| bm25::BM25Index::from_bytes(&b, cache.chunks.len()))
|
||||
};
|
||||
|
||||
Ok(Self {
|
||||
config,
|
||||
cache,
|
||||
@@ -306,7 +327,7 @@ impl HDF5Memory {
|
||||
hnsw_dirty: true,
|
||||
#[cfg(feature = "hnsw")]
|
||||
hnsw_synced_len: 0,
|
||||
bm25_cache: None,
|
||||
bm25_cache,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -326,9 +347,35 @@ impl HDF5Memory {
|
||||
if let Some(ref mut w) = self.wal {
|
||||
w.truncate()?;
|
||||
}
|
||||
// Persist the BM25 index alongside the .h5 file so the next open()
|
||||
// can skip the O(N × terms) rebuild. Only write when we have a cached
|
||||
// index; if there is none, leave any existing sidecar in place.
|
||||
if let Some(ref idx) = self.bm25_cache {
|
||||
let sidecar_path = bm25_sidecar_path(&self.config.path);
|
||||
let bytes = idx.to_bytes();
|
||||
// Best-effort: a sidecar write failure is not fatal — the caller
|
||||
// will rebuild from scratch on the next open().
|
||||
let _ = std::fs::write(&sidecar_path, &bytes);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Path to the `.bm25` sidecar file for this memory store.
|
||||
fn bm25_sidecar_path(&self) -> std::path::PathBuf {
|
||||
bm25_sidecar_path(&self.config.path)
|
||||
}
|
||||
|
||||
/// Try to load the BM25 index from the `.bm25` sidecar file.
|
||||
///
|
||||
/// Returns `Some(index)` if the sidecar exists and is valid for the current
|
||||
/// cache state (same total chunk count including tombstones). Returns
|
||||
/// `None` if the sidecar is absent, malformed, or stale.
|
||||
fn load_bm25_sidecar(&self) -> Option<bm25::BM25Index> {
|
||||
let sidecar_path = self.bm25_sidecar_path();
|
||||
let bytes = std::fs::read(&sidecar_path).ok()?;
|
||||
bm25::BM25Index::from_bytes(&bytes, self.cache.chunks.len())
|
||||
}
|
||||
|
||||
// ---- HNSW index maintenance --------------------------------------------
|
||||
//
|
||||
// The index mirrors the cache: HNSW node id == cache index, kept aligned by
|
||||
|
||||
Reference in New Issue
Block a user