llm_kernel/search/
rrf.rs

1//! Reciprocal Rank Fusion (RRF) for combining multiple ranked result sets.
2//!
3//! RRF is a simple but effective fusion method:
4//! `score(d) = Σ 1/(k + rank_i(d))` where k is a constant (typically 60).
5
6use std::collections::HashMap;
7
8use crate::search::types::SearchResult;
9
10/// Fuse multiple ranked result lists using Reciprocal Rank Fusion.
11///
12/// `k` is the RRF constant (typically 60). Higher values smooth the
13/// contribution of top-ranked results.
14pub fn rrf_fuse(result_sets: &[Vec<SearchResult>], k: u32) -> Vec<SearchResult> {
15    let mut scores: HashMap<String, f32> = HashMap::new();
16    let mut texts: HashMap<String, String> = HashMap::new();
17
18    for results in result_sets {
19        for (rank, result) in results.iter().enumerate() {
20            let entry = scores.entry(result.id.clone()).or_insert(0.0);
21            *entry += 1.0 / (k as f32 + rank as f32 + 1.0);
22            texts
23                .entry(result.id.clone())
24                .or_insert_with(|| result.text.clone());
25        }
26    }
27
28    let mut fused: Vec<SearchResult> = scores
29        .into_iter()
30        .map(|(id, score)| SearchResult {
31            text: texts.remove(&id).unwrap_or_default(),
32            id,
33            score,
34        })
35        .collect();
36
37    fused.sort_by(|a, b| {
38        b.score
39            .partial_cmp(&a.score)
40            .unwrap_or(std::cmp::Ordering::Equal)
41    });
42    fused
43}
44
45#[cfg(test)]
46mod tests {
47    use super::*;
48
49    fn make_results(ids: &[(&str, f32)]) -> Vec<SearchResult> {
50        ids.iter()
51            .map(|(id, score)| SearchResult {
52                id: id.to_string(),
53                score: *score,
54                text: format!("text for {id}"),
55            })
56            .collect()
57    }
58
59    #[test]
60    fn empty_inputs() {
61        let result = rrf_fuse(&[], 60);
62        assert!(result.is_empty());
63    }
64
65    #[test]
66    fn single_list() {
67        let list = make_results(&[("a", 0.9), ("b", 0.5)]);
68        let fused = rrf_fuse(&[list], 60);
69        assert_eq!(fused.len(), 2);
70        assert_eq!(fused[0].id, "a"); // rank 0 → highest RRF score
71    }
72
73    #[test]
74    fn two_lists_overlap() {
75        let bm25 = make_results(&[("a", 0.9), ("b", 0.7)]);
76        let vector = make_results(&[("b", 0.95), ("c", 0.6)]);
77
78        let fused = rrf_fuse(&[bm25, vector], 60);
79        assert_eq!(fused.len(), 3);
80
81        // "b" appears in both lists at rank 0 and 1 → should rank highest
82        assert_eq!(fused[0].id, "b");
83    }
84
85    #[test]
86    fn rrf_score_formula() {
87        let list = make_results(&[("x", 1.0)]);
88        let fused = rrf_fuse(&[list], 60);
89        // rank 0 → 1/(60 + 0 + 1) = 1/61 ≈ 0.01639
90        let expected = 1.0 / 61.0;
91        assert!((fused[0].score - expected).abs() < 1e-6);
92    }
93
94    #[test]
95    fn three_lists() {
96        let a = make_results(&[("doc1", 1.0), ("doc2", 0.8)]);
97        let b = make_results(&[("doc2", 0.9), ("doc3", 0.7)]);
98        let c = make_results(&[("doc1", 0.8), ("doc3", 0.9)]);
99
100        let fused = rrf_fuse(&[a, b, c], 60);
101        assert_eq!(fused.len(), 3);
102
103        // doc1: rank 0 in a, not in b, rank 1 in c → 1/61 + 0 + 1/62
104        // doc2: rank 1 in a, rank 0 in b, not in c → 1/62 + 1/61 + 0
105        // doc3: not in a, rank 1 in b, rank 1 in c → 0 + 1/62 + 1/62
106        // doc1 and doc2 tie (same sum), doc3 is lowest
107        let doc3_score = fused.iter().find(|r| r.id == "doc3").unwrap().score;
108        let doc1_score = fused.iter().find(|r| r.id == "doc1").unwrap().score;
109        assert!(doc1_score > doc3_score);
110    }
111}
llm_kernel/search/rrf.rs

llm_kernel/search/
rrf.rs
