Raise matching floor: boundary tokenizer, stemming/folding/stopwords, boosts, snippets, engine seam

This commit is contained in:
George Coles
2026-09-15 08:11:55 -04:00
parent 0d35036168
commit 382acc36a9
10 changed files with 294 additions and 42 deletions
+67
View File
@@ -0,0 +1,67 @@
use anyhow::Result;
use crate::index::{LocalIndex, SearchHit};
pub struct EngineOutput {
pub hits: Vec<SearchHit>,
pub total: Option<u64>,
}
pub trait SearchEngine: Send + Sync {
fn search(&self, text: &str, budget: usize, only_shared: bool) -> Result<EngineOutput>;
fn doc_count(&self) -> u64;
}
pub struct TantivyEngine {
pub index: LocalIndex,
pub min_coverage: f64,
}
impl SearchEngine for TantivyEngine {
fn search(&self, text: &str, budget: usize, only_shared: bool) -> Result<EngineOutput> {
let (hits, total) =
self.index
.search_with_coverage(text, budget, only_shared, self.min_coverage)?;
Ok(EngineOutput {
hits,
total: Some(total),
})
}
fn doc_count(&self) -> u64 {
self.index.doc_count()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::index::Collection;
use crate::message::EXPOSURE_FULL;
use std::fs;
#[test]
fn tantivy_engine_reports_hits_and_total() {
let temp = tempfile::tempdir().unwrap();
let dir = temp.path().join("corpus");
fs::create_dir_all(&dir).unwrap();
fs::write(dir.join("a.txt"), "rust ownership guide").unwrap();
let engine = TantivyEngine {
index: LocalIndex::open(&temp.path().join("index")).unwrap(),
min_coverage: 0.4,
};
engine
.index
.add_collection(&Collection {
name: "docs".to_string(),
path: dir.display().to_string(),
shared: true,
exposure: EXPOSURE_FULL.to_string(),
})
.unwrap();
let output = engine.search("rust", 5, false).unwrap();
assert_eq!(output.hits.len(), 1);
assert_eq!(output.total, Some(1));
assert_eq!(engine.doc_count(), 1);
}
}
+51 -11
View File
@@ -7,13 +7,19 @@ use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use tantivy::collector::{Count, TopDocs};
use tantivy::directory::MmapDirectory;
use tantivy::query::{AllQuery, BooleanQuery, Occur, Query, TermQuery};
use tantivy::schema::{Field, IndexRecordOption, STORED, STRING, Schema, TEXT, Value};
use tantivy::query::{
AllQuery, BooleanQuery, BoostQuery, EmptyQuery, Occur, PhraseQuery, Query, TermQuery,
};
use tantivy::schema::{
Field, IndexRecordOption, STORED, STRING, Schema, TextFieldIndexing, TextOptions, Value,
};
use tantivy::snippet::SnippetGenerator;
use tantivy::tokenizer::TokenStream;
use tantivy::{Index, IndexReader, IndexWriter, TantivyDocument, Term, doc};
use crate::extract::{extract_file, summary_of, supported};
use crate::message::{EXPOSURE_FULL, ResponseItem};
use crate::tokenizer::{TOKENIZER_NAME, analyzer};
const WRITER_BUDGET: usize = 50_000_000;
const MAX_QUERY_TERMS: usize = 10;
@@ -107,6 +113,7 @@ impl LocalIndex {
let schema = build_schema();
let index = Index::open_or_create(MmapDirectory::open(dir)?, schema)
.context("opening tantivy index")?;
index.tokenizers().register(TOKENIZER_NAME, analyzer());
let schema = index.schema();
let fields = Fields {
url: schema.get_field("url")?,
@@ -198,10 +205,28 @@ impl LocalIndex {
self.reader.reload()?;
let searcher = self.reader.searcher();
let terms = self.query_terms(text);
let user_query: Box<dyn Query> = if terms.is_empty() {
let user_query: Box<dyn Query> = if text.trim().is_empty() {
Box::new(AllQuery)
} else if terms.is_empty() {
Box::new(EmptyQuery)
} else {
self.coverage_query(&terms, min_coverage)
let gate = self.coverage_query(&terms, min_coverage);
let mut parts: Vec<(Occur, Box<dyn Query>)> = vec![(Occur::Must, gate)];
if terms.len() >= 2 {
let phrase_terms: Vec<(usize, Term)> = terms
.iter()
.enumerate()
.map(|(index, term)| (index, Term::from_field_text(self.fields.body, term)))
.collect();
parts.push((
Occur::Should,
Box::new(BoostQuery::new(
Box::new(PhraseQuery::new_with_offset_and_slop(phrase_terms, 1)),
3.0,
)),
));
}
Box::new(BooleanQuery::new(parts))
};
let query: Box<dyn Query> = if only_shared {
let shared_term = TermQuery::new(
@@ -217,13 +242,18 @@ impl LocalIndex {
};
let total = searcher.search(&*query, &Count)? as u64;
let top = searcher.search(&*query, &TopDocs::with_limit(limit).order_by_score())?;
let snippet_generator = SnippetGenerator::create(&searcher, &*query, self.fields.body).ok();
let mut hits = Vec::with_capacity(top.len());
for (_score, address) in top {
let document: TantivyDocument = searcher.doc(address)?;
let snippet = snippet_generator
.as_ref()
.map(|generator| generator.snippet_from_doc(&document).fragment().to_string())
.filter(|fragment| !fragment.trim().is_empty());
hits.push(SearchHit {
url: text_value(&document, self.fields.url),
title: text_value(&document, self.fields.title),
summary: text_value(&document, self.fields.summary),
summary: snippet.unwrap_or_else(|| text_value(&document, self.fields.summary)),
published: text_value(&document, self.fields.published),
exposure: text_value(&document, self.fields.exposure),
collection: text_value(&document, self.fields.collection),
@@ -237,7 +267,7 @@ impl LocalIndex {
fn query_terms(&self, text: &str) -> Vec<String> {
let mut terms = Vec::new();
let mut seen = HashSet::new();
if let Some(mut tokenizer) = self.index.tokenizers().get("default") {
if let Some(mut tokenizer) = self.index.tokenizers().get(TOKENIZER_NAME) {
let mut stream = tokenizer.token_stream(text);
while let Some(token) = stream.next() {
let term = token.text.to_string();
@@ -270,9 +300,12 @@ impl LocalIndex {
let term_query: Box<dyn Query> = Box::new(BooleanQuery::new(vec![
(
Occur::Should,
Box::new(TermQuery::new(
Term::from_field_text(self.fields.title, term),
IndexRecordOption::WithFreqs,
Box::new(BoostQuery::new(
Box::new(TermQuery::new(
Term::from_field_text(self.fields.title, term),
IndexRecordOption::WithFreqs,
)),
2.0,
)),
),
(
@@ -345,9 +378,16 @@ fn text_value(document: &TantivyDocument, field: Field) -> String {
fn build_schema() -> Schema {
let mut builder = Schema::builder();
let frx_options = TextOptions::default()
.set_indexing_options(
TextFieldIndexing::default()
.set_tokenizer(TOKENIZER_NAME)
.set_index_option(IndexRecordOption::WithFreqsAndPositions),
)
.set_stored();
builder.add_text_field("url", STRING | STORED);
builder.add_text_field("title", TEXT | STORED);
builder.add_text_field("body", TEXT | STORED);
builder.add_text_field("title", frx_options.clone());
builder.add_text_field("body", frx_options);
builder.add_text_field("summary", STORED);
builder.add_text_field("published", STORED);
builder.add_text_field("exposure", STRING | STORED);
+2
View File
@@ -1,6 +1,7 @@
pub mod commands;
pub mod config;
pub mod crypto;
pub mod engine;
pub mod extract;
pub mod index;
pub mod message;
@@ -9,6 +10,7 @@ pub mod node;
pub mod registry;
pub mod relay;
pub mod render;
pub mod tokenizer;
pub const VERSION: &str = env!("CARGO_PKG_VERSION");
pub const PROTOCOL: &str = "FRX/0.5";
+13 -7
View File
@@ -102,17 +102,23 @@ pub struct ResponseBody {
pub fn build_response(
qid: &str,
hits: Vec<ResponseItem>,
total: u64,
total: Option<u64>,
max_results: usize,
) -> ResponseBody {
let max_results = max_results.max(1);
let mut results = hits;
results.truncate(max_results);
let more_available = total.saturating_sub(results.len() as u64);
let (truncated, more_available) = match total {
Some(total) => {
let more = total.saturating_sub(results.len() as u64);
(more > 0, more)
}
None => (true, 0),
};
ResponseBody {
qid: qid.to_string(),
results,
truncated: more_available > 0,
truncated,
more_available,
cursor: None,
}
@@ -153,24 +159,24 @@ mod tests {
#[test]
fn max_results_is_respected() {
let response = build_response("q1", vec![item(1), item(2), item(3)], 10, 2);
let response = build_response("q1", vec![item(1), item(2), item(3)], Some(10), 2);
assert_eq!(response.results.len(), 2);
}
#[test]
fn truncation_is_honest() {
let response = build_response("q1", vec![item(1), item(2), item(3)], 10, 2);
let response = build_response("q1", vec![item(1), item(2), item(3)], Some(10), 2);
assert!(response.truncated);
assert_eq!(response.more_available, 8);
let response = build_response("q2", vec![item(1), item(2)], 2, 5);
let response = build_response("q2", vec![item(1), item(2)], Some(2), 5);
assert!(!response.truncated);
assert_eq!(response.more_available, 0);
}
#[test]
fn response_carries_no_scores() {
let response = build_response("q1", vec![item(1)], 1, 5);
let response = build_response("q1", vec![item(1)], Some(1), 5);
let json = serde_json::to_string(&response).unwrap();
assert!(!json.contains("score"));
assert!(!json.contains("relevance"));
+16 -16
View File
@@ -18,6 +18,7 @@ use tokio::task::JoinHandle;
use crate::config::{CLASS_ENRICHMENT, Config, Member, load_members};
use crate::crypto::{Keypair, now_ts, poll_signing_bytes};
use crate::engine::{SearchEngine, TantivyEngine};
use crate::index::{LocalIndex, SearchHit, response_items};
use crate::message::{
AggregateBody, Envelope, QueryBody, ResponseBody, ResponseItem, TYPE_AGGREGATE, TYPE_QUERY,
@@ -34,7 +35,7 @@ struct Aggregates {
pub struct Node {
pub config: Config,
pub key: Keypair,
index: LocalIndex,
engine: Arc<dyn SearchEngine>,
members: RwLock<Vec<Member>>,
members_mtime: Mutex<Option<SystemTime>>,
registry: Option<Arc<Watcher>>,
@@ -113,7 +114,10 @@ pub struct LocalPart {
impl Node {
pub fn open(config: Config) -> Result<Arc<Self>> {
let key = config.load_key()?;
let index = LocalIndex::open(&config.index_dir())?;
let engine: Arc<dyn SearchEngine> = Arc::new(TantivyEngine {
index: LocalIndex::open(&config.index_dir())?,
min_coverage: config.matching.min_coverage,
});
let members_path = config.members_path();
let members = load_members(&members_path)?;
let members_mtime = fs::metadata(&members_path)
@@ -155,10 +159,10 @@ impl Node {
Ok(Arc::new(Self {
config,
key,
index,
members: RwLock::new(members),
members_mtime: Mutex::new(members_mtime),
registry,
engine,
enc_secret,
enc_public,
aggregates: Mutex::new(Aggregates::default()),
@@ -229,7 +233,7 @@ impl Node {
}
pub fn doc_count(&self) -> u64 {
self.index.doc_count()
self.engine.doc_count()
}
pub fn identifier(&self) -> String {
@@ -249,7 +253,8 @@ impl Node {
}
pub fn local_search(&self, text: &str, limit: usize) -> Result<(Vec<SearchHit>, u64)> {
self.index.search(text, limit, false)
let output = self.engine.search(text, limit, false)?;
Ok((output.hits, output.total.unwrap_or(0)))
}
pub async fn start(config: Config) -> Result<NodeHandle> {
@@ -310,9 +315,9 @@ impl Node {
let max = max_results.unwrap_or(self.config.query.max_results).max(1);
let query = QueryBody::new(text, max);
let qid = query.qid.clone();
let (local_hits, local_total) =
self.index
.search_with_coverage(text, max, false, self.config.matching.min_coverage)?;
let output = self.engine.search(text, max, false)?;
let local_hits = output.hits;
let local_total = output.total.unwrap_or(0);
let local_items = response_items(&local_hits);
let mut responses = Vec::new();
if network {
@@ -562,16 +567,11 @@ impl Node {
async fn respond(&self, query: &QueryBody, querier: &str, relay: &str) -> Result<()> {
let max = query.budget.max_results.clamp(1, 1000);
let (hits, total) = self.index.search_with_coverage(
&query.text,
max,
true,
self.config.matching.min_coverage,
)?;
if hits.is_empty() {
let output = self.engine.search(&query.text, max, true)?;
if output.hits.is_empty() {
return Ok(());
}
let body = build_response(&query.qid, response_items(&hits), total, max);
let body = build_response(&query.qid, response_items(&output.hits), output.total, max);
self.send_payload(querier, TYPE_RESPONSE, serde_json::to_value(&body)?, relay)
.await
}
+131
View File
@@ -0,0 +1,131 @@
use tantivy::tokenizer::{
AsciiFoldingFilter, Language, LowerCaser, RemoveLongFilter, Stemmer, StopWordFilter,
TextAnalyzer, Token, TokenStream, Tokenizer,
};
#[derive(Clone, Default)]
pub struct BoundaryTokenizer {
token: Token,
}
impl BoundaryTokenizer {
pub fn new() -> Self {
Self {
token: Token::default(),
}
}
}
impl Tokenizer for BoundaryTokenizer {
type TokenStream<'a> = BoundaryTokenStream<'a>;
fn token_stream<'a>(&'a mut self, text: &'a str) -> Self::TokenStream<'a> {
self.token.reset();
BoundaryTokenStream {
text,
chars: text.char_indices().peekable(),
token: &mut self.token,
position: 0,
}
}
}
pub struct BoundaryTokenStream<'a> {
text: &'a str,
chars: std::iter::Peekable<std::str::CharIndices<'a>>,
token: &'a mut Token,
position: usize,
}
impl TokenStream for BoundaryTokenStream<'_> {
fn advance(&mut self) -> bool {
let mut start = None;
let mut alphabetic = None;
while let Some(&(index, ch)) = self.chars.peek() {
if !ch.is_alphanumeric() {
if start.is_none() {
self.chars.next();
continue;
}
break;
}
match alphabetic {
None => {
start = Some(index);
alphabetic = Some(ch.is_alphabetic());
self.chars.next();
}
Some(prev) if prev == ch.is_alphabetic() => {
self.chars.next();
}
Some(_) => break,
}
}
let Some(start) = start else {
return false;
};
let end = self
.chars
.peek()
.map(|&(index, _)| index)
.unwrap_or(self.text.len());
self.token.text = self.text[start..end].to_string();
self.token.offset_from = start;
self.token.offset_to = end;
self.token.position = self.position;
self.position += 1;
true
}
fn token(&self) -> &Token {
self.token
}
fn token_mut(&mut self) -> &mut Token {
self.token
}
}
pub fn analyzer() -> TextAnalyzer {
TextAnalyzer::builder(BoundaryTokenizer::new())
.filter(LowerCaser)
.filter(AsciiFoldingFilter)
.filter(StopWordFilter::new(Language::English).expect("english stopwords"))
.filter(Stemmer::new(Language::English))
.filter(RemoveLongFilter::limit(255))
.build()
}
pub const TOKENIZER_NAME: &str = "frx";
#[cfg(test)]
mod tests {
use super::*;
fn tokens(text: &str) -> Vec<String> {
let mut analyzer = analyzer();
let mut stream = analyzer.token_stream(text);
let mut out = Vec::new();
while let Some(token) = stream.next() {
out.push(token.text.clone());
}
out
}
#[test]
fn splits_model_numbers_and_folds_case() {
let tokens = tokens("DLEX5555 Café");
assert!(tokens.contains(&"dlex".to_string()), "{tokens:?}");
assert!(tokens.contains(&"5555".to_string()), "{tokens:?}");
assert!(tokens.contains(&"cafe".to_string()), "{tokens:?}");
}
#[test]
fn stems_and_drops_stopwords() {
let tokens = tokens("the dryers are searching");
assert!(!tokens.contains(&"the".to_string()), "{tokens:?}");
assert!(!tokens.contains(&"are".to_string()), "{tokens:?}");
assert!(tokens.iter().any(|t| t.starts_with("dryer")), "{tokens:?}");
assert!(tokens.iter().any(|t| t.starts_with("search")), "{tokens:?}");
}
}