tools/search/src/index.rs

use crate::{error, SearchDocument};
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};
use tantivy::indexer::UserOperation;
use tracing::{info, instrument, span, Level};

#[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 {
            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,
        })
    }
}

#[derive(Clone)]
pub struct SearchIndex {
    index: Index,
    reader: IndexReader,
}

impl SearchIndex {
    pub fn builder() -> SearchIndexBuilder {
        SearchIndexBuilder::default()
    }

    pub fn writer(&self) -> crate::Result<IndexWriter> {
        Ok(self.index.writer(128_000_000)?)
    }

    pub fn remove_all_of_type(&mut self, type_: u64) -> crate::Result<()> {
        let mut writer: IndexWriter<TantivyDocument> = self.index.writer(64_000_000)?;
        writer.delete_term(Term::from_field_u64(Schema::type_field(), type_));
        writer.commit()?;
        Ok(())
    }

    pub fn query(
        &self,
        type_: Option<u64>,
        user_query: &str,
        default_fields: Vec<Field>,
    ) -> error::Result<Vec<NamedFieldDocument>> {
        let span = span!(Level::INFO, "Search Query");
        let _enter = span.enter();

        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)
    }
}