Add reciprocal rank fusion skill selection (#34547)

## What changed

- Add a skill selector that uses reciprocal rank fusion to combine the top 50 results from the weighted lexical and character n-gram selectors.
- Prefer candidates supported by both rankings, with deterministic rank and identifier tie-breakers, and propagate truncation metadata from both selectors.
- Include the fused selector in the shadow-selection experiment under `rrf_lexical_char_v1`.

## Testing

- Add unit coverage for cross-ranking support, deterministic tie-breaking, and combined input-bound reporting.

GitOrigin-RevId: 3165136b8729ea3a9484f25c7f2ce0f91fb8aaa9
This commit is contained in:
jif
2026-07-21 13:49:22 +00:00
committed by copyberry
parent 2d5c259202
commit 6915bac7ba
4 changed files with 123 additions and 0 deletions
@@ -1,10 +1,12 @@
mod character_ngram;
mod fielded_bm25;
mod multi_query_lexical;
mod rrf_lexical_char;
mod weighted_lexical;
pub(crate) use character_ngram::CharacterNgramSkillSelector;
pub(crate) use fielded_bm25::FieldedBm25SkillSelector;
pub(crate) use multi_query_lexical::MultiQueryLexicalSkillSelector;
pub(crate) use rrf_lexical_char::RrfLexicalCharSkillSelector;
pub(crate) use weighted_lexical::WeightedLexicalSkillSelector;
/// Metadata searched by a cheap skill selector.
@@ -0,0 +1,85 @@
use std::collections::HashMap;
use super::CharacterNgramSkillSelector;
use super::CheapSkillSelection;
use super::CheapSkillSelector;
use super::SkillSelectionDocument;
use super::WeightedLexicalSkillSelector;
const MAX_FUSION_CANDIDATES: usize = 50;
const RRF_K: usize = 60;
#[derive(Clone, Copy, Debug, Default)]
pub(crate) struct RrfLexicalCharSkillSelector;
impl CheapSkillSelector for RrfLexicalCharSkillSelector {
fn method(&self) -> &'static str {
"rrf_lexical_char_v1"
}
fn select(
&self,
query: &str,
documents: &[SkillSelectionDocument<'_>],
limit: usize,
) -> CheapSkillSelection {
if limit == 0 {
return CheapSkillSelection::default();
}
let lexical = WeightedLexicalSkillSelector.select(
query,
documents,
/*limit*/ MAX_FUSION_CANDIDATES,
);
let character = CharacterNgramSkillSelector.select(
query,
documents,
/*limit*/ MAX_FUSION_CANDIDATES,
);
let candidate_ids = fuse_rankings(
[&lexical.candidate_ids, &character.candidate_ids],
limit.min(MAX_FUSION_CANDIDATES),
);
CheapSkillSelection {
candidate_ids,
query_term_count: lexical.query_term_count.max(character.query_term_count),
query_truncated: lexical.query_truncated || character.query_truncated,
candidate_set_truncated: lexical.candidate_set_truncated
|| character.candidate_set_truncated,
}
}
}
fn fuse_rankings<const N: usize>(rankings: [&[usize]; N], limit: usize) -> Vec<usize> {
let mut scores = HashMap::<usize, f64>::new();
let mut best_ranks = HashMap::<usize, usize>::new();
for ranking in rankings {
for (index, id) in ranking.iter().copied().enumerate() {
let rank = index + 1;
*scores.entry(id).or_default() += 1.0 / (RRF_K + rank) as f64;
best_ranks
.entry(id)
.and_modify(|best_rank| *best_rank = (*best_rank).min(rank))
.or_insert(rank);
}
}
let mut candidates = scores.into_iter().collect::<Vec<_>>();
candidates.sort_by(|(left_id, left_score), (right_id, right_score)| {
right_score
.total_cmp(left_score)
.then_with(|| best_ranks[left_id].cmp(&best_ranks[right_id]))
.then_with(|| left_id.cmp(right_id))
});
candidates
.into_iter()
.take(limit)
.map(|(id, _)| id)
.collect()
}
#[cfg(test)]
#[path = "rrf_lexical_char_tests.rs"]
mod tests;
@@ -0,0 +1,34 @@
use pretty_assertions::assert_eq;
use super::*;
#[test]
fn fusion_prefers_candidates_supported_by_both_rankings() {
let fused = fuse_rankings([&[1, 2, 3], &[2, 3, 4]], /*limit*/ 4);
assert_eq!(vec![2, 3, 1, 4], fused);
}
#[test]
fn fusion_uses_rank_and_id_to_break_ties() {
let fused = fuse_rankings([&[20, 10], &[30]], /*limit*/ 3);
assert_eq!(vec![20, 30, 10], fused);
}
#[test]
fn selector_reports_the_combined_input_bounds() {
let documents = [SkillSelectionDocument {
id: 7,
name: "presentations",
short_description: None,
description: "Create visual decks.",
}];
let query = "presentation ".repeat(4 * 1024);
let selection = RrfLexicalCharSkillSelector.select(&query, &documents, /*limit*/ 20);
assert_eq!(vec![7], selection.candidate_ids);
assert!(selection.query_truncated);
assert!(!selection.candidate_set_truncated);
}
@@ -16,6 +16,7 @@ use crate::dynamic_skill_selector::CheapSkillSelection;
use crate::dynamic_skill_selector::CheapSkillSelector;
use crate::dynamic_skill_selector::FieldedBm25SkillSelector;
use crate::dynamic_skill_selector::MultiQueryLexicalSkillSelector;
use crate::dynamic_skill_selector::RrfLexicalCharSkillSelector;
use crate::dynamic_skill_selector::SkillSelectionDocument;
use crate::dynamic_skill_selector::WeightedLexicalSkillSelector;
@@ -43,6 +44,7 @@ impl ShadowSelectionExperiment {
Box::new(FieldedBm25SkillSelector),
Box::new(CharacterNgramSkillSelector),
Box::new(MultiQueryLexicalSkillSelector),
Box::new(RrfLexicalCharSkillSelector),
],
metrics_client,
}