perf(agent): persistent incremental BM25 index; no store rewrite per query
hybrid_search rebuilt the BM25 index from scratch (re-tokenising every record) and rewrote the whole .h5 file on every single query, so a query cost O(store size) in both CPU and disk I/O. Steady-state p50 per the search harness: 5.5 -> 0.24 ms (1K), 49 -> 2.1 ms (10K), 884 -> 23 ms (100K). - BM25Index is incremental: add_document / remove_document keep it exactly equivalent to a fresh build over the same live documents (property test: 60 random op sequences compared against BM25Index::build after every step). IDF moves to query time since it depends on the live document count. Top-k uses a bounded heap, ties break by doc id (results were HashMap-ordered), and the "WAND" code that computed a bound and then discarded it is removed. - HDF5Memory keeps one index for its lifetime, built lazily. Appends are picked up by ensure_bm25_fresh whatever path added them; delete and in-place update report themselves; compaction drops the index. A test drives every mutation and compares against a fresh build. - A query no longer calls flush(). Activation boosts are marked dirty and persisted by the next checkpoint, including a best-effort one on drop so a search-only session keeps them (approved behaviour change). Activation weights are capped at 16.0; they previously grew without bound. The archived mission branch's BM25 cache was reviewed and not used: it was invalidated by every write, so interleaved save/search still rebuilt per query, and it changed the default fusion weights. Co-Authored-By: Claude Fable 5.1 <[email protected]>
This commit is contained in:
osobh
co-authored by
Claude Fable 5.1
Claude Fable 5.1
parent
61424d1418
commit
2bfbb7fb4b
+217
-113
@@ -3,10 +3,15 @@
|
||||
//! Provides a standard BM25 (Okapi BM25) implementation with an in-memory
|
||||
//! inverted index. Tombstoned documents are excluded from indexing and search.
|
||||
//!
|
||||
//! Optimizations:
|
||||
//! - Cached IDF scores (don't recompute per query)
|
||||
//! - Sorted posting lists by doc_id for cache-friendly access
|
||||
//! - Block-Max WAND early termination
|
||||
//! The index is **incremental**: [`BM25Index::add_document`] and
|
||||
//! [`BM25Index::remove_document`] keep it exactly equivalent to one built from
|
||||
//! scratch over the same live documents, so a store can maintain one index for
|
||||
//! its lifetime instead of re-tokenising the whole corpus per query. To make
|
||||
//! that possible IDF is computed at query time (it depends on the live
|
||||
//! document count) rather than cached at build time.
|
||||
//!
|
||||
//! - Posting lists sorted by doc id
|
||||
//! - Bounded-heap top-k; results ordered by score, then doc id (deterministic)
|
||||
|
||||
use std::cmp::Reverse;
|
||||
use std::collections::{BinaryHeap, HashMap};
|
||||
@@ -41,10 +46,11 @@ const DEFAULT_B: f32 = 0.75;
|
||||
pub struct BM25Index {
|
||||
/// Inverted index: token -> sorted list of (doc_id, term_frequency).
|
||||
inverted: HashMap<String, Vec<(usize, u32)>>,
|
||||
/// Cached IDF scores per token.
|
||||
idf_cache: HashMap<String, f32>,
|
||||
/// Number of tokens in each document (0 for tombstoned docs).
|
||||
doc_lengths: Vec<u32>,
|
||||
/// Sum of `doc_lengths` over live documents (keeps `avg_dl` exact under
|
||||
/// incremental updates).
|
||||
total_length: u64,
|
||||
/// Average document length across non-tombstoned docs.
|
||||
avg_dl: f32,
|
||||
/// Number of non-tombstoned documents.
|
||||
@@ -60,8 +66,8 @@ impl BM25Index {
|
||||
pub fn build(documents: &[String], tombstones: &[u8]) -> Self {
|
||||
let mut index = Self {
|
||||
inverted: HashMap::new(),
|
||||
idf_cache: HashMap::new(),
|
||||
doc_lengths: vec![0; documents.len()],
|
||||
total_length: 0,
|
||||
avg_dl: 0.0,
|
||||
num_docs: 0,
|
||||
k1: DEFAULT_K1,
|
||||
@@ -81,103 +87,135 @@ impl BM25Index {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
let tokens = tokenize(query);
|
||||
if tokens.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
// Collect posting lists and cached IDF scores for query tokens
|
||||
type QueryTerm<'a> = (&'a str, f32, &'a [(usize, u32)]);
|
||||
let mut query_terms: Vec<QueryTerm<'_>> = Vec::new();
|
||||
for token in &tokens {
|
||||
if let (Some(postings), Some(&idf)) = (
|
||||
self.inverted.get(token.as_str()),
|
||||
self.idf_cache.get(token.as_str()),
|
||||
) {
|
||||
query_terms.push((token, idf, postings));
|
||||
}
|
||||
}
|
||||
|
||||
if query_terms.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
// Accumulate BM25 scores per document using WAND-style scoring
|
||||
// Term-at-a-time accumulation. IDF is computed here rather than cached
|
||||
// at build time: it depends on the live document count, which changes
|
||||
// with every incremental add/remove, and costs one `ln` per query term.
|
||||
let mut scores: HashMap<usize, f32> = HashMap::new();
|
||||
|
||||
// Compute maximum possible contribution per term for WAND
|
||||
let max_tf_score: Vec<f32> = query_terms
|
||||
.iter()
|
||||
.map(|(_, idf, _)| {
|
||||
// Upper bound: max TF contribution when tf is high and dl is short
|
||||
let max_tf_num = 10.0 * (self.k1 + 1.0);
|
||||
let max_tf_den = 10.0 + self.k1 * (1.0 - self.b);
|
||||
idf * max_tf_num / max_tf_den
|
||||
})
|
||||
.collect();
|
||||
|
||||
let total_max_contribution: f32 = max_tf_score.iter().sum();
|
||||
|
||||
// Threshold for WAND early termination. `top_k_heap` is a min-heap of
|
||||
// size k (worst-of-the-top-k at the head) so it can be maintained in
|
||||
// O(log k) per update instead of re-sorting the whole buffer.
|
||||
let mut threshold = 0.0f32;
|
||||
let mut top_k_heap: BinaryHeap<Reverse<HeapScore>> = BinaryHeap::with_capacity(k);
|
||||
|
||||
for (term_idx, (_, idf, postings)) in query_terms.iter().enumerate() {
|
||||
for &(doc_id, freq) in *postings {
|
||||
for token in tokenize(query) {
|
||||
let Some(postings) = self.inverted.get(token.as_str()) else {
|
||||
continue;
|
||||
};
|
||||
let df = postings.len() as f32;
|
||||
let idf = ((self.num_docs as f32 - df + 0.5) / (df + 0.5) + 1.0).ln();
|
||||
for &(doc_id, freq) in postings {
|
||||
let dl = self.doc_lengths[doc_id] as f32;
|
||||
let freq_f = freq as f32;
|
||||
let tf = (freq_f * (self.k1 + 1.0))
|
||||
/ (freq_f + self.k1 * (1.0 - self.b + self.b * dl / self.avg_dl));
|
||||
let contribution = idf * tf;
|
||||
|
||||
let entry = scores.entry(doc_id).or_insert(0.0);
|
||||
*entry += contribution;
|
||||
|
||||
// WAND check: if this doc's current partial score + remaining
|
||||
// max terms can't beat threshold, we can skip (but we still
|
||||
// accumulate since we process term-at-a-time)
|
||||
if term_idx == query_terms.len() - 1 {
|
||||
// Last term: check if this doc beats threshold
|
||||
let final_score = *entry;
|
||||
if top_k_heap.len() >= k {
|
||||
if final_score > threshold {
|
||||
// Replace the current worst-of-top-k.
|
||||
top_k_heap.pop();
|
||||
top_k_heap.push(Reverse(HeapScore(final_score)));
|
||||
threshold = top_k_heap.peek().map(|Reverse(s)| s.0).unwrap_or(0.0);
|
||||
}
|
||||
} else {
|
||||
top_k_heap.push(Reverse(HeapScore(final_score)));
|
||||
if top_k_heap.len() == k {
|
||||
threshold = top_k_heap.peek().map(|Reverse(s)| s.0).unwrap_or(0.0);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
// After processing each term, check if remaining terms can
|
||||
// possibly produce results above threshold
|
||||
let remaining_max: f32 = max_tf_score[term_idx + 1..].iter().sum();
|
||||
if remaining_max < threshold && total_max_contribution > 0.0 {
|
||||
// Early termination: remaining terms can't produce new top-k
|
||||
// entries on their own. But existing partial scores may still
|
||||
// be updated, so we continue (WAND is approximate here).
|
||||
let _ = remaining_max; // hint to compiler
|
||||
*scores.entry(doc_id).or_insert(0.0) += idf * tf;
|
||||
}
|
||||
}
|
||||
|
||||
let mut results: Vec<(usize, f32)> = scores.into_iter().collect();
|
||||
results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
|
||||
results.truncate(k);
|
||||
// Top-k with a bounded min-heap: O(matches * log k) instead of sorting
|
||||
// every match. Ties break towards the lower doc id so results are
|
||||
// deterministic (the accumulator is a HashMap).
|
||||
let mut heap: BinaryHeap<Reverse<(HeapScore, Reverse<usize>)>> =
|
||||
BinaryHeap::with_capacity(k + 1);
|
||||
for (doc_id, score) in scores {
|
||||
heap.push(Reverse((HeapScore(score), Reverse(doc_id))));
|
||||
if heap.len() > k {
|
||||
heap.pop();
|
||||
}
|
||||
}
|
||||
let mut results: Vec<(usize, f32)> = heap
|
||||
.into_iter()
|
||||
.map(|Reverse((HeapScore(score), Reverse(doc_id)))| (doc_id, score))
|
||||
.collect();
|
||||
results.sort_by(|a, b| b.1.total_cmp(&a.1).then(a.0.cmp(&b.0)));
|
||||
results
|
||||
}
|
||||
|
||||
/// Number of document slots (live or not) the index covers. Ids are
|
||||
/// positions in the document list it mirrors.
|
||||
pub fn len(&self) -> usize {
|
||||
self.doc_lengths.len()
|
||||
}
|
||||
|
||||
/// `true` when the index covers no document slots.
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.doc_lengths.is_empty()
|
||||
}
|
||||
|
||||
/// Index `text` as document `doc_id`, which must be the next free id
|
||||
/// (`self.len()`) or an existing slot that is currently empty (removed or
|
||||
/// tombstoned). After any sequence of `add_document` / `remove_document`
|
||||
/// calls the index scores exactly as one freshly built from the same live
|
||||
/// documents.
|
||||
pub fn add_document(&mut self, doc_id: usize, text: &str) {
|
||||
if doc_id >= self.doc_lengths.len() {
|
||||
self.doc_lengths.resize(doc_id + 1, 0);
|
||||
}
|
||||
debug_assert_eq!(self.doc_lengths[doc_id], 0, "slot {doc_id} is occupied");
|
||||
|
||||
let tokens = tokenize(text);
|
||||
let mut term_freqs: HashMap<&str, u32> = HashMap::new();
|
||||
for token in &tokens {
|
||||
*term_freqs.entry(token).or_insert(0) += 1;
|
||||
}
|
||||
for (token, freq) in term_freqs {
|
||||
let postings = self.inverted.entry(token.to_string()).or_default();
|
||||
// Posting lists stay sorted by doc id; appends are the common case.
|
||||
match postings.last() {
|
||||
Some(&(last, _)) if last >= doc_id => {
|
||||
let at = postings.partition_point(|&(id, _)| id < doc_id);
|
||||
postings.insert(at, (doc_id, freq));
|
||||
}
|
||||
_ => postings.push((doc_id, freq)),
|
||||
}
|
||||
}
|
||||
self.doc_lengths[doc_id] = tokens.len() as u32;
|
||||
self.total_length += tokens.len() as u64;
|
||||
self.num_docs += 1;
|
||||
self.refresh_avg_dl();
|
||||
}
|
||||
|
||||
/// Extend the index to cover `len` document slots, leaving new ones empty.
|
||||
/// Used for slots that hold no live document (tombstoned records).
|
||||
pub fn pad_to(&mut self, len: usize) {
|
||||
if len > self.doc_lengths.len() {
|
||||
self.doc_lengths.resize(len, 0);
|
||||
}
|
||||
}
|
||||
|
||||
/// Remove document `doc_id`, whose indexed text was `text`. The text is
|
||||
/// needed to find its postings; pass exactly what was added.
|
||||
pub fn remove_document(&mut self, doc_id: usize, text: &str) {
|
||||
let tokens = tokenize(text);
|
||||
let mut seen: std::collections::HashSet<&str> = std::collections::HashSet::new();
|
||||
for token in &tokens {
|
||||
if !seen.insert(token) {
|
||||
continue;
|
||||
}
|
||||
if let Some(postings) = self.inverted.get_mut(token.as_str()) {
|
||||
if let Ok(at) = postings.binary_search_by_key(&doc_id, |&(id, _)| id) {
|
||||
postings.remove(at);
|
||||
}
|
||||
if postings.is_empty() {
|
||||
self.inverted.remove(token.as_str());
|
||||
}
|
||||
}
|
||||
}
|
||||
if let Some(len) = self.doc_lengths.get_mut(doc_id) {
|
||||
self.total_length = self.total_length.saturating_sub(u64::from(*len));
|
||||
*len = 0;
|
||||
}
|
||||
self.num_docs = self.num_docs.saturating_sub(1);
|
||||
self.refresh_avg_dl();
|
||||
}
|
||||
|
||||
fn refresh_avg_dl(&mut self) {
|
||||
self.avg_dl = if self.num_docs > 0 {
|
||||
self.total_length as f32 / self.num_docs as f32
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
}
|
||||
|
||||
/// Rebuild the index from scratch (e.g., after compaction).
|
||||
pub fn rebuild(&mut self, documents: &[String], tombstones: &[u8]) {
|
||||
self.inverted.clear();
|
||||
self.idf_cache.clear();
|
||||
self.doc_lengths = vec![0; documents.len()];
|
||||
self.total_length = 0;
|
||||
self.avg_dl = 0.0;
|
||||
self.num_docs = 0;
|
||||
self.index_documents(documents, tombstones);
|
||||
@@ -214,23 +252,13 @@ impl BM25Index {
|
||||
}
|
||||
|
||||
self.num_docs = count;
|
||||
self.avg_dl = if count > 0 {
|
||||
total_length as f32 / count as f32
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
self.total_length = total_length;
|
||||
self.refresh_avg_dl();
|
||||
|
||||
// Sort posting lists by doc_id for cache-friendly access
|
||||
for postings in self.inverted.values_mut() {
|
||||
postings.sort_by_key(|&(doc_id, _)| doc_id);
|
||||
}
|
||||
|
||||
// Pre-compute and cache IDF scores
|
||||
for (token, postings) in &self.inverted {
|
||||
let df = postings.len() as f32;
|
||||
let idf = ((self.num_docs as f32 - df + 0.5) / (df + 0.5) + 1.0).ln();
|
||||
self.idf_cache.insert(token.clone(), idf);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -386,24 +414,21 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cached_idf_consistent_with_computed() {
|
||||
fn score_matches_the_bm25_formula() {
|
||||
let docs = vec![
|
||||
"rust programming".to_string(),
|
||||
"rust systems".to_string(),
|
||||
"python scripting".to_string(),
|
||||
];
|
||||
let tombstones = vec![0, 0, 0];
|
||||
let index = BM25Index::build(&docs, &tombstones);
|
||||
let index = BM25Index::build(&docs, &[0, 0, 0]);
|
||||
|
||||
// IDF for "rust" (appears in 2 of 3 docs)
|
||||
let idf_rust = index.idf_cache.get("rust").unwrap();
|
||||
let expected_idf = ((3.0f32 - 2.0 + 0.5) / (2.0 + 0.5) + 1.0).ln();
|
||||
assert!(
|
||||
(idf_rust - expected_idf).abs() < 1e-6,
|
||||
"cached IDF mismatch: {} vs {}",
|
||||
idf_rust,
|
||||
expected_idf
|
||||
);
|
||||
// "python": df = 1 of N = 3. Every doc has the average length (2) and
|
||||
// tf = 1, so the tf factor is exactly 1 and the score is the IDF.
|
||||
let results = index.search("python", 3);
|
||||
let expected_idf = ((3.0f32 - 1.0 + 0.5) / (1.0 + 0.5) + 1.0).ln();
|
||||
assert_eq!(results.len(), 1);
|
||||
assert_eq!(results[0].0, 2);
|
||||
assert!((results[0].1 - expected_idf).abs() < 1e-6, "{results:?}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -467,4 +492,83 @@ mod tests {
|
||||
);
|
||||
}
|
||||
}
|
||||
/// Documents drawn from a small vocabulary so terms collide heavily.
|
||||
fn random_doc(state: &mut u64) -> String {
|
||||
const VOCAB: &[&str] = &[
|
||||
"alpha", "beta", "gamma", "delta", "eps", "zeta", "eta", "x1",
|
||||
];
|
||||
let mut next = || {
|
||||
*state = state
|
||||
.wrapping_mul(6364136223846793005)
|
||||
.wrapping_add(1442695040888963407);
|
||||
(*state >> 33) as usize
|
||||
};
|
||||
let len = 1 + next() % 9;
|
||||
(0..len)
|
||||
.map(|_| VOCAB[next() % VOCAB.len()])
|
||||
.collect::<Vec<_>>()
|
||||
.join(" ")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn incremental_updates_match_a_fresh_build_exactly() {
|
||||
for seed in 0..60u64 {
|
||||
let mut state = seed.wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1;
|
||||
let mut docs: Vec<String> = Vec::new();
|
||||
let mut tombstones: Vec<u8> = Vec::new();
|
||||
let mut index = BM25Index::build(&docs, &tombstones);
|
||||
|
||||
for step in 0..80 {
|
||||
state = state.wrapping_mul(6364136223846793005).wrapping_add(1);
|
||||
let live: Vec<usize> = (0..docs.len()).filter(|&i| tombstones[i] == 0).collect();
|
||||
match (state >> 40) % 4 {
|
||||
0 if !live.is_empty() => {
|
||||
// delete
|
||||
let id = live[(state >> 20) as usize % live.len()];
|
||||
index.remove_document(id, &docs[id]);
|
||||
tombstones[id] = 1;
|
||||
}
|
||||
1 if !live.is_empty() => {
|
||||
// update in place
|
||||
let id = live[(state >> 20) as usize % live.len()];
|
||||
let new_text = random_doc(&mut state);
|
||||
index.remove_document(id, &docs[id]);
|
||||
index.add_document(id, &new_text);
|
||||
docs[id] = new_text;
|
||||
}
|
||||
_ => {
|
||||
let text = random_doc(&mut state);
|
||||
index.add_document(docs.len(), &text);
|
||||
docs.push(text);
|
||||
tombstones.push(0);
|
||||
}
|
||||
}
|
||||
|
||||
let fresh = BM25Index::build(&docs, &tombstones);
|
||||
for query in ["alpha", "beta gamma", "x1 zeta alpha delta", "missing"] {
|
||||
let got = index.search(query, 5);
|
||||
let want = fresh.search(query, 5);
|
||||
assert_eq!(got.len(), want.len(), "seed {seed} step {step} {query:?}");
|
||||
for (g, w) in got.iter().zip(&want) {
|
||||
assert_eq!(
|
||||
g.0, w.0,
|
||||
"seed {seed} step {step} {query:?}: {got:?} vs {want:?}"
|
||||
);
|
||||
assert!(
|
||||
(g.1 - w.1).abs() < 1e-5,
|
||||
"seed {seed} step {step} {query:?}"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ties_break_towards_the_lower_doc_id() {
|
||||
let docs: Vec<String> = (0..6).map(|_| "same text".to_string()).collect();
|
||||
let index = BM25Index::build(&docs, &[0; 6]);
|
||||
let ids: Vec<usize> = index.search("same", 3).into_iter().map(|r| r.0).collect();
|
||||
assert_eq!(ids, [0, 1, 2]);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user