|
| 1 | +// SPDX-License-Identifier: Apache-2.0 |
| 2 | +// SPDX-FileCopyrightText: Copyright The Lance Authors |
| 3 | + |
| 4 | +//! `rrf` — Reciprocal Rank Fusion (Cormack, Clarke & Büttcher, SIGIR 2009). |
| 5 | +//! |
| 6 | +//! Fuses N independently-ranked result lists into one ranking, scoring each id |
| 7 | +//! by `Σ_lists 1 / (k + rank)` where `rank` is the **1-based position** of the |
| 8 | +//! id in that list. RRF fuses by *rank position*, never by the source scores — |
| 9 | +//! which is exactly why it combines lists whose scores are **not |
| 10 | +//! commensurable**: [`Bm25Index::rank`](super::bm25::Bm25Index::rank) (tf-idf |
| 11 | +//! `f64`), [`PersonalizedPageRank::ranked`](super::ppr::PersonalizedPageRank) |
| 12 | +//! (unit-sum probability), and CAM-PQ (`i8` distance) share no scale, yet their |
| 13 | +//! rank orders fuse cleanly. `k` (default [`DEFAULT_RRF_K`] = 60, the paper's |
| 14 | +//! constant) damps deep-ranked items so a strong agreement near the top of |
| 15 | +//! several lists dominates a lone top-1 in one list. |
| 16 | +//! |
| 17 | +//! This is the fusion primitive named as the D-GR-2 retrieval **keystone** in |
| 18 | +//! `.claude/knowledge/graphrag-representations-inventory.md` (the SAP "Practical |
| 19 | +//! GraphRAG" reader's headline gap: every ranked leg exists — |
| 20 | +//! `Bm25Index`/`PersonalizedPageRank`/CAM-PQ — but nothing fused them). It is a |
| 21 | +//! **pure, reversible** capability: it computes a fused ranking and reads no |
| 22 | +//! carrier state. Wiring it into `OsintRetriever::retrieve` (so the retriever |
| 23 | +//! actually fuses its legs) stays gated on the G0 load-bearing verdict — this |
| 24 | +//! module only lands the algorithm, ahead of that gate, exactly as |
| 25 | +//! `Bm25Index`/`PersonalizedPageRank`/`Communities` landed as pure capabilities. |
| 26 | +
|
| 27 | +use std::collections::{BTreeMap, BTreeSet}; |
| 28 | + |
| 29 | +use lance_graph_contract::doc_graph::ScoredId; |
| 30 | + |
| 31 | +/// The paper's default rank-fusion constant (`k = 60`). |
| 32 | +pub const DEFAULT_RRF_K: f64 = 60.0; |
| 33 | + |
| 34 | +/// Fuse `ranked_lists` — each already ordered **best-first** — into one ranking |
| 35 | +/// by Reciprocal Rank Fusion. |
| 36 | +/// |
| 37 | +/// Each id's fused score is `Σ_lists 1 / (k + rank)` with `rank` 1-based. The |
| 38 | +/// returned [`ScoredId`]s are sorted by fused score **descending**, ties broken |
| 39 | +/// by id **ascending** (deterministic). `depth` carries the *shallowest* depth |
| 40 | +/// the id appeared at across the inputs (strongest provenance wins); an id |
| 41 | +/// absent from every list simply does not appear. |
| 42 | +/// |
| 43 | +/// Because fusion is by RANK, the per-list [`ScoredId::score`] values need not |
| 44 | +/// be commensurable — the reason RRF can combine the BM25 / PPR / CAM-PQ legs, |
| 45 | +/// whose scores live on unrelated scales. |
| 46 | +/// |
| 47 | +/// `k` should be positive; the paper uses `60` ([`DEFAULT_RRF_K`]). An empty |
| 48 | +/// `ranked_lists` (or all-empty lists) yields an empty result. |
| 49 | +/// |
| 50 | +/// # Examples |
| 51 | +/// ``` |
| 52 | +/// use lance_graph::graph::arigraph::rrf::{reciprocal_rank_fusion, DEFAULT_RRF_K}; |
| 53 | +/// use lance_graph_contract::doc_graph::ScoredId; |
| 54 | +/// |
| 55 | +/// // `x` is rank-1 in BOTH lists; `y`/`z` are rank-2 in only one each. |
| 56 | +/// let list_a = [ScoredId::new("x", 1.0, 0), ScoredId::new("y", 1.0, 0)]; |
| 57 | +/// let list_b = [ScoredId::new("x", 1.0, 0), ScoredId::new("z", 1.0, 0)]; |
| 58 | +/// let fused = reciprocal_rank_fusion(&[&list_a, &list_b], DEFAULT_RRF_K); |
| 59 | +/// |
| 60 | +/// assert_eq!(fused.len(), 3); |
| 61 | +/// assert_eq!(fused[0].id, "x"); // consensus at the top wins |
| 62 | +/// ``` |
| 63 | +#[must_use] |
| 64 | +pub fn reciprocal_rank_fusion(ranked_lists: &[&[ScoredId]], k: f64) -> Vec<ScoredId> { |
| 65 | + // id -> (accumulated RRF score, shallowest depth seen). BTreeMap gives a |
| 66 | + // deterministic id-ascending iteration order, which the stable sort below |
| 67 | + // preserves within equal fused scores. |
| 68 | + let mut acc: BTreeMap<&str, (f64, u8)> = BTreeMap::new(); |
| 69 | + for list in ranked_lists { |
| 70 | + // RRF gives each list AT MOST ONE vote per id, at its best (first) rank. |
| 71 | + // A leg that surfaces the same entity more than once (e.g. several |
| 72 | + // relations to one node, before caller-side dedup) must not stack |
| 73 | + // `1/(k+rank)` contributions and swamp consensus across the other legs. |
| 74 | + let mut voted: BTreeSet<&str> = BTreeSet::new(); |
| 75 | + for (pos, item) in list.iter().enumerate() { |
| 76 | + let entry = acc.entry(item.id.as_str()).or_insert((0.0, u8::MAX)); |
| 77 | + entry.1 = entry.1.min(item.depth); // shallowest depth across all occurrences |
| 78 | + if voted.insert(item.id.as_str()) { |
| 79 | + // best-first order ⇒ the first occurrence is the best rank |
| 80 | + entry.0 += 1.0 / (k + pos as f64 + 1.0); |
| 81 | + } |
| 82 | + } |
| 83 | + } |
| 84 | + let mut fused: Vec<ScoredId> = acc |
| 85 | + .into_iter() |
| 86 | + .map(|(id, (score, depth))| { |
| 87 | + ScoredId::new(id, score as f32, if depth == u8::MAX { 0 } else { depth }) |
| 88 | + }) |
| 89 | + .collect(); |
| 90 | + // Fused score descending; `sort_by` is stable, so the BTreeMap's id-ascending |
| 91 | + // order breaks ties deterministically. |
| 92 | + fused.sort_by(|a, b| { |
| 93 | + b.score |
| 94 | + .partial_cmp(&a.score) |
| 95 | + .unwrap_or(std::cmp::Ordering::Equal) |
| 96 | + }); |
| 97 | + fused |
| 98 | +} |
| 99 | + |
| 100 | +#[cfg(test)] |
| 101 | +mod tests { |
| 102 | + use super::*; |
| 103 | + |
| 104 | + fn ids(v: &[ScoredId]) -> Vec<&str> { |
| 105 | + v.iter().map(|s| s.id.as_str()).collect() |
| 106 | + } |
| 107 | + |
| 108 | + #[test] |
| 109 | + fn consensus_near_top_beats_lone_top_one() { |
| 110 | + // `x` rank-1 in both; `y`,`z` rank-2 in one each. |
| 111 | + let a = [ScoredId::new("x", 1.0, 0), ScoredId::new("y", 1.0, 0)]; |
| 112 | + let b = [ScoredId::new("x", 1.0, 0), ScoredId::new("z", 1.0, 0)]; |
| 113 | + let fused = reciprocal_rank_fusion(&[&a, &b], DEFAULT_RRF_K); |
| 114 | + assert_eq!(fused.len(), 3); |
| 115 | + assert_eq!(fused[0].id, "x"); |
| 116 | + } |
| 117 | + |
| 118 | + #[test] |
| 119 | + fn fuses_incommensurable_scores_by_rank_only() { |
| 120 | + // Wildly different score scales; only the RANK order matters. |
| 121 | + let bm25 = [ScoredId::new("a", 9000.0, 0), ScoredId::new("b", 4000.0, 0)]; |
| 122 | + let ppr = [ScoredId::new("b", 0.51, 1), ScoredId::new("a", 0.49, 2)]; |
| 123 | + let fused = reciprocal_rank_fusion(&[&bm25, &ppr], DEFAULT_RRF_K); |
| 124 | + // a: 1/61 + 1/62 ; b: 1/62 + 1/61 — equal → deterministic id-asc order. |
| 125 | + assert_eq!(ids(&fused), ["a", "b"]); |
| 126 | + // score is symmetric, not dominated by bm25's huge raw magnitudes. |
| 127 | + assert!((fused[0].score - fused[1].score).abs() < 1e-6); |
| 128 | + } |
| 129 | + |
| 130 | + #[test] |
| 131 | + fn shallowest_depth_wins() { |
| 132 | + let a = [ScoredId::new("a", 1.0, 3)]; |
| 133 | + let b = [ScoredId::new("a", 1.0, 1)]; |
| 134 | + let fused = reciprocal_rank_fusion(&[&a, &b], DEFAULT_RRF_K); |
| 135 | + assert_eq!(fused.len(), 1); |
| 136 | + assert_eq!(fused[0].depth, 1); // min(3, 1) |
| 137 | + } |
| 138 | + |
| 139 | + #[test] |
| 140 | + fn rank_position_dominates_within_one_list() { |
| 141 | + let only = [ |
| 142 | + ScoredId::new("first", 0.1, 0), |
| 143 | + ScoredId::new("second", 0.1, 0), |
| 144 | + ScoredId::new("third", 0.1, 0), |
| 145 | + ]; |
| 146 | + let fused = reciprocal_rank_fusion(&[&only], DEFAULT_RRF_K); |
| 147 | + assert_eq!(ids(&fused), ["first", "second", "third"]); |
| 148 | + // strictly decreasing: 1/61 > 1/62 > 1/63 |
| 149 | + assert!(fused[0].score > fused[1].score && fused[1].score > fused[2].score); |
| 150 | + } |
| 151 | + |
| 152 | + #[test] |
| 153 | + fn smaller_k_sharpens_top_rank_advantage() { |
| 154 | + // A rank-1 hit vs a rank-10 hit: smaller k widens their score ratio. |
| 155 | + let mk = |k: f64| { |
| 156 | + let a = [ScoredId::new("top", 1.0, 0)]; |
| 157 | + let mut deep: Vec<ScoredId> = (0..10) |
| 158 | + .map(|i| ScoredId::new(format!("d{i}"), 1.0, 0)) |
| 159 | + .collect(); |
| 160 | + deep[9] = ScoredId::new("low", 1.0, 0); // "low" at rank 10 |
| 161 | + let f = reciprocal_rank_fusion(&[&a, &deep], k); |
| 162 | + let top = f.iter().find(|s| s.id == "top").unwrap().score; |
| 163 | + let low = f.iter().find(|s| s.id == "low").unwrap().score; |
| 164 | + top / low |
| 165 | + }; |
| 166 | + assert!(mk(10.0) > mk(60.0)); // smaller k → bigger top-vs-deep ratio |
| 167 | + } |
| 168 | + |
| 169 | + #[test] |
| 170 | + fn empty_inputs_are_safe() { |
| 171 | + assert!(reciprocal_rank_fusion(&[], DEFAULT_RRF_K).is_empty()); |
| 172 | + let empty: [ScoredId; 0] = []; |
| 173 | + assert!(reciprocal_rank_fusion(&[&empty, &empty], DEFAULT_RRF_K).is_empty()); |
| 174 | + } |
| 175 | + |
| 176 | + #[test] |
| 177 | + fn deterministic() { |
| 178 | + let a = [ScoredId::new("a", 1.0, 0), ScoredId::new("b", 1.0, 0)]; |
| 179 | + let b = [ScoredId::new("c", 1.0, 0), ScoredId::new("a", 1.0, 0)]; |
| 180 | + let x = reciprocal_rank_fusion(&[&a, &b], DEFAULT_RRF_K); |
| 181 | + let y = reciprocal_rank_fusion(&[&a, &b], DEFAULT_RRF_K); |
| 182 | + assert_eq!(ids(&x), ids(&y)); |
| 183 | + assert_eq!( |
| 184 | + x.iter().map(|s| s.score).collect::<Vec<_>>(), |
| 185 | + y.iter().map(|s| s.score).collect::<Vec<_>>() |
| 186 | + ); |
| 187 | + } |
| 188 | + |
| 189 | + #[test] |
| 190 | + fn duplicate_id_in_one_leg_votes_once() { |
| 191 | + // A leg surfaces "a" twice (rank 1 and rank 3) plus "b" at rank 2 — as a |
| 192 | + // node with several relations would before caller-side dedup. RRF must |
| 193 | + // credit "a" ONCE at its best rank (1), not 1/(k+1)+1/(k+3). |
| 194 | + let leg = [ |
| 195 | + ScoredId::new("a", 1.0, 0), |
| 196 | + ScoredId::new("b", 1.0, 0), |
| 197 | + ScoredId::new("a", 1.0, 0), |
| 198 | + ]; |
| 199 | + let fused = reciprocal_rank_fusion(&[&leg], DEFAULT_RRF_K); |
| 200 | + assert_eq!(fused.len(), 2); |
| 201 | + let a = fused.iter().find(|s| s.id == "a").unwrap().score; |
| 202 | + let b = fused.iter().find(|s| s.id == "b").unwrap().score; |
| 203 | + let expected_a = (1.0f64 / (DEFAULT_RRF_K + 1.0)) as f32; // best rank only |
| 204 | + assert!( |
| 205 | + (a - expected_a).abs() < 1e-7, |
| 206 | + "a double-counted: {a} vs {expected_a}" |
| 207 | + ); |
| 208 | + assert!(a > b); // single best-rank vote still beats b's rank-2 vote |
| 209 | + } |
| 210 | + |
| 211 | + #[test] |
| 212 | + fn duplicate_id_still_folds_shallowest_depth() { |
| 213 | + // Duplicate occurrences don't re-vote, but depth still takes the min. |
| 214 | + let leg = [ScoredId::new("a", 1.0, 5), ScoredId::new("a", 1.0, 2)]; |
| 215 | + let fused = reciprocal_rank_fusion(&[&leg], DEFAULT_RRF_K); |
| 216 | + assert_eq!(fused.len(), 1); |
| 217 | + assert_eq!(fused[0].depth, 2); // min(5, 2), even though the 2nd didn't vote |
| 218 | + } |
| 219 | +} |
0 commit comments