Add multi-query lexical skill selection (#33614)

## What changed

- Add a lexical selector that splits compound queries into sentence and connector-delimited views.
- Merge candidates across the full query and its component views, prioritizing their best rank while using full-query rank, view coverage, and stable identifiers as tie-breakers.
- Bound query decomposition and result counts, propagate truncation signals, and include the selector in the shadow-selection experiment.

## Testing

- Add unit coverage for compound-query ranking, single-query compatibility, query splitting, and bounded-input signals.

GitOrigin-RevId: 2442bf7198fc24295cbc270b9d222275ae8a7cc3
This commit is contained in:
jif
2026-07-16 14:42:41 +00:00
committed by copyberry
parent c983a53f20
commit 0f44bca915
4 changed files with 260 additions and 0 deletions
@@ -1,8 +1,10 @@
mod character_ngram;
mod fielded_bm25;
mod multi_query_lexical;
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 weighted_lexical::WeightedLexicalSkillSelector;
/// Metadata searched by a cheap skill selector.
@@ -0,0 +1,166 @@
use std::collections::HashMap;
use std::collections::HashSet;
use super::CheapSkillSelection;
use super::CheapSkillSelector;
use super::SkillSelectionDocument;
use super::WeightedLexicalSkillSelector;
const MAX_QUERY_VIEWS: usize = 8;
const MAX_DECOMPOSITION_BYTES: usize = 4 * 1024;
const MAX_RESULTS: usize = 50;
const CONNECTORS: &[&str] = &[" and then ", " and ", " then ", " also "];
#[derive(Clone, Copy, Debug, Default)]
pub(crate) struct MultiQueryLexicalSkillSelector;
impl CheapSkillSelector for MultiQueryLexicalSkillSelector {
fn method(&self) -> &'static str {
"multi_query_lexical_v1"
}
fn select(
&self,
query: &str,
documents: &[SkillSelectionDocument<'_>],
limit: usize,
) -> CheapSkillSelection {
let full_selection =
WeightedLexicalSkillSelector.select(query, documents, limit.min(MAX_RESULTS));
let views = query_views(query);
if views.len() <= 1 || limit == 0 {
return full_selection;
}
let mut candidates = HashMap::new();
record_candidates(&mut candidates, &full_selection, /*view_index*/ 0);
let mut query_truncated = full_selection.query_truncated;
let mut candidate_set_truncated = full_selection.candidate_set_truncated;
for (view_index, view) in views.into_iter().enumerate().skip(1) {
let selection =
WeightedLexicalSkillSelector.select(view, documents, limit.min(MAX_RESULTS));
query_truncated |= selection.query_truncated;
candidate_set_truncated |= selection.candidate_set_truncated;
record_candidates(&mut candidates, &selection, view_index);
}
let mut candidates = candidates.into_values().collect::<Vec<_>>();
candidates.sort_by(|left, right| {
left.best_rank
.cmp(&right.best_rank)
.then_with(|| {
left.full_query_rank
.unwrap_or(usize::MAX)
.cmp(&right.full_query_rank.unwrap_or(usize::MAX))
})
.then_with(|| right.view_count.cmp(&left.view_count))
.then_with(|| left.first_view.cmp(&right.first_view))
.then_with(|| left.id.cmp(&right.id))
});
CheapSkillSelection {
candidate_ids: candidates
.into_iter()
.take(limit.min(MAX_RESULTS))
.map(|candidate| candidate.id)
.collect(),
query_term_count: full_selection.query_term_count,
query_truncated,
candidate_set_truncated,
}
}
}
struct RankedCandidate {
id: usize,
best_rank: usize,
full_query_rank: Option<usize>,
view_count: usize,
first_view: usize,
}
fn record_candidates(
candidates: &mut HashMap<usize, RankedCandidate>,
selection: &CheapSkillSelection,
view_index: usize,
) {
for (rank, id) in selection.candidate_ids.iter().copied().enumerate() {
let rank = rank + 1;
candidates
.entry(id)
.and_modify(|candidate| {
candidate.best_rank = candidate.best_rank.min(rank);
candidate.view_count = candidate.view_count.saturating_add(1);
if view_index == 0 {
candidate.full_query_rank = Some(rank);
}
})
.or_insert(RankedCandidate {
id,
best_rank: rank,
full_query_rank: (view_index == 0).then_some(rank),
view_count: 1,
first_view: view_index,
});
}
}
fn query_views(query: &str) -> Vec<&str> {
let full_query = bounded(query, MAX_DECOMPOSITION_BYTES).trim();
if full_query.is_empty() {
return Vec::new();
}
let mut views = vec![full_query];
let mut seen = HashSet::from([full_query]);
for sentence in full_query.split(['\n', '\r', '.', '!', '?', ';']) {
for clause in split_connectors(sentence) {
let clause = clause.trim();
if clause.chars().count() >= 2 && seen.insert(clause) {
views.push(clause);
if views.len() == MAX_QUERY_VIEWS {
return views;
}
}
}
}
views
}
fn bounded(value: &str, max_bytes: usize) -> &str {
if value.len() <= max_bytes {
return value;
}
let mut end = max_bytes;
while !value.is_char_boundary(end) {
end = end.saturating_sub(1);
}
&value[..end]
}
fn split_connectors(value: &str) -> Vec<&str> {
let lowercase = value.to_ascii_lowercase();
let mut segments = Vec::new();
let mut start = 0;
while start < value.len() {
let next = CONNECTORS
.iter()
.filter_map(|connector| {
lowercase[start..]
.find(connector)
.map(|offset| (start + offset, connector.len()))
})
.min_by_key(|(position, _)| *position);
let Some((position, connector_length)) = next else {
break;
};
segments.push(&value[start..position]);
start = position + connector_length;
}
segments.push(&value[start..]);
segments
}
#[cfg(test)]
#[path = "multi_query_lexical_tests.rs"]
mod tests;
@@ -0,0 +1,90 @@
use super::*;
use pretty_assertions::assert_eq;
#[test]
fn multi_query_promotes_each_clause_leader() {
let documents = [
document(/*id*/ 1, "rust-format", "Format Rust source code."),
document(/*id*/ 2, "rust-lint", "Fix Rust source lint errors."),
document(/*id*/ 3, "rust-review", "Review Rust source code."),
document(
/*id*/ 4,
"ci-fix",
"Diagnose failing GitHub Actions checks.",
),
];
let selection = MultiQueryLexicalSkillSelector.select(
"format and review Rust source code, and then diagnose failing GitHub Actions checks",
&documents,
/*limit*/ 4,
);
assert!(selection.candidate_ids[..3].contains(&1));
assert!(selection.candidate_ids[..3].contains(&4));
}
#[test]
fn single_query_matches_the_underlying_selector() {
let documents = [
document(/*id*/ 1, "presentations", "Create visual decks."),
document(/*id*/ 2, "spreadsheets", "Analyze tabular data."),
];
let expected =
WeightedLexicalSkillSelector.select("create presentations", &documents, /*limit*/ 20);
let actual = MultiQueryLexicalSkillSelector.select(
"create presentations",
&documents,
/*limit*/ 20,
);
assert_eq!(expected, actual);
}
#[test]
fn query_views_split_sentences_and_connectors() {
assert_eq!(
vec![
"format code and then fix tests; write a summary",
"format code",
"fix tests",
"write a summary",
],
query_views("format code and then fix tests; write a summary"),
);
}
#[test]
fn multi_query_preserves_bounded_input_signals() {
let long_query = format!("{}\nand inspect logs", "match ".repeat(4 * 1024));
let names = (0..=1_000)
.map(|index| format!("candidate-{index}"))
.collect::<Vec<_>>();
let documents = names
.iter()
.enumerate()
.map(|(id, name)| SkillSelectionDocument {
id,
name,
short_description: None,
description: "match logs",
})
.collect::<Vec<_>>();
let selection =
MultiQueryLexicalSkillSelector.select(&long_query, &documents, /*limit*/ 20);
assert!(selection.query_truncated);
assert!(selection.candidate_set_truncated);
assert_eq!(20, selection.candidate_ids.len());
}
fn document<'a>(id: usize, name: &'a str, description: &'a str) -> SkillSelectionDocument<'a> {
SkillSelectionDocument {
id,
name,
short_description: None,
description,
}
}
@@ -15,6 +15,7 @@ use crate::dynamic_skill_selector::CharacterNgramSkillSelector;
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::SkillSelectionDocument;
use crate::dynamic_skill_selector::WeightedLexicalSkillSelector;
@@ -41,6 +42,7 @@ impl ShadowSelectionExperiment {
Box::new(WeightedLexicalSkillSelector),
Box::new(FieldedBm25SkillSelector),
Box::new(CharacterNgramSkillSelector),
Box::new(MultiQueryLexicalSkillSelector),
],
metrics_client,
}