Raise matching floor: boundary tokenizer, stemming/folding/stopwords, boosts, snippets, engine seam
This commit is contained in:
@@ -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
@@ -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);
|
||||
|
||||
@@ -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
@@ -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
@@ -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
|
||||
}
|
||||
|
||||
@@ -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:?}");
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user