103 lines
3.2 KiB
Rust
103 lines
3.2 KiB
Rust
use std::fs;
|
|
use crate::error;
|
|
use crate::schema::Schema;
|
|
use std::path::PathBuf;
|
|
use tantivy::collector::TopDocs;
|
|
use tantivy::directory::{ManagedDirectory, MmapDirectory};
|
|
use tantivy::query::{BooleanQuery, Occur, QueryParser, TermQuery};
|
|
use tantivy::schema::{Field, IndexRecordOption};
|
|
use tantivy::tokenizer::{LowerCaser, NgramTokenizer, TextAnalyzer, TokenizerManager};
|
|
use tantivy::{Index, IndexReader, IndexWriter, ReloadPolicy, TantivyDocument, Term};
|
|
use tracing::debug_span;
|
|
|
|
#[derive(Default)]
|
|
pub struct SearchIndexBuilder {
|
|
path: Option<PathBuf>,
|
|
}
|
|
|
|
impl SearchIndexBuilder {
|
|
pub fn with_path(mut self, path: impl Into<PathBuf>) -> Self {
|
|
self.path = Some(path.into());
|
|
self
|
|
}
|
|
|
|
pub fn build(self) -> error::Result<SearchIndex> {
|
|
let ngram_32 = NgramTokenizer::new(1, 32, false)?;
|
|
let ngram_32_lowercase = TextAnalyzer::builder(ngram_32)
|
|
.filter(LowerCaser)
|
|
.build();
|
|
|
|
let tokenizer_manager = TokenizerManager::default();
|
|
tokenizer_manager.register("ngram_32", ngram_32_lowercase);
|
|
|
|
let builder = Index::builder()
|
|
.schema(Schema::schema().clone())
|
|
.tokenizers(tokenizer_manager);
|
|
|
|
let index = if let Some(path) = self.path {
|
|
fs::create_dir_all(&path)?;
|
|
let mmap_directory = MmapDirectory::open(path)?;
|
|
let managed_directory = ManagedDirectory::wrap(Box::new(mmap_directory))?;
|
|
builder.open_or_create(managed_directory)?
|
|
} else {
|
|
builder.create_in_ram()?
|
|
};
|
|
|
|
let reader = index
|
|
.reader_builder()
|
|
.reload_policy(ReloadPolicy::OnCommitWithDelay)
|
|
.try_into()?;
|
|
|
|
Ok(SearchIndex {
|
|
index,
|
|
reader,
|
|
})
|
|
}
|
|
}
|
|
|
|
pub struct SearchIndex {
|
|
index: Index,
|
|
reader: IndexReader,
|
|
}
|
|
|
|
impl SearchIndex {
|
|
pub fn builder() -> SearchIndexBuilder {
|
|
SearchIndexBuilder::default()
|
|
}
|
|
|
|
pub fn writer(&mut self) -> crate::Result<IndexWriter> {
|
|
Ok(self.index.writer(128 * 1024 * 1024)?)
|
|
}
|
|
|
|
pub fn query(
|
|
&self,
|
|
type_: Option<u64>,
|
|
user_query: &str,
|
|
default_fields: Vec<Field>,
|
|
limit: usize,
|
|
) -> error::Result<Vec<TantivyDocument>> {
|
|
let _enter = debug_span!("Search Query").entered();
|
|
|
|
let parser = QueryParser::for_index(&self.index, default_fields);
|
|
let (user_query, _) = parser.parse_query_lenient(user_query);
|
|
|
|
let mut subqueries = vec![
|
|
(Occur::Must, user_query)
|
|
];
|
|
|
|
if let Some(type_) = type_ {
|
|
let doc_type_term = Term::from_field_u64(Schema::type_field(), type_);
|
|
let doc_type_query = Box::new(TermQuery::new(doc_type_term, IndexRecordOption::Basic));
|
|
subqueries.push((Occur::Must, doc_type_query));
|
|
}
|
|
|
|
let query = BooleanQuery::new(subqueries);
|
|
let searcher = self.reader.searcher();
|
|
let results: Vec<TantivyDocument> = searcher.search(&query, &TopDocs::with_limit(limit).order_by_score())?
|
|
.iter()
|
|
.map(|(_, address)| searcher.doc(*address))
|
|
.filter_map(Result::ok)
|
|
.collect();
|
|
Ok(results)
|
|
}
|
|
} |