tools/search/src/index.rs

use std::cmp::Ordering;
use std::collections::{BTreeMap, BTreeSet};
use crate::{error, update};
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, ReloadPolicy, TantivyDocument, Term};
use tracing::{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<update::IndexWriter> {
        let inner = self.index.writer(128 * 1024usize.pow(2))?;
        Ok(crate::IndexWriter { inner })
    }

    pub fn query(
        &self,
        type_: Option<u64>,
        user_query: &str,
        default_fields: Vec<Field>,
        limit: usize,
    ) -> error::Result<Vec<TantivyDocument>> {
        let span = span!(Level::DEBUG, "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: 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)
    }
}