osobhandClaude Opus 5
e9aeb110b7
Traversal recursed one frame per level with the depth taken from the file (a u16), and followed child addresses without asking whether they were shared. Two crafted inputs, both reproduced before fixing: - A node listing itself as its own child, under a header claiming 65 535 levels, overflowed the stack and aborted the process — SIGABRT, not an error a caller can handle — from under 100 bytes. - Levels whose children all point at one shared node below reached it fan-out^depth times: 29.5 million records in 8 s from ~5 KB, and one more level would exhaust memory. Depth is now capped at 64, as the fractal heap already was; no real tree approaches it, since even at the minimum fan-out of two that is over 2^64 records. And traversal stops once it has produced more records than the file has bytes to hold them — a valid tree stores each record once in its own bytes, so this bounds shared subtrees without trusting the header's own `total_records`. Both inputs now fail in under a millisecond. Every B-tree v2 user goes through this collector: dense attributes, v2 groups, shared messages and chunk indexes. To show the budget never refuses a real file, a new interop test has HDF5 2.0 write a depth-2 chunk index with 40 000 records and reads back all 160 000 values; it fails when the budget is deliberately made too tight. Also corrects `BM25Index::search`, which claimed to use Block-Max WAND. It scores exhaustively, and pruning would not help the store: `hybrid_search` needs every score because fusion normalises over them. Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
749 lines
28 KiB
Rust
749 lines
28 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,
|
|
/// Applied to every document and query token, so the two always agree.
|
|
filter: TokenFilter,
|
|
}
|
|
|
|
impl BM25Index {
|
|
/// Build a BM25 index from a set of documents, excluding tombstoned entries.
|
|
pub fn build(documents: &[String], tombstones: &[u8]) -> Self {
|
|
Self::build_with(documents, tombstones, TokenFilter::default())
|
|
}
|
|
|
|
/// [`BM25Index::build`] with the token filter chosen explicitly.
|
|
pub fn build_with(documents: &[String], tombstones: &[u8], filter: TokenFilter) -> 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,
|
|
filter,
|
|
};
|
|
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.
|
|
///
|
|
/// Scores every matching document exhaustively, then keeps the top `k`.
|
|
/// There is no early termination (WAND, MaxScore): the store's hot path
|
|
/// is [`scores`](Self::scores), because score fusion normalises over the
|
|
/// whole matching set and so needs every score, which no pruning scheme
|
|
/// can skip. This method is for BM25-only callers.
|
|
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_with(query, self.filter) {
|
|
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()
|
|
}
|
|
|
|
/// The token filter this index was built with.
|
|
pub fn token_filter(&self) -> TokenFilter {
|
|
self.filter
|
|
}
|
|
|
|
/// 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_with(text, self.filter);
|
|
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_with(text, self.filter);
|
|
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_with(doc, self.filter);
|
|
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.
|
|
/// What [`tokenize_with`] does to each token after splitting.
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
|
|
pub enum TokenFilter {
|
|
/// Lowercase and split only — the original behaviour.
|
|
#[default]
|
|
Plain,
|
|
/// Also strip common English inflections, so "running" and "runs" match
|
|
/// "run". Conservative on purpose: only plural and past/continuous verb
|
|
/// endings, and only on tokens long enough that stripping leaves a real
|
|
/// stem. A stemmer earns its keep by conflating *related* words; an
|
|
/// aggressive one also conflates unrelated ones ("universe"/"university"),
|
|
/// which costs precision.
|
|
Stemmed,
|
|
}
|
|
|
|
/// Strip common English inflections from an already-lowercased token.
|
|
///
|
|
/// Applied identically to documents and queries, so the pair only has to agree
|
|
/// with itself — the stem need not be a real word.
|
|
fn stem(token: &str) -> &str {
|
|
// Below this, stripping does more harm than good ("bed" -> "b").
|
|
const MIN_STEM: usize = 4;
|
|
let strip = |suffix: &str, min_len: usize| -> Option<&str> {
|
|
let stem = token.strip_suffix(suffix)?;
|
|
(stem.len() >= min_len).then_some(stem)
|
|
};
|
|
|
|
// Plurals first: "studies" -> "studi", "classes" -> "class", "cats" -> "cat".
|
|
// "ies" keeps its "i" so the result meets "-ied" ("studied" -> "studi").
|
|
if let Some(stem) = strip("ies", 2) {
|
|
return &token[..stem.len() + 1];
|
|
}
|
|
for suffix in ["sses", "shes", "ches", "xes", "zes"] {
|
|
if let Some(stem) = strip(suffix, MIN_STEM - 1) {
|
|
// Keep the sibilant: "classes" -> "class", not "clas".
|
|
return &token[..stem.len() + 2];
|
|
}
|
|
}
|
|
// Verb endings before the bare plural, so "raced" doesn't become "raced".
|
|
if let Some(stem) = strip("ing", MIN_STEM - 1).or_else(|| strip("ed", MIN_STEM - 1)) {
|
|
return undouble(stem);
|
|
}
|
|
if !token.ends_with("ss")
|
|
&& !token.ends_with("us")
|
|
&& !token.ends_with("is")
|
|
&& let Some(stem) = strip("s", MIN_STEM - 1)
|
|
{
|
|
return stem;
|
|
}
|
|
token
|
|
}
|
|
|
|
/// "runn" -> "run": undo the consonant doubling that "-ing"/"-ed" introduce.
|
|
fn undouble(stem: &str) -> &str {
|
|
let mut chars = stem.chars().rev();
|
|
let (Some(last), Some(prev)) = (chars.next(), chars.next()) else {
|
|
return stem;
|
|
};
|
|
let doubled = last == prev && !"aeiou".contains(last) && last.is_ascii_alphabetic();
|
|
if doubled && stem.len() > 3 {
|
|
&stem[..stem.len() - 1]
|
|
} else {
|
|
stem
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
fn tokenize(text: &str) -> Vec<String> {
|
|
tokenize_with(text, TokenFilter::Plain)
|
|
}
|
|
|
|
/// Split `text` into scoring tokens under `filter`.
|
|
pub fn tokenize_with(text: &str, filter: TokenFilter) -> Vec<String> {
|
|
text.to_lowercase()
|
|
.split(|c: char| !c.is_alphanumeric())
|
|
.filter(|s| !s.is_empty())
|
|
.map(|token| match filter {
|
|
TokenFilter::Plain => token.to_string(),
|
|
TokenFilter::Stemmed => stem(token).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 top_k_search_matches_ranking_every_score() {
|
|
// `search` must agree with ranking the full `scores` set — the
|
|
// bounded heap is an optimisation over sorting, not an approximation.
|
|
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 stemming_conflates_inflections_of_the_same_word() {
|
|
let stem_of = |w: &str| tokenize_with(w, TokenFilter::Stemmed).pop().unwrap();
|
|
// Pairs that should meet.
|
|
for (a, b) in [
|
|
("running", "runs"),
|
|
("trained", "training"),
|
|
("miles", "mile"),
|
|
("studies", "studied"),
|
|
("mentioned", "mentioning"),
|
|
("classes", "class"),
|
|
("planned", "planning"),
|
|
] {
|
|
assert_eq!(stem_of(a), stem_of(b), "{a} / {b} should share a stem");
|
|
}
|
|
// Pairs that must stay apart. Note which pairs are deliberately absent:
|
|
// "bed"/"bedding" and "gas"/"gassed" both collapse to one stem, which
|
|
// is what Porter does too and is right — they are related words.
|
|
for (a, b) in [
|
|
("universe", "university"),
|
|
("business", "busy"),
|
|
("this", "thing"),
|
|
] {
|
|
assert_ne!(stem_of(a), stem_of(b), "{a} / {b} must not be conflated");
|
|
}
|
|
// Short words and non-inflections are left alone.
|
|
for word in ["run", "bus", "is", "his", "data", "gas"] {
|
|
assert_eq!(stem_of(word), word, "{word} should be untouched");
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn stemming_is_off_by_default_and_applied_consistently() {
|
|
assert_eq!(tokenize("Running miles"), ["running", "miles"]);
|
|
assert_eq!(
|
|
tokenize_with("Running miles", TokenFilter::Stemmed),
|
|
["run", "mile"]
|
|
);
|
|
|
|
// A query inflected differently from the document still matches.
|
|
let docs = vec!["I ran while training for the marathon".to_string()];
|
|
let plain = BM25Index::build_with(&docs, &[0], TokenFilter::Plain);
|
|
let stemmed = BM25Index::build_with(&docs, &[0], TokenFilter::Stemmed);
|
|
assert!(plain.search("trains", 1).is_empty());
|
|
assert_eq!(stemmed.search("trains", 1).len(), 1);
|
|
}
|
|
|
|
#[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]);
|
|
}
|
|
}
|