tools/search/src/index.rs

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, doc};

#[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 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 remove_all_annotated_iris(&mut self) {
        self.writer
            .delete_term(Term::from_field_u64(Schema::type_field(), 0u64));
        self.writer.commit().unwrap();
    }

    pub fn add_annotated_iri(
        &self,
        iri: &str,
        label: Option<&str>,
        comment: Option<&str>,
    ) -> crate::Result<()> {
        let mut document = doc!(
            Schema::type_field() => 0u64,
            Schema::iri_field() => iri,
        );

        if let Some(label) = label {
            document.add_text(Schema::label_field(), label.to_lowercase());
        }

        if let Some(comment) = comment {
            document.add_text(Schema::comment_field(), comment);
        }

        self.writer.add_document(document)?;
        Ok(())
    }

    pub fn add<'a>(&self, document: TantivyDocument) -> crate::Result<()> {
        self.writer.add_document(document)?;
        Ok(())
    }

    pub fn commit(&mut self) -> crate::Result<()> {
        self.writer.commit()?;
        Ok(())
    }

    pub fn query(
        &self,
        type_: u64,
        user_query: &str,
        default_fields: Vec<Field>,
    ) -> error::Result<Vec<String>> {
        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));

        let parser = QueryParser::for_index(&self.index, default_fields);
        let (user_query, _) = parser.parse_query_lenient(user_query);
        let query = BooleanQuery::new(vec![
            (Occur::Must, doc_type_query),
            (Occur::Must, user_query),
        ]);
        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)
    }
}