perf(agent): unranked BM25 scores and a top-k merge — same rankings, 4-5x faster
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:
co-authored by
Claude Fable 5.1
parent
f15bf2eb22
commit
390a2e3836
@@ -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)]);
|
||||
|
||||
Reference in New Issue
Block a user