Skip to main content

rightkit_embed/
fusion.rs

1//! Reciprocal-rank fusion (RRF) of ranked lists. Pure, no model and no feature gate.
2
3use std::collections::{HashMap, HashSet};
4use std::hash::Hash;
5
6/// The damping constant of the original RRF paper (Cormack, Clarke and Büttcher, 2009).
7pub const RRF_K: f64 = 60.0;
8
9/// Fuses ranked lists, best first, by reciprocal rank. An item scores `sum(1 / (k + rank))`
10/// over every list that holds it, with rank starting at 1, so an item near the top of several
11/// lists beats one at the top of only one. Returns each item once, best score first. Items with
12/// equal scores keep the order in which they first appear across the lists.
13///
14/// An item may appear at most once per list; a repeat in one list counts at its first rank.
15/// A `k` that is not finite or is negative returns an empty result. Use [`RRF_K`] for the
16/// conventional value.
17pub fn reciprocal_rank_fusion<K: Eq + Hash + Clone>(rankings: &[&[K]], k: f64) -> Vec<(K, f64)> {
18    if !k.is_finite() || k < 0.0 {
19        return Vec::new();
20    }
21    let mut slots: HashMap<K, usize> = HashMap::new();
22    let mut fused: Vec<(K, f64)> = Vec::new();
23    for ranking in rankings {
24        let mut seen: HashSet<&K> = HashSet::with_capacity(ranking.len());
25        for (position, item) in ranking.iter().enumerate() {
26            if !seen.insert(item) {
27                continue;
28            }
29            let contribution = 1.0 / (k + position as f64 + 1.0);
30            match slots.get(item) {
31                Some(&slot) => fused[slot].1 += contribution,
32                None => {
33                    slots.insert(item.clone(), fused.len());
34                    fused.push((item.clone(), contribution));
35                }
36            }
37        }
38    }
39    // Stable sort: equal scores keep first-appearance order.
40    fused.sort_by(|left, right| right.1.total_cmp(&left.1));
41    fused
42}
43
44#[cfg(test)]
45mod tests {
46    use super::*;
47
48    #[test]
49    fn items_in_several_lists_outrank_a_single_top_hit() {
50        let lexical = ["a", "b", "c"];
51        let semantic = ["b", "d", "a"];
52        let fused = reciprocal_rank_fusion(&[&lexical[..], &semantic[..]], RRF_K);
53        let order: Vec<&str> = fused.iter().map(|(key, _)| *key).collect();
54        assert_eq!(order, vec!["b", "a", "d", "c"]);
55        let b = 1.0 / 62.0 + 1.0 / 61.0;
56        assert!((fused[0].1 - b).abs() < 1e-12);
57    }
58
59    #[test]
60    fn ties_keep_first_appearance_and_repeats_count_once() {
61        let first = ["x", "y", "x"];
62        let second = ["y", "x"];
63        let fused = reciprocal_rank_fusion(&[&first[..], &second[..]], 0.0);
64        let order: Vec<&str> = fused.iter().map(|(key, _)| *key).collect();
65        assert_eq!(order, vec!["x", "y"]);
66        // k = 0: x = 1/1 (first list) + 1/2 (second list); y = 1/2 + 1/1.
67        assert!((fused[0].1 - 1.5).abs() < 1e-12);
68        assert!((fused[1].1 - 1.5).abs() < 1e-12);
69    }
70
71    #[test]
72    fn empty_input_and_invalid_k_are_empty() {
73        let none: [&[u32]; 0] = [];
74        assert!(reciprocal_rank_fusion(&none, RRF_K).is_empty());
75        let one = [1u32, 2];
76        assert!(reciprocal_rank_fusion(&[&one[..]], -1.0).is_empty());
77        assert!(reciprocal_rank_fusion(&[&one[..]], f64::NAN).is_empty());
78        assert_eq!(reciprocal_rank_fusion(&[&one[..]], RRF_K).len(), 2);
79    }
80}