use crate::error; use crate::error::SearchError; 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, Value}; use tantivy::tokenizer::{NgramTokenizer, TokenizerManager}; use tantivy::{Index, IndexReader, IndexWriter, ReloadPolicy, TantivyDocument, Term}; #[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 { if let Some(path) = self.path { let ngram_32 = NgramTokenizer::new(1, 32, false)?; let tokenizer_manager = TokenizerManager::default(); tokenizer_manager.register("ngram_32", ngram_32); let mmap_directory = MmapDirectory::open(path)?; let managed_directory = ManagedDirectory::wrap(Box::new(mmap_directory))?; let index = Index::builder() .schema(Schema::schema().clone()) .tokenizers(tokenizer_manager) .open_or_create(managed_directory)?; let reader = index .reader_builder() .reload_policy(ReloadPolicy::OnCommitWithDelay) .try_into()?; let writer = index.writer(50_000_000)?; Ok(SearchIndex { index, reader, writer, }) } else { Err(SearchError::IndexPathNotSpecified) } } } pub struct SearchIndex { index: Index, reader: IndexReader, writer: IndexWriter, } impl SearchIndex { pub fn builder() -> SearchIndexBuilder { SearchIndexBuilder::default() } pub fn add<'a>(&self, document: TantivyDocument) -> crate::Result<()> { self.writer.add_document(document)?; Ok(()) } pub fn remove_all_of_type(&mut self, type_: u64) { self.writer .delete_term(Term::from_field_u64(Schema::type_field(), type_)); } pub fn commit(&mut self) -> crate::Result<()> { self.writer.commit()?; Ok(()) } pub fn query( &self, type_: Option, user_query: &str, default_fields: Vec, ) -> error::Result> { 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 = searcher.search(&query, &TopDocs::with_limit(10000).order_by_score())?; let mut iris = vec![]; for (_score, address) in results.iter() { let doc: TantivyDocument = searcher.doc(*address)?; if let Some(doc_iri) = doc.get_first(Schema::iri_field()) { let doc_iri_string = doc_iri.as_str().unwrap_or("???").to_string(); iris.push(doc_iri_string); } } Ok(iris) } }