nodedb_query/
fusion.rs

1// SPDX-License-Identifier: Apache-2.0
2
3use std::collections::HashMap;
4use std::hash::Hash;
5
6/// Reciprocal Rank Fusion (RRF) for combining ranked results from multiple engines.
7///
8/// RRF is used when a query hits multiple engines (e.g., vector similarity +
9/// metadata filter + BM25 text search). Each engine returns a ranked list;
10/// RRF combines them into a single ranked list.
11///
12/// Formula: RRF_score(d) = Σ 1 / (k + rank_i(d))
13/// where k is a smoothing constant (default 60).
14/// RRF smoothing constant. Standard value from Cormack et al. (2009).
15pub const DEFAULT_RRF_K: f64 = 60.0;
16
17/// A scored result from a single engine.
18///
19/// `K` is the key the engine ranks under. It defaults to `String` for
20/// callers whose key is opaque text.
21#[derive(Debug, Clone)]
22pub struct RankedResult<K = String> {
23    /// Document identifier (engine-specific).
24    pub document_id: K,
25    /// Rank within the engine's result list (0-based).
26    pub rank: usize,
27    /// Original score from the engine (for diagnostics).
28    pub score: f32,
29    /// Source engine identifier.
30    pub source: &'static str,
31}
32
33/// A fused result after RRF combination.
34#[derive(Debug, Clone)]
35pub struct FusedResult<K = String> {
36    pub document_id: K,
37    pub rrf_score: f64,
38    /// Per-engine contributions for explainability.
39    pub contributions: Vec<(&'static str, f64)>,
40}
41
42/// Sum each key's contributions, sort by RRF score, and cut to `top_k`.
43///
44/// The tie-break on `document_id` is what makes the output deterministic:
45/// RRF produces many equal scores and the score map iterates in arbitrary
46/// order, so a unique secondary key gives a total order.
47fn finish_fusion<K: Clone + Eq + Hash + Ord>(
48    scores: HashMap<K, Vec<(&'static str, f64)>>,
49    top_k: usize,
50) -> Vec<FusedResult<K>> {
51    let mut fused: Vec<FusedResult<K>> = scores
52        .into_iter()
53        .map(|(doc_id, contributions)| {
54            let rrf_score = contributions.iter().map(|(_, s)| s).sum();
55            FusedResult {
56                document_id: doc_id,
57                rrf_score,
58                contributions,
59            }
60        })
61        .collect();
62
63    fused.sort_unstable_by(|a, b| {
64        crate::numeric_cmp::cmp_f64(b.rrf_score, a.rrf_score)
65            .then_with(|| a.document_id.cmp(&b.document_id))
66    });
67    fused.truncate(top_k);
68    fused
69}
70
71/// Fuse multiple ranked result lists using Reciprocal Rank Fusion.
72///
73/// Each inner Vec is a ranked list from one engine (ordered by relevance).
74/// Returns the top_k fused results sorted by RRF score (descending).
75pub fn reciprocal_rank_fusion<K: Clone + Eq + Hash + Ord>(
76    ranked_lists: &[Vec<RankedResult<K>>],
77    k: Option<f64>,
78    top_k: usize,
79) -> Vec<FusedResult<K>> {
80    let k = k.unwrap_or(DEFAULT_RRF_K);
81
82    let mut scores: HashMap<K, Vec<(&'static str, f64)>> = HashMap::new();
83
84    for list in ranked_lists {
85        for result in list {
86            let contribution = 1.0 / (k + result.rank as f64 + 1.0);
87            scores
88                .entry(result.document_id.clone())
89                .or_default()
90                .push((result.source, contribution));
91        }
92    }
93
94    finish_fusion(scores, top_k)
95}
96
97/// Fuse ranked lists with per-list **linear weights**.
98///
99/// Each list's reciprocal-rank contribution is scaled by its weight, so a
100/// more-trusted source can dominate: `contribution = weight_i / (k + rank + 1)`.
101/// Unlike [`reciprocal_rank_fusion_weighted`] (which varies the `k` decay
102/// constant per list), this scales contribution magnitude directly, which is
103/// the right lever when one source (e.g. BM25) is far more reliable than another
104/// (e.g. a weak dense index) — equal-weight RRF would let the weak source drag
105/// down the strong source's ranking. `weights.len()` must equal
106/// `ranked_lists.len()`.
107///
108/// # Panics
109///
110/// Panics if `weights.len() != ranked_lists.len()`.
111pub fn reciprocal_rank_fusion_linear<K: Clone + Eq + Hash + Ord>(
112    ranked_lists: &[Vec<RankedResult<K>>],
113    k: Option<f64>,
114    weights: &[f64],
115    top_k: usize,
116) -> Vec<FusedResult<K>> {
117    assert_eq!(
118        ranked_lists.len(),
119        weights.len(),
120        "weights length must match ranked_lists length"
121    );
122    let k = k.unwrap_or(DEFAULT_RRF_K);
123
124    let mut scores: HashMap<K, Vec<(&'static str, f64)>> = HashMap::new();
125
126    for (list_idx, list) in ranked_lists.iter().enumerate() {
127        let w = weights[list_idx];
128        for result in list {
129            let contribution = w / (k + result.rank as f64 + 1.0);
130            scores
131                .entry(result.document_id.clone())
132                .or_default()
133                .push((result.source, contribution));
134        }
135    }
136
137    finish_fusion(scores, top_k)
138}
139
140/// Fuse ranked lists with per-list k-constants for weighted influence.
141///
142/// Each list gets its own k value: lower k → steeper rank discount → more
143/// influence. Typical usage: `k_i = base_k / weight_i`.
144///
145/// # Panics
146///
147/// Panics if `k_per_list.len() != ranked_lists.len()`.
148pub fn reciprocal_rank_fusion_weighted<K: Clone + Eq + Hash + Ord>(
149    ranked_lists: &[Vec<RankedResult<K>>],
150    k_per_list: &[f64],
151    top_k: usize,
152) -> Vec<FusedResult<K>> {
153    assert_eq!(
154        ranked_lists.len(),
155        k_per_list.len(),
156        "k_per_list length must match ranked_lists length"
157    );
158
159    let mut scores: HashMap<K, Vec<(&'static str, f64)>> = HashMap::new();
160
161    for (list_idx, list) in ranked_lists.iter().enumerate() {
162        let k = k_per_list[list_idx];
163        for result in list {
164            let contribution = 1.0 / (k + result.rank as f64 + 1.0);
165            scores
166                .entry(result.document_id.clone())
167                .or_default()
168                .push((result.source, contribution));
169        }
170    }
171
172    finish_fusion(scores, top_k)
173}
174
175#[cfg(test)]
176mod tests {
177    use super::*;
178
179    fn make_ranked(doc_ids: &[&str], source: &'static str) -> Vec<RankedResult> {
180        doc_ids
181            .iter()
182            .enumerate()
183            .map(|(rank, &id)| RankedResult {
184                document_id: id.to_string(),
185                rank,
186                score: 1.0 - (rank as f32 * 0.1),
187                source,
188            })
189            .collect()
190    }
191
192    #[test]
193    fn single_list_preserves_order() {
194        let list = make_ranked(&["d1", "d2", "d3"], "vector");
195        let fused = reciprocal_rank_fusion(&[list], None, 10);
196        assert_eq!(fused.len(), 3);
197        assert_eq!(fused[0].document_id, "d1");
198    }
199
200    #[test]
201    fn overlapping_lists_boost_common_docs() {
202        let vector = make_ranked(&["d1", "d2", "d3"], "vector");
203        let sparse = make_ranked(&["d2", "d1", "d4"], "sparse");
204        let fused = reciprocal_rank_fusion(&[vector, sparse], None, 10);
205        let top2_ids: Vec<&str> = fused[..2].iter().map(|f| f.document_id.as_str()).collect();
206        assert!(top2_ids.contains(&"d1"));
207        assert!(top2_ids.contains(&"d2"));
208    }
209
210    /// A non-`String` key fuses the same way: the key type is the caller's.
211    #[test]
212    fn typed_key_fuses_and_tie_breaks_on_ord() {
213        let make = |ids: &[u32], source: &'static str| -> Vec<RankedResult<u32>> {
214            ids.iter()
215                .enumerate()
216                .map(|(rank, &id)| RankedResult {
217                    document_id: id,
218                    rank,
219                    score: 0.0,
220                    source,
221                })
222                .collect()
223        };
224        let a = make(&[7, 3], "a");
225        let b = make(&[3, 9], "b");
226        let fused = reciprocal_rank_fusion(&[a, b], None, 10);
227        assert_eq!(fused[0].document_id, 3);
228        // 7 and 9 each rank #1 in one list, so they tie; `Ord` breaks it.
229        assert_eq!(fused[1].document_id, 7);
230        assert_eq!(fused[2].document_id, 9);
231    }
232
233    #[test]
234    fn weighted_rrf() {
235        let list_a = make_ranked(&["a1", "a2"], "vector");
236        let list_b = make_ranked(&["b1", "a1"], "text");
237        let fused = reciprocal_rank_fusion_weighted(&[list_a, list_b], &[30.0, 120.0], 10);
238        let a1 = fused.iter().find(|f| f.document_id == "a1").unwrap();
239        assert_eq!(a1.contributions.len(), 2);
240    }
241
242    #[test]
243    fn linear_weight_lets_strong_source_dominate() {
244        // `a1` is ranked #1 by the strong source and #2 by the weak source;
245        // `b1` is ranked #1 only by the weak source. With a heavy weight on
246        // the strong source, `a1` must outrank `b1` — equal-weight RRF would
247        // tie/flip them on the weak source's #1.
248        let strong = make_ranked(&["a1", "a2"], "strong");
249        let weak = make_ranked(&["b1", "a1"], "weak");
250        let fused = reciprocal_rank_fusion_linear(&[strong, weak], None, &[4.0, 0.25], 10);
251
252        let a1 = fused.iter().position(|f| f.document_id == "a1").unwrap();
253        let b1 = fused.iter().position(|f| f.document_id == "b1").unwrap();
254        assert!(a1 < b1, "a1 (rank {a1}) should outrank b1 (rank {b1})");
255
256        let a1_res = &fused[a1];
257        assert_eq!(a1_res.contributions.len(), 2);
258        // Contribution magnitude scales linearly with the per-list weight.
259        let strong_contrib = a1_res
260            .contributions
261            .iter()
262            .find(|(src, _)| *src == "strong")
263            .map(|(_, s)| *s)
264            .unwrap();
265        let expected = 4.0 / (DEFAULT_RRF_K + 0.0 + 1.0);
266        assert!((strong_contrib - expected).abs() < 1e-12);
267    }
268
269    #[test]
270    #[should_panic(expected = "weights length must match ranked_lists length")]
271    fn linear_mismatched_weights_panics() {
272        let list = make_ranked(&["d1"], "vector");
273        let _ = reciprocal_rank_fusion_linear(&[list], None, &[1.0, 2.0], 10);
274    }
275
276    #[test]
277    fn empty() {
278        assert!(reciprocal_rank_fusion::<String>(&[], None, 10).is_empty());
279        assert!(reciprocal_rank_fusion_linear::<String>(&[], None, &[], 10).is_empty());
280        assert!(reciprocal_rank_fusion_weighted::<String>(&[], &[], 10).is_empty());
281    }
282}
nodedb_query/fusion.rs

nodedb_query/
fusion.rs
