quantumclaw/clawhdf5
3
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 1 Wiki Activity

perf(agent): unranked BM25 scores and a top-k merge — same rankings, 4-5x faster

Browse Source 
Fusion min-max normalises over every keyword match, so hybrid_search asked
BM25 for a ranked list of the whole corpus: a hash insert per posting, then a
sort of every match, then the merge sorted every candidate again to keep k.

- BM25Index::scores returns every match unsorted, accumulated in a dense array
  (contributions are strictly positive, so zero means untouched). search() is
  built on it with the bounded heap.
- merge_vector_keyword partitions out its top k (select_nth) and orders only
  those, with the same score-then-id order.
- Both hybrid paths use scores().

Rankings are identical (equivalence tests for both changes). p50 0.24 -> 0.07
ms (1K), 2.1 -> 0.49 ms (10K), 23 -> 4.65 ms (100K).

The harness gains --fusion-study, which measured the alternative — capping the
keyword pool — and found it changes the top-10 for most queries (overlap
0.83-0.92, different #1 for 10-35%) for only a 2x saving. Not adopted.

Co-Authored-By: Claude Fable 5.1 <[email protected]>
This commit is contained in:
osobh
2026-09-19 13:27:31 -07:00
co-authored by Claude Fable 5.1
parent f15bf2eb22
commit 390a2e3836
6 changed files with 216 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/clawhdf5-agent/src/bm25.rs
+59 -24
View File
@@ -83,35 +83,15 @@ impl BM25Index {
/// Uses Block-Max WAND for early termination when remaining documents
/// cannot beat the current top-k threshold.
pub fn search(&self, query: &str, k: usize) -> Vec<(usize, f32)> {
if self.num_docs == 0 || k == 0 {
if k == 0 {
return Vec::new();
}
// 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();
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));
*scores.entry(doc_id).or_insert(0.0) += idf * tf;
}
}
// 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).
// deterministic.
let mut heap: BinaryHeap<Reverse<(HeapScore, Reverse<usize>)>> =
BinaryHeap::with_capacity(k + 1);
for (doc_id, score) in scores {
BinaryHeap::with_capacity(k.min(1024) + 1);
for (doc_id, score) in self.scores(query) {
heap.push(Reverse((HeapScore(score), Reverse(doc_id))));
if heap.len() > k {
heap.pop();
@@ -125,6 +105,47 @@ impl BM25Index {
results
}
/// The BM25 score of **every** matching document, in doc-id order, unsorted
/// by score. Score fusion normalises over the whole matching set, so it
/// needs all of these but not their ranking; producing a ranked list of
/// every match (`search(query, corpus_len)`) spent most of its time sorting.
pub fn scores(&self, query: &str) -> Vec<(usize, f32)> {
if self.num_docs == 0 {
return Vec::new();
}
// Term-at-a-time accumulation into a dense array: a common term has a
// posting per document, and hashing each one dominated query time.
// 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 acc = vec![0.0f32; self.doc_lengths.len()];
let mut matched = false;
for token in tokenize(query) {
let Some(postings) = self.inverted.get(token.as_str()) else {
continue;
};
matched = true;
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));
acc[doc_id] += idf * tf;
}
}
if !matched {
return Vec::new();
}
// Every contribution is strictly positive (idf = ln(1 + x), x > 0), so
// a zero entry is a document no query term touched.
acc.into_iter()
.enumerate()
.filter(|&(_, score)| score > 0.0)
.collect()
}
/// Number of document slots (live or not) the index covers. Ids are
/// positions in the document list it mirrors.
pub fn len(&self) -> usize {
@@ -564,6 +585,20 @@ mod tests {
}
}
#[test]
fn scores_is_the_unranked_form_of_a_full_search() {
let mut state = 99u64;
let docs: Vec<String> = (0..200).map(|_| random_doc(&mut state)).collect();
let tombstones: Vec<u8> = (0..200).map(|i| u8::from(i % 7 == 0)).collect();
let index = BM25Index::build(&docs, &tombstones);
for query in ["alpha", "beta gamma x1", "missing", ""] {
let mut all = index.scores(query);
all.sort_by(|a, b| b.1.total_cmp(&a.1).then(a.0.cmp(&b.0)));
assert_eq!(all, index.search(query, docs.len()), "{query:?}");
assert!(all.iter().all(|(id, _)| tombstones[*id] == 0));
}
}
#[test]
fn ties_break_towards_the_lower_doc_id() {
let docs: Vec<String> = (0..6).map(|_| "same text".to_string()).collect();
crates/clawhdf5-agent/src/hybrid.rs
+33 -5
View File
@@ -58,7 +58,7 @@ pub fn hybrid_search(
vector_search::cosine_similarity_batch(query_embedding, vectors, tombstones)
}
};
let kw_scores = bm25_index.search(query_text, vectors.len());
let kw_scores = bm25_index.scores(query_text);
merge_vector_keyword(vec_scores, kw_scores, vector_weight, keyword_weight, k)
}
@@ -92,13 +92,22 @@ pub fn merge_vector_keyword(
let mut results: Vec<(usize, f32)> = merged.into_iter().collect();
// Index tie-break: `merged` is a HashMap, so without it the ties that
// survive `truncate` differ from run to run.
results.sort_by(|a, b| {
// survive differ from run to run.
let by_score_then_id = |a: &(usize, f32), b: &(usize, f32)| {
b.1.partial_cmp(&a.1)
.unwrap_or(std::cmp::Ordering::Equal)
.then(a.0.cmp(&b.0))
});
results.truncate(k);
};
// Only the top k are wanted: partition them out, then order just those,
// instead of sorting every candidate (the keyword side can be the corpus).
if k == 0 {
return Vec::new();
}
if results.len() > k {
results.select_nth_unstable_by(k - 1, by_score_then_id);
results.truncate(k);
}
results.sort_by(by_score_then_id);
results
}
@@ -344,6 +353,25 @@ mod tests {
assert_eq!(result[0].1, 1.0);
}
#[test]
fn merge_top_k_matches_a_full_sort() {
// Many ties (scores repeat) so the index tie-break is exercised.
let vec_scores: Vec<(usize, f32)> = (0..300).map(|i| (i, ((i * 7) % 13) as f32)).collect();
let kw_scores: Vec<(usize, f32)> = (100..500).map(|i| (i, ((i * 5) % 11) as f32)).collect();
let everything =
merge_vector_keyword(vec_scores.clone(), kw_scores.clone(), 0.7, 0.3, 10_000);
assert_eq!(everything.len(), 500);
assert!(
everything
.windows(2)
.all(|w| { w[0].1 > w[1].1 || (w[0].1 == w[1].1 && w[0].0 < w[1].0) })
);
for k in [0, 1, 7, 50, 499, 500, 501] {
let top = merge_vector_keyword(vec_scores.clone(), kw_scores.clone(), 0.7, 0.3, k);
assert_eq!(top, everything[..k.min(500)], "k = {k}");
}
}
#[test]
fn normalize_scores_all_equal() {
let matched = normalize_scores(&[(0, 0.4), (1, 0.4)]);
crates/clawhdf5-agent/src/search.rs
+3 -1
View File
@@ -35,7 +35,9 @@ impl HDF5Memory {
.into_iter()
.map(|(id, dist)| (id, 1.0 - dist))
.collect();
let kw_scores = bm25.search(query_text, self.cache.len());
// Fusion normalises over every keyword match, so it needs all
// the scores — but not ranked.
let kw_scores = bm25.scores(query_text);
hybrid::merge_vector_keyword(
vec_scores,
kw_scores,