osobhandClaude Fable 5.1
390a2e3836
Fusion min-max normalises over every keyword match, so hybrid_search asked BM25 for a ranked list of the whole corpus: a hash insert per posting, then a sort of every match, then the merge sorted every candidate again to keep k. - BM25Index::scores returns every match unsorted, accumulated in a dense array (contributions are strictly positive, so zero means untouched). search() is built on it with the bounded heap. - merge_vector_keyword partitions out its top k (select_nth) and orders only those, with the same score-then-id order. - Both hybrid paths use scores(). Rankings are identical (equivalence tests for both changes). p50 0.24 -> 0.07 ms (1K), 2.1 -> 0.49 ms (10K), 23 -> 4.65 ms (100K). The harness gains --fusion-study, which measured the alternative — capping the keyword pool — and found it changes the top-10 for most queries (overlap 0.83-0.92, different #1 for 10-35%) for only a 2x saving. Not adopted. Co-Authored-By: Claude Fable 5.1 <[email protected]>
610 lines
22 KiB
Rust
610 lines
22 KiB
Rust
//! BM25 keyword search engine.
|
|
//!
|
|
//! Provides a standard BM25 (Okapi BM25) implementation with an in-memory
|
|
//! inverted index. Tombstoned documents are excluded from indexing and search.
|
|
//!
|
|
//! The index is **incremental**: [`BM25Index::add_document`] and
|
|
//! [`BM25Index::remove_document`] keep it exactly equivalent to one built from
|
|
//! scratch over the same live documents, so a store can maintain one index for
|
|
//! its lifetime instead of re-tokenising the whole corpus per query. To make
|
|
//! that possible IDF is computed at query time (it depends on the live
|
|
//! document count) rather than cached at build time.
|
|
//!
|
|
//! - Posting lists sorted by doc id
|
|
//! - Bounded-heap top-k; results ordered by score, then doc id (deterministic)
|
|
|
|
use std::cmp::Reverse;
|
|
use std::collections::{BinaryHeap, HashMap};
|
|
|
|
/// `f32` wrapper providing a total order (via `total_cmp`) so BM25 scores can
|
|
/// be kept in a `BinaryHeap`. Scores are always finite in practice (no NaN
|
|
/// inputs reach this path), so `total_cmp`'s NaN ordering is never exercised.
|
|
#[derive(Debug, Clone, Copy, PartialEq)]
|
|
struct HeapScore(f32);
|
|
|
|
impl Eq for HeapScore {}
|
|
|
|
impl PartialOrd for HeapScore {
|
|
fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
|
|
Some(self.cmp(other))
|
|
}
|
|
}
|
|
|
|
impl Ord for HeapScore {
|
|
fn cmp(&self, other: &Self) -> std::cmp::Ordering {
|
|
self.0.total_cmp(&other.0)
|
|
}
|
|
}
|
|
|
|
/// Default BM25 term-frequency saturation parameter.
|
|
const DEFAULT_K1: f32 = 1.2;
|
|
|
|
/// Default BM25 document-length normalization parameter.
|
|
const DEFAULT_B: f32 = 0.75;
|
|
|
|
/// An in-memory BM25 index for keyword search.
|
|
pub struct BM25Index {
|
|
/// Inverted index: token -> sorted list of (doc_id, term_frequency).
|
|
inverted: HashMap<String, Vec<(usize, u32)>>,
|
|
/// Number of tokens in each document (0 for tombstoned docs).
|
|
doc_lengths: Vec<u32>,
|
|
/// Sum of `doc_lengths` over live documents (keeps `avg_dl` exact under
|
|
/// incremental updates).
|
|
total_length: u64,
|
|
/// Average document length across non-tombstoned docs.
|
|
avg_dl: f32,
|
|
/// Number of non-tombstoned documents.
|
|
num_docs: usize,
|
|
/// BM25 k1 parameter.
|
|
k1: f32,
|
|
/// BM25 b parameter.
|
|
b: f32,
|
|
}
|
|
|
|
impl BM25Index {
|
|
/// Build a BM25 index from a set of documents, excluding tombstoned entries.
|
|
pub fn build(documents: &[String], tombstones: &[u8]) -> Self {
|
|
let mut index = Self {
|
|
inverted: HashMap::new(),
|
|
doc_lengths: vec![0; documents.len()],
|
|
total_length: 0,
|
|
avg_dl: 0.0,
|
|
num_docs: 0,
|
|
k1: DEFAULT_K1,
|
|
b: DEFAULT_B,
|
|
};
|
|
index.index_documents(documents, tombstones);
|
|
index
|
|
}
|
|
|
|
/// Search the index for a query, returning the top `k` results
|
|
/// as `(doc_id, score)` pairs sorted by score descending.
|
|
///
|
|
/// Uses Block-Max WAND for early termination when remaining documents
|
|
/// cannot beat the current top-k threshold.
|
|
pub fn search(&self, query: &str, k: usize) -> Vec<(usize, f32)> {
|
|
if k == 0 {
|
|
return Vec::new();
|
|
}
|
|
// Top-k with a bounded min-heap: O(matches * log k) instead of sorting
|
|
// every match. Ties break towards the lower doc id so results are
|
|
// deterministic.
|
|
let mut heap: BinaryHeap<Reverse<(HeapScore, Reverse<usize>)>> =
|
|
BinaryHeap::with_capacity(k.min(1024) + 1);
|
|
for (doc_id, score) in self.scores(query) {
|
|
heap.push(Reverse((HeapScore(score), Reverse(doc_id))));
|
|
if heap.len() > k {
|
|
heap.pop();
|
|
}
|
|
}
|
|
let mut results: Vec<(usize, f32)> = heap
|
|
.into_iter()
|
|
.map(|Reverse((HeapScore(score), Reverse(doc_id)))| (doc_id, score))
|
|
.collect();
|
|
results.sort_by(|a, b| b.1.total_cmp(&a.1).then(a.0.cmp(&b.0)));
|
|
results
|
|
}
|
|
|
|
/// The BM25 score of **every** matching document, in doc-id order, unsorted
|
|
/// by score. Score fusion normalises over the whole matching set, so it
|
|
/// needs all of these but not their ranking; producing a ranked list of
|
|
/// every match (`search(query, corpus_len)`) spent most of its time sorting.
|
|
pub fn scores(&self, query: &str) -> Vec<(usize, f32)> {
|
|
if self.num_docs == 0 {
|
|
return Vec::new();
|
|
}
|
|
// Term-at-a-time accumulation into a dense array: a common term has a
|
|
// posting per document, and hashing each one dominated query time.
|
|
// IDF is computed here rather than cached at build time: it depends on
|
|
// the live document count, which changes with every incremental
|
|
// add/remove, and costs one `ln` per query term.
|
|
let mut acc = vec![0.0f32; self.doc_lengths.len()];
|
|
let mut matched = false;
|
|
for token in tokenize(query) {
|
|
let Some(postings) = self.inverted.get(token.as_str()) else {
|
|
continue;
|
|
};
|
|
matched = true;
|
|
let df = postings.len() as f32;
|
|
let idf = ((self.num_docs as f32 - df + 0.5) / (df + 0.5) + 1.0).ln();
|
|
for &(doc_id, freq) in postings {
|
|
let dl = self.doc_lengths[doc_id] as f32;
|
|
let freq_f = freq as f32;
|
|
let tf = (freq_f * (self.k1 + 1.0))
|
|
/ (freq_f + self.k1 * (1.0 - self.b + self.b * dl / self.avg_dl));
|
|
acc[doc_id] += idf * tf;
|
|
}
|
|
}
|
|
if !matched {
|
|
return Vec::new();
|
|
}
|
|
// Every contribution is strictly positive (idf = ln(1 + x), x > 0), so
|
|
// a zero entry is a document no query term touched.
|
|
acc.into_iter()
|
|
.enumerate()
|
|
.filter(|&(_, score)| score > 0.0)
|
|
.collect()
|
|
}
|
|
|
|
/// Number of document slots (live or not) the index covers. Ids are
|
|
/// positions in the document list it mirrors.
|
|
pub fn len(&self) -> usize {
|
|
self.doc_lengths.len()
|
|
}
|
|
|
|
/// `true` when the index covers no document slots.
|
|
pub fn is_empty(&self) -> bool {
|
|
self.doc_lengths.is_empty()
|
|
}
|
|
|
|
/// Index `text` as document `doc_id`, which must be the next free id
|
|
/// (`self.len()`) or an existing slot that is currently empty (removed or
|
|
/// tombstoned). After any sequence of `add_document` / `remove_document`
|
|
/// calls the index scores exactly as one freshly built from the same live
|
|
/// documents.
|
|
pub fn add_document(&mut self, doc_id: usize, text: &str) {
|
|
if doc_id >= self.doc_lengths.len() {
|
|
self.doc_lengths.resize(doc_id + 1, 0);
|
|
}
|
|
debug_assert_eq!(self.doc_lengths[doc_id], 0, "slot {doc_id} is occupied");
|
|
|
|
let tokens = tokenize(text);
|
|
let mut term_freqs: HashMap<&str, u32> = HashMap::new();
|
|
for token in &tokens {
|
|
*term_freqs.entry(token).or_insert(0) += 1;
|
|
}
|
|
for (token, freq) in term_freqs {
|
|
let postings = self.inverted.entry(token.to_string()).or_default();
|
|
// Posting lists stay sorted by doc id; appends are the common case.
|
|
match postings.last() {
|
|
Some(&(last, _)) if last >= doc_id => {
|
|
let at = postings.partition_point(|&(id, _)| id < doc_id);
|
|
postings.insert(at, (doc_id, freq));
|
|
}
|
|
_ => postings.push((doc_id, freq)),
|
|
}
|
|
}
|
|
self.doc_lengths[doc_id] = tokens.len() as u32;
|
|
self.total_length += tokens.len() as u64;
|
|
self.num_docs += 1;
|
|
self.refresh_avg_dl();
|
|
}
|
|
|
|
/// Extend the index to cover `len` document slots, leaving new ones empty.
|
|
/// Used for slots that hold no live document (tombstoned records).
|
|
pub fn pad_to(&mut self, len: usize) {
|
|
if len > self.doc_lengths.len() {
|
|
self.doc_lengths.resize(len, 0);
|
|
}
|
|
}
|
|
|
|
/// Remove document `doc_id`, whose indexed text was `text`. The text is
|
|
/// needed to find its postings; pass exactly what was added.
|
|
pub fn remove_document(&mut self, doc_id: usize, text: &str) {
|
|
let tokens = tokenize(text);
|
|
let mut seen: std::collections::HashSet<&str> = std::collections::HashSet::new();
|
|
for token in &tokens {
|
|
if !seen.insert(token) {
|
|
continue;
|
|
}
|
|
if let Some(postings) = self.inverted.get_mut(token.as_str()) {
|
|
if let Ok(at) = postings.binary_search_by_key(&doc_id, |&(id, _)| id) {
|
|
postings.remove(at);
|
|
}
|
|
if postings.is_empty() {
|
|
self.inverted.remove(token.as_str());
|
|
}
|
|
}
|
|
}
|
|
if let Some(len) = self.doc_lengths.get_mut(doc_id) {
|
|
self.total_length = self.total_length.saturating_sub(u64::from(*len));
|
|
*len = 0;
|
|
}
|
|
self.num_docs = self.num_docs.saturating_sub(1);
|
|
self.refresh_avg_dl();
|
|
}
|
|
|
|
fn refresh_avg_dl(&mut self) {
|
|
self.avg_dl = if self.num_docs > 0 {
|
|
self.total_length as f32 / self.num_docs as f32
|
|
} else {
|
|
0.0
|
|
};
|
|
}
|
|
|
|
/// Rebuild the index from scratch (e.g., after compaction).
|
|
pub fn rebuild(&mut self, documents: &[String], tombstones: &[u8]) {
|
|
self.inverted.clear();
|
|
self.doc_lengths = vec![0; documents.len()];
|
|
self.total_length = 0;
|
|
self.avg_dl = 0.0;
|
|
self.num_docs = 0;
|
|
self.index_documents(documents, tombstones);
|
|
}
|
|
|
|
/// Internal: populate the inverted index from documents.
|
|
fn index_documents(&mut self, documents: &[String], tombstones: &[u8]) {
|
|
let mut total_length: u64 = 0;
|
|
let mut count: usize = 0;
|
|
|
|
for (i, doc) in documents.iter().enumerate() {
|
|
if i < tombstones.len() && tombstones[i] != 0 {
|
|
continue;
|
|
}
|
|
|
|
let tokens = tokenize(doc);
|
|
let doc_len = tokens.len() as u32;
|
|
self.doc_lengths[i] = doc_len;
|
|
total_length += doc_len as u64;
|
|
count += 1;
|
|
|
|
// Count term frequencies for this document.
|
|
let mut term_freqs: HashMap<&str, u32> = HashMap::new();
|
|
for token in &tokens {
|
|
*term_freqs.entry(token).or_insert(0) += 1;
|
|
}
|
|
|
|
for (token, freq) in term_freqs {
|
|
self.inverted
|
|
.entry(token.to_string())
|
|
.or_default()
|
|
.push((i, freq));
|
|
}
|
|
}
|
|
|
|
self.num_docs = count;
|
|
self.total_length = total_length;
|
|
self.refresh_avg_dl();
|
|
|
|
// Sort posting lists by doc_id for cache-friendly access
|
|
for postings in self.inverted.values_mut() {
|
|
postings.sort_by_key(|&(doc_id, _)| doc_id);
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Tokenize a string: lowercase, split on non-alphanumeric characters,
|
|
/// filter empty tokens.
|
|
fn tokenize(text: &str) -> Vec<String> {
|
|
text.to_lowercase()
|
|
.split(|c: char| !c.is_alphanumeric())
|
|
.filter(|s| !s.is_empty())
|
|
.map(|s| s.to_string())
|
|
.collect()
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn single_document_match() {
|
|
let docs = vec!["the quick brown fox jumps over the lazy dog".to_string()];
|
|
let tombstones = vec![0u8];
|
|
let index = BM25Index::build(&docs, &tombstones);
|
|
|
|
let results = index.search("fox", 10);
|
|
assert_eq!(results.len(), 1);
|
|
assert_eq!(results[0].0, 0);
|
|
assert!(results[0].1 > 0.0);
|
|
}
|
|
|
|
#[test]
|
|
fn multi_document_ranking() {
|
|
let docs = vec![
|
|
"rust programming language systems".to_string(),
|
|
"rust rust rust is great for systems programming".to_string(),
|
|
"python is a scripting language".to_string(),
|
|
];
|
|
let tombstones = vec![0, 0, 0];
|
|
let index = BM25Index::build(&docs, &tombstones);
|
|
|
|
let results = index.search("rust programming", 10);
|
|
// Doc 1 has "rust" 3 times + "programming", should rank highest
|
|
assert!(results.len() >= 2);
|
|
assert_eq!(
|
|
results[0].0, 1,
|
|
"doc with most 'rust' mentions should rank first"
|
|
);
|
|
assert_eq!(results[1].0, 0);
|
|
}
|
|
|
|
#[test]
|
|
fn no_matches_returns_empty() {
|
|
let docs = vec!["hello world".to_string()];
|
|
let tombstones = vec![0u8];
|
|
let index = BM25Index::build(&docs, &tombstones);
|
|
|
|
let results = index.search("nonexistent", 10);
|
|
assert!(results.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn tombstoned_documents_excluded() {
|
|
let docs = vec![
|
|
"rust programming".to_string(),
|
|
"rust systems language".to_string(),
|
|
];
|
|
let tombstones = vec![0, 1]; // doc 1 tombstoned
|
|
|
|
let index = BM25Index::build(&docs, &tombstones);
|
|
let results = index.search("rust", 10);
|
|
|
|
assert_eq!(results.len(), 1);
|
|
assert_eq!(results[0].0, 0);
|
|
}
|
|
|
|
#[test]
|
|
fn rebuild_after_changes() {
|
|
let docs = vec!["hello world".to_string(), "goodbye world".to_string()];
|
|
let tombstones = vec![0, 0];
|
|
let mut index = BM25Index::build(&docs, &tombstones);
|
|
|
|
// Initially both docs match "world"
|
|
let results = index.search("world", 10);
|
|
assert_eq!(results.len(), 2);
|
|
|
|
// Tombstone doc 0 and rebuild
|
|
let new_tombstones = vec![1, 0];
|
|
index.rebuild(&docs, &new_tombstones);
|
|
|
|
let results = index.search("world", 10);
|
|
assert_eq!(results.len(), 1);
|
|
assert_eq!(results[0].0, 1);
|
|
}
|
|
|
|
#[test]
|
|
fn empty_query_returns_empty() {
|
|
let docs = vec!["hello world".to_string()];
|
|
let tombstones = vec![0u8];
|
|
let index = BM25Index::build(&docs, &tombstones);
|
|
|
|
let results = index.search("", 10);
|
|
assert!(results.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn empty_documents_returns_empty() {
|
|
let docs: Vec<String> = Vec::new();
|
|
let tombstones: Vec<u8> = Vec::new();
|
|
let index = BM25Index::build(&docs, &tombstones);
|
|
|
|
let results = index.search("anything", 10);
|
|
assert!(results.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn tokenizer_handles_punctuation() {
|
|
let tokens = tokenize("Hello, World! This is a test.");
|
|
assert_eq!(tokens, vec!["hello", "world", "this", "is", "a", "test"]);
|
|
}
|
|
|
|
#[test]
|
|
fn tokenizer_handles_mixed_case_and_numbers() {
|
|
let tokens = tokenize("HTTP 200 OK");
|
|
assert_eq!(tokens, vec!["http", "200", "ok"]);
|
|
}
|
|
|
|
#[test]
|
|
fn top_k_limits_results() {
|
|
let docs: Vec<String> = (0..20)
|
|
.map(|i| format!("document number {i} about rust"))
|
|
.collect();
|
|
let tombstones = vec![0u8; 20];
|
|
let index = BM25Index::build(&docs, &tombstones);
|
|
|
|
let results = index.search("rust", 5);
|
|
assert_eq!(results.len(), 5);
|
|
}
|
|
|
|
#[test]
|
|
fn idf_weights_rare_terms_higher() {
|
|
let docs = vec![
|
|
"common common common rare".to_string(),
|
|
"common common common".to_string(),
|
|
"common common".to_string(),
|
|
];
|
|
let tombstones = vec![0, 0, 0];
|
|
let index = BM25Index::build(&docs, &tombstones);
|
|
|
|
// "rare" only appears in doc 0, should get a high score
|
|
let results = index.search("rare", 10);
|
|
assert_eq!(results.len(), 1);
|
|
assert_eq!(results[0].0, 0);
|
|
assert!(results[0].1 > 0.0);
|
|
}
|
|
|
|
#[test]
|
|
fn score_matches_the_bm25_formula() {
|
|
let docs = vec![
|
|
"rust programming".to_string(),
|
|
"rust systems".to_string(),
|
|
"python scripting".to_string(),
|
|
];
|
|
let index = BM25Index::build(&docs, &[0, 0, 0]);
|
|
|
|
// "python": df = 1 of N = 3. Every doc has the average length (2) and
|
|
// tf = 1, so the tf factor is exactly 1 and the score is the IDF.
|
|
let results = index.search("python", 3);
|
|
let expected_idf = ((3.0f32 - 1.0 + 0.5) / (1.0 + 0.5) + 1.0).ln();
|
|
assert_eq!(results.len(), 1);
|
|
assert_eq!(results[0].0, 2);
|
|
assert!((results[0].1 - expected_idf).abs() < 1e-6, "{results:?}");
|
|
}
|
|
|
|
#[test]
|
|
fn postings_sorted_by_doc_id() {
|
|
let docs: Vec<String> = (0..20)
|
|
.map(|i| format!("document {i} about rust"))
|
|
.collect();
|
|
let tombstones = vec![0u8; 20];
|
|
let index = BM25Index::build(&docs, &tombstones);
|
|
|
|
if let Some(postings) = index.inverted.get("rust") {
|
|
for w in postings.windows(2) {
|
|
assert!(
|
|
w[0].0 <= w[1].0,
|
|
"postings not sorted: {} > {}",
|
|
w[0].0,
|
|
w[1].0
|
|
);
|
|
}
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn wand_returns_same_results_as_exhaustive() {
|
|
// WAND-style search should produce same scores as exhaustive
|
|
let docs: Vec<String> = (0..100)
|
|
.map(|i| {
|
|
if i % 3 == 0 {
|
|
format!("rust programming language {i}")
|
|
} else if i % 3 == 1 {
|
|
format!("python scripting language {i}")
|
|
} else {
|
|
format!("javascript web development {i}")
|
|
}
|
|
})
|
|
.collect();
|
|
let tombstones = vec![0u8; 100];
|
|
let index = BM25Index::build(&docs, &tombstones);
|
|
|
|
let results_10 = index.search("rust programming", 10);
|
|
let results_100 = index.search("rust programming", 100);
|
|
|
|
// Top-10 from k=10 should have same scores as first 10 from k=100
|
|
assert_eq!(results_10.len(), 10);
|
|
let scores_10: Vec<f32> = results_10.iter().map(|r| r.1).collect();
|
|
let scores_100: Vec<f32> = results_100.iter().take(10).map(|r| r.1).collect();
|
|
for (s10, s100) in scores_10.iter().zip(&scores_100) {
|
|
assert!(
|
|
(s10 - s100).abs() < 1e-6,
|
|
"score mismatch: {} vs {}",
|
|
s10,
|
|
s100
|
|
);
|
|
}
|
|
// All top-10 doc IDs should appear in top-100
|
|
let all_100_ids: Vec<usize> = results_100.iter().map(|r| r.0).collect();
|
|
for (idx, _) in &results_10 {
|
|
assert!(
|
|
all_100_ids.contains(idx),
|
|
"doc {idx} missing from k=100 results"
|
|
);
|
|
}
|
|
}
|
|
/// Documents drawn from a small vocabulary so terms collide heavily.
|
|
fn random_doc(state: &mut u64) -> String {
|
|
const VOCAB: &[&str] = &[
|
|
"alpha", "beta", "gamma", "delta", "eps", "zeta", "eta", "x1",
|
|
];
|
|
let mut next = || {
|
|
*state = state
|
|
.wrapping_mul(6364136223846793005)
|
|
.wrapping_add(1442695040888963407);
|
|
(*state >> 33) as usize
|
|
};
|
|
let len = 1 + next() % 9;
|
|
(0..len)
|
|
.map(|_| VOCAB[next() % VOCAB.len()])
|
|
.collect::<Vec<_>>()
|
|
.join(" ")
|
|
}
|
|
|
|
#[test]
|
|
fn incremental_updates_match_a_fresh_build_exactly() {
|
|
for seed in 0..60u64 {
|
|
let mut state = seed.wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1;
|
|
let mut docs: Vec<String> = Vec::new();
|
|
let mut tombstones: Vec<u8> = Vec::new();
|
|
let mut index = BM25Index::build(&docs, &tombstones);
|
|
|
|
for step in 0..80 {
|
|
state = state.wrapping_mul(6364136223846793005).wrapping_add(1);
|
|
let live: Vec<usize> = (0..docs.len()).filter(|&i| tombstones[i] == 0).collect();
|
|
match (state >> 40) % 4 {
|
|
0 if !live.is_empty() => {
|
|
// delete
|
|
let id = live[(state >> 20) as usize % live.len()];
|
|
index.remove_document(id, &docs[id]);
|
|
tombstones[id] = 1;
|
|
}
|
|
1 if !live.is_empty() => {
|
|
// update in place
|
|
let id = live[(state >> 20) as usize % live.len()];
|
|
let new_text = random_doc(&mut state);
|
|
index.remove_document(id, &docs[id]);
|
|
index.add_document(id, &new_text);
|
|
docs[id] = new_text;
|
|
}
|
|
_ => {
|
|
let text = random_doc(&mut state);
|
|
index.add_document(docs.len(), &text);
|
|
docs.push(text);
|
|
tombstones.push(0);
|
|
}
|
|
}
|
|
|
|
let fresh = BM25Index::build(&docs, &tombstones);
|
|
for query in ["alpha", "beta gamma", "x1 zeta alpha delta", "missing"] {
|
|
let got = index.search(query, 5);
|
|
let want = fresh.search(query, 5);
|
|
assert_eq!(got.len(), want.len(), "seed {seed} step {step} {query:?}");
|
|
for (g, w) in got.iter().zip(&want) {
|
|
assert_eq!(
|
|
g.0, w.0,
|
|
"seed {seed} step {step} {query:?}: {got:?} vs {want:?}"
|
|
);
|
|
assert!(
|
|
(g.1 - w.1).abs() < 1e-5,
|
|
"seed {seed} step {step} {query:?}"
|
|
);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn scores_is_the_unranked_form_of_a_full_search() {
|
|
let mut state = 99u64;
|
|
let docs: Vec<String> = (0..200).map(|_| random_doc(&mut state)).collect();
|
|
let tombstones: Vec<u8> = (0..200).map(|i| u8::from(i % 7 == 0)).collect();
|
|
let index = BM25Index::build(&docs, &tombstones);
|
|
for query in ["alpha", "beta gamma x1", "missing", ""] {
|
|
let mut all = index.scores(query);
|
|
all.sort_by(|a, b| b.1.total_cmp(&a.1).then(a.0.cmp(&b.0)));
|
|
assert_eq!(all, index.search(query, docs.len()), "{query:?}");
|
|
assert!(all.iter().all(|(id, _)| tombstones[*id] == 0));
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn ties_break_towards_the_lower_doc_id() {
|
|
let docs: Vec<String> = (0..6).map(|_| "same text".to_string()).collect();
|
|
let index = BM25Index::build(&docs, &[0; 6]);
|
|
let ids: Vec<usize> = index.search("same", 3).into_iter().map(|r| r.0).collect();
|
|
assert_eq!(ids, [0, 1, 2]);
|
|
}
|
|
}
|