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, } impl SearchIndexBuilder { pub fn with_path(mut self, path: impl Into) -> Self { self.path = Some(path.into()); self } pub fn build(self) -> error::Result { 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 { Ok(self.index.writer(128 * 1024 * 1024)?) } pub fn query( &self, type_: Option, user_query: &str, default_fields: Vec, limit: usize, ) -> error::Result> { 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 = searcher.search(&query, &TopDocs::with_limit(limit).order_by_score())? .iter() .map(|(_, address)| searcher.doc(*address)) .filter_map(Result::ok) .collect(); Ok(results) } }