mirror of
https://git.pumpin.online/funkpower/catena.git
synced 2026-10-23 03:07:13 +00:00
rrf: rank fusion without the database
This commit is contained in:
funkpower
@@ -0,0 +1,16 @@
|
||||
[package]
|
||||
name = "rrf"
|
||||
description = "Reciprocal rank fusion: several ranked lists in, one ranked list out."
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
rust-version.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
clap.workspace = true
|
||||
corpus-core.workspace = true
|
||||
serde_json.workspace = true
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
@@ -0,0 +1,220 @@
|
||||
//! Reciprocal rank fusion: several ranked lists in, one ranked list out.
|
||||
//!
|
||||
//! The useful half of hybrid search, with none of the database. Vector search
|
||||
//! and full-text search disagree about everything except what matters, and
|
||||
//! their scores are not comparable — a cosine distance and a `ts_rank` are
|
||||
//! different quantities on different scales. Fusion sidesteps that by
|
||||
//! throwing the scores away and keeping only the ranks: a result that both
|
||||
//! channels put near the top beats one that either channel loves alone.
|
||||
//!
|
||||
//! A channel may be weighted, which is how a metadata match outranks a
|
||||
//! body-text tie: when somebody searches for a book by its title, the book
|
||||
//! rarely repeats its own title in its body, so the channel that matched the
|
||||
//! title is the one telling the truth.
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use corpus_core::{Channel, Hit};
|
||||
|
||||
/// The constant that flattens the difference between rank 1 and rank 2.
|
||||
///
|
||||
/// Sixty, from the paper the method comes from and from every implementation
|
||||
/// since. It is large enough that the top of one list does not swamp the rest
|
||||
/// of another, and it is the number every other system's fused scores are
|
||||
/// comparable against, which is worth more than tuning it.
|
||||
pub const K: f64 = 60.0;
|
||||
|
||||
/// How much a metadata match counts for, against a body-text match at the
|
||||
/// same rank.
|
||||
pub const META_WEIGHT: f64 = 1.5;
|
||||
|
||||
/// What each channel's rank is worth.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct Weights {
|
||||
/// The rank constant.
|
||||
pub k: f64,
|
||||
/// Per-channel weights.
|
||||
pub channels: HashMap<Channel, f64>,
|
||||
}
|
||||
|
||||
impl Default for Weights {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
k: K,
|
||||
channels: HashMap::from([
|
||||
(Channel::Vector, 1.0),
|
||||
(Channel::Fts, 1.0),
|
||||
(Channel::Meta, META_WEIGHT),
|
||||
]),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Weights {
|
||||
/// What one channel's rank is worth. An unnamed channel counts for one.
|
||||
pub fn of(&self, channel: Channel) -> f64 {
|
||||
self.channels.get(&channel).copied().unwrap_or(1.0)
|
||||
}
|
||||
}
|
||||
|
||||
/// Fuse ranked lists into one.
|
||||
///
|
||||
/// Every hit is grouped by its `id`, and each contributes `weight / (k +
|
||||
/// rank)` — the rank the producer gave it, not its position in the stream, so
|
||||
/// the lists may arrive interleaved. The record that comes back for an id is
|
||||
/// the one that ranked it best, carrying its payload, with `rank` and `score`
|
||||
/// replaced by its place and score in the fused list.
|
||||
///
|
||||
/// Ties keep the order the ids were first seen in, so the same input always
|
||||
/// produces the same output.
|
||||
pub fn fuse(hits: &[Hit], weights: &Weights) -> Vec<Hit> {
|
||||
let mut order: Vec<&str> = Vec::new();
|
||||
let mut fused: HashMap<&str, Fusing<'_>> = HashMap::new();
|
||||
|
||||
for hit in hits {
|
||||
let entry = fused.entry(hit.id.as_str()).or_insert_with(|| {
|
||||
order.push(hit.id.as_str());
|
||||
Fusing {
|
||||
best: hit,
|
||||
score: 0.0,
|
||||
seen: HashMap::new(),
|
||||
}
|
||||
});
|
||||
// A channel that ranks the same id twice is a producer contradicting
|
||||
// itself; the later word is taken as the corrected one.
|
||||
if let Some(previous) = entry.seen.insert(hit.channel, hit.rank) {
|
||||
entry.score -= weights.of(hit.channel) / (weights.k + previous as f64);
|
||||
}
|
||||
entry.score += weights.of(hit.channel) / (weights.k + hit.rank as f64);
|
||||
if hit.rank < entry.best.rank {
|
||||
entry.best = hit;
|
||||
}
|
||||
}
|
||||
|
||||
let mut ranked: Vec<(&str, f64)> = order
|
||||
.iter()
|
||||
.map(|id| (*id, fused[id].score))
|
||||
.collect::<Vec<_>>();
|
||||
// Stable, so equal scores stay in first-seen order.
|
||||
ranked.sort_by(|a, b| b.1.total_cmp(&a.1));
|
||||
|
||||
ranked
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
.map(|(at, (id, score))| {
|
||||
let mut hit = fused[id].best.clone();
|
||||
hit.rank = at as u32 + 1;
|
||||
hit.score = score;
|
||||
hit
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
struct Fusing<'a> {
|
||||
best: &'a Hit,
|
||||
score: f64,
|
||||
seen: HashMap<Channel, u32>,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn hit(id: &str, channel: Channel, rank: u32) -> Hit {
|
||||
Hit::new(id, channel, rank, 0.0)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_result_both_channels_like_beats_one_either_channel_loves() {
|
||||
let hits = [
|
||||
hit("both", Channel::Vector, 2),
|
||||
hit("vector-only", Channel::Vector, 1),
|
||||
hit("both", Channel::Fts, 2),
|
||||
hit("fts-only", Channel::Fts, 1),
|
||||
];
|
||||
let fused = fuse(&hits, &Weights::default());
|
||||
assert_eq!(fused[0].id, "both");
|
||||
assert_eq!(fused[0].rank, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_metadata_match_outranks_a_body_tie() {
|
||||
let hits = [hit("title", Channel::Meta, 1), hit("body", Channel::Fts, 1)];
|
||||
let fused = fuse(&hits, &Weights::default());
|
||||
assert_eq!(fused[0].id, "title");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_scores_are_the_ones_the_method_defines() {
|
||||
let hits = [hit("a", Channel::Vector, 1), hit("a", Channel::Fts, 3)];
|
||||
let fused = fuse(&hits, &Weights::default());
|
||||
assert_eq!(fused[0].score, 1.0 / 61.0 + 1.0 / 63.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_record_that_comes_back_is_the_one_that_ranked_best() {
|
||||
let mut best = hit("a", Channel::Vector, 1);
|
||||
best.payload.insert("from".into(), "vector".into());
|
||||
let mut worse = hit("a", Channel::Fts, 9);
|
||||
worse.payload.insert("from".into(), "fts".into());
|
||||
|
||||
let fused = fuse(&[worse, best], &Weights::default());
|
||||
assert_eq!(fused[0].payload["from"], "vector");
|
||||
assert_eq!(fused[0].channel, Channel::Vector);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ranks_come_out_consecutive_from_one() {
|
||||
let hits = [
|
||||
hit("a", Channel::Vector, 1),
|
||||
hit("b", Channel::Vector, 2),
|
||||
hit("c", Channel::Vector, 3),
|
||||
];
|
||||
let fused = fuse(&hits, &Weights::default());
|
||||
assert_eq!(
|
||||
fused.iter().map(|hit| hit.rank).collect::<Vec<_>>(),
|
||||
[1, 2, 3]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_tie_keeps_the_order_the_ids_arrived_in() {
|
||||
let hits = [
|
||||
hit("second", Channel::Vector, 5),
|
||||
hit("first", Channel::Fts, 5),
|
||||
];
|
||||
let fused = fuse(&hits, &Weights::default());
|
||||
assert_eq!(fused[0].id, "second");
|
||||
assert_eq!(fused[1].id, "first");
|
||||
assert_eq!(fused[0].score, fused[1].score);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_channel_that_ranks_the_same_id_twice_counts_once() {
|
||||
let twice = fuse(
|
||||
&[hit("a", Channel::Vector, 9), hit("a", Channel::Vector, 1)],
|
||||
&Weights::default(),
|
||||
);
|
||||
let once = fuse(&[hit("a", Channel::Vector, 1)], &Weights::default());
|
||||
assert_eq!(twice[0].score, once[0].score);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn weights_can_be_changed() {
|
||||
let weights = Weights {
|
||||
k: 1.0,
|
||||
channels: HashMap::from([(Channel::Fts, 10.0)]),
|
||||
};
|
||||
let fused = fuse(
|
||||
&[hit("meta", Channel::Meta, 1), hit("fts", Channel::Fts, 1)],
|
||||
&weights,
|
||||
);
|
||||
assert_eq!(fused[0].id, "fts");
|
||||
assert_eq!(fused[0].score, 10.0 / 2.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn nothing_in_nothing_out() {
|
||||
assert!(fuse(&[], &Weights::default()).is_empty());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,138 @@
|
||||
//! Reciprocal rank fusion: several ranked lists in, one ranked list out.
|
||||
|
||||
use std::io::{BufWriter, Read, Write};
|
||||
use std::path::PathBuf;
|
||||
use std::process::ExitCode;
|
||||
|
||||
use clap::Parser;
|
||||
use corpus_core::{Channel, Hit, exit, ndjson};
|
||||
use rrf::Weights;
|
||||
|
||||
/// Fuse several ranked lists of hits into one.
|
||||
///
|
||||
/// Reads hit records on stdin, in any order and interleaved if you like —
|
||||
/// each carries the rank its own channel gave it, so the stream does not have
|
||||
/// to be grouped. Writes the fused list, best first, with `rank` and `score`
|
||||
/// replaced by the fused ones.
|
||||
///
|
||||
/// Database-agnostic on purpose. By the time a hit reaches here it is already
|
||||
/// yours, so this is the one tool in the suite that is worth calling in
|
||||
/// process rather than as a subprocess.
|
||||
#[derive(Parser)]
|
||||
#[command(
|
||||
version,
|
||||
about,
|
||||
after_help = "\
|
||||
Example:
|
||||
|
||||
cat vector.ndjson fts.ndjson | rrf --limit 10
|
||||
|
||||
Weights are given as channel=number, and default to vec=1, fts=1, meta=1.5 —
|
||||
a metadata match outranks a body-text match at the same rank, because a book
|
||||
rarely repeats its own title in its body."
|
||||
)]
|
||||
struct Args {
|
||||
/// Hit stream to read. Reads stdin when absent.
|
||||
path: Option<PathBuf>,
|
||||
|
||||
/// How many results to keep. All of them when absent.
|
||||
#[arg(long, short = 'n', value_name = "COUNT")]
|
||||
limit: Option<usize>,
|
||||
|
||||
/// The rank constant.
|
||||
#[arg(long = "rrf-k", default_value_t = rrf::K, value_name = "K")]
|
||||
k: f64,
|
||||
|
||||
/// What a channel's rank is worth, as `channel=number`. Repeatable.
|
||||
#[arg(long, value_name = "CHANNEL=WEIGHT")]
|
||||
weight: Vec<String>,
|
||||
}
|
||||
|
||||
fn main() -> ExitCode {
|
||||
let args = Args::parse();
|
||||
|
||||
let weights = match weights(&args) {
|
||||
Ok(weights) => weights,
|
||||
Err(why) => {
|
||||
eprintln!("rrf: {why}");
|
||||
return ExitCode::from(exit::USAGE);
|
||||
}
|
||||
};
|
||||
|
||||
let hits: Result<Vec<Hit>, _> = match args.path.as_deref() {
|
||||
Some(path) => match std::fs::File::open(path) {
|
||||
Ok(file) => ndjson::read(std::io::BufReader::new(file)).collect(),
|
||||
Err(e) => {
|
||||
eprintln!("rrf: {}: {e}", path.display());
|
||||
return ExitCode::from(exit::IO);
|
||||
}
|
||||
},
|
||||
None => {
|
||||
let mut text = Vec::new();
|
||||
if let Err(e) = std::io::stdin().read_to_end(&mut text) {
|
||||
eprintln!("rrf: {e}");
|
||||
return ExitCode::from(exit::IO);
|
||||
}
|
||||
ndjson::read(text.as_slice()).collect()
|
||||
}
|
||||
};
|
||||
let hits = match hits {
|
||||
Ok(hits) => hits,
|
||||
Err(e) => {
|
||||
eprintln!("rrf: {e}");
|
||||
return ExitCode::from(if e.is_broken_pipe() {
|
||||
exit::OK
|
||||
} else {
|
||||
exit::DATA
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
let fused = rrf::fuse(&hits, &weights);
|
||||
let fused = &fused[..args.limit.unwrap_or(fused.len()).min(fused.len())];
|
||||
|
||||
let mut out = BufWriter::new(std::io::stdout().lock());
|
||||
for hit in fused {
|
||||
if let Err(e) = ndjson::write_line(&mut out, hit) {
|
||||
// A consumer that stopped reading is not a failure.
|
||||
if e.is_broken_pipe() {
|
||||
return ExitCode::from(exit::OK);
|
||||
}
|
||||
eprintln!("rrf: {e}");
|
||||
return ExitCode::from(exit::IO);
|
||||
}
|
||||
}
|
||||
match out.flush() {
|
||||
Ok(()) => ExitCode::from(exit::OK),
|
||||
Err(e) if e.kind() == std::io::ErrorKind::BrokenPipe => ExitCode::from(exit::OK),
|
||||
Err(e) => {
|
||||
eprintln!("rrf: {e}");
|
||||
ExitCode::from(exit::IO)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn weights(args: &Args) -> Result<Weights, String> {
|
||||
if !args.k.is_finite() || args.k <= 0.0 {
|
||||
return Err(format!("--rrf-k {} must be a positive number", args.k));
|
||||
}
|
||||
let mut weights = Weights {
|
||||
k: args.k,
|
||||
..Weights::default()
|
||||
};
|
||||
for given in &args.weight {
|
||||
let (channel, weight) = given
|
||||
.split_once('=')
|
||||
.ok_or_else(|| format!("--weight {given:?} is not channel=number"))?;
|
||||
let channel: Channel = serde_json::from_value(serde_json::Value::String(channel.into()))
|
||||
.map_err(|_| format!("--weight {given:?} names no channel this knows"))?;
|
||||
let weight: f64 = weight
|
||||
.parse()
|
||||
.map_err(|_| format!("--weight {given:?} is not channel=number"))?;
|
||||
if !weight.is_finite() {
|
||||
return Err(format!("--weight {given:?} is not a number"));
|
||||
}
|
||||
weights.channels.insert(channel, weight);
|
||||
}
|
||||
Ok(weights)
|
||||
}
|
||||
Reference in New Issue
Block a user