tools/search/src/index.rs

use crate::{error, SearchDocument};
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, NamedFieldDocument, Value};
use tantivy::tokenizer::{LowerCaser, NgramTokenizer, TextAnalyzer, TokenizerManager};
use tantivy::{Document, Index, IndexReader, IndexWriter, ReloadPolicy, Score, TantivyDocument, Term};

#[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> {
        if let Some(path) = self.path {
            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 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<u64>,
        user_query: &str,
        default_fields: Vec<Field>,
    ) -> error::Result<Vec<NamedFieldDocument>> {
        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())?
            .iter()
            .map(|(_, address)| searcher.doc(*address))
            .filter_map(Result::ok)
            .map(|doc: SearchDocument| doc.to_named_doc(Schema::schema()))
            .collect();
        Ok(results)
    }
}