MeiliSearch/meilidb-core/src/query_builder.rs

use std::hash::Hash;
use std::ops::Range;
use std::rc::Rc;
use std::time::Instant;
use std::{cmp, mem};

use rayon::slice::ParallelSliceMut;
use slice_group_by::GroupByMut;
use meilidb_tokenizer::{is_cjk, split_query_string};
use hashbrown::{HashMap, HashSet};
use fst::Streamer;
use log::info;

use crate::automaton::{self, DfaExt, AutomatonExt};
use crate::distinct_map::{DistinctMap, BufferedDistinctMap};
use crate::criterion::Criteria;
use crate::raw_documents_from_matches;
use crate::{Match, DocumentId, Store, RawDocument, Document};

fn generate_automatons(query: &str, synonyms: &HashMap<&str, &[&str]>) -> Vec<(usize, DfaExt)> {
    let has_end_whitespace = query.chars().last().map_or(false, char::is_whitespace);
    let mut groups = split_query_string(query).map(str::to_lowercase).peekable();
    let mut automatons = Vec::new();
    let mut index = 0;

    while let Some(word) = groups.next() {
        let word = word.as_str();
        let has_following_word = groups.peek().is_some();
        let is_prefix_dfa = has_following_word || has_end_whitespace || word.chars().all(is_cjk);
        let words = synonyms.get(word).cloned().unwrap_or_default().iter().chain(Some(&word));

        for word in words {
            let lev = if is_prefix_dfa {
                automaton::build_dfa(word)
            } else {
                automaton::build_prefix_dfa(word)
            };
            automatons.push((index, lev));
        }

        index += 1;
    }

    automatons
}

pub struct QueryBuilder<'c, S, FI = fn(DocumentId) -> bool> {
    store: S,
    criteria: Criteria<'c>,
    searchable_attrs: Option<HashSet<u16>>,
    filter: Option<FI>,
}

impl<'c, S> QueryBuilder<'c, S, fn(DocumentId) -> bool> {
    pub fn new(store: S) -> Self {
        QueryBuilder::with_criteria(store, Criteria::default())
    }

    pub fn with_criteria(store: S, criteria: Criteria<'c>) -> Self {
        QueryBuilder { store, criteria, searchable_attrs: None, filter: None }
    }
}

impl<'c, S, FI> QueryBuilder<'c, S, FI>
{
    pub fn with_filter<F>(self, function: F) -> QueryBuilder<'c, S, F>
    where F: Fn(DocumentId) -> bool,
    {
        QueryBuilder {
            store: self.store,
            criteria: self.criteria,
            searchable_attrs: self.searchable_attrs,
            filter: Some(function)
        }
    }

    pub fn with_distinct<F, K>(self, function: F, size: usize) -> DistinctQueryBuilder<'c, S, FI, F>
    where F: Fn(DocumentId) -> Option<K>,
          K: Hash + Eq,
    {
        DistinctQueryBuilder { inner: self, function, size }
    }

    pub fn add_searchable_attribute(&mut self, attribute: u16) {
        let attributes = self.searchable_attrs.get_or_insert_with(HashSet::new);
        attributes.insert(attribute);
    }
}

impl<'c, S, FI> QueryBuilder<'c, S, FI>
where S: Store,
{
    fn query_all(&self, query: &str) -> Result<Vec<RawDocument>, S::Error> {
        let map = {
            let mut map = HashMap::new();

            map.insert("hello", &["bonjour", "salut"][..]);
            map.insert("bonjour", &["hello", "salut"]);
            map.insert("salut", &["hello", "bonjour"]);

            map
        };

        let automatons = generate_automatons(query, &map);
        let words = self.store.words()?.as_fst();

        let mut stream = {
            let mut op_builder = fst::raw::OpBuilder::new();
            for (_index, automaton) in &automatons {
                let stream = words.search(automaton);
                op_builder.push(stream);
            }
            op_builder.r#union()
        };

        let mut matches = Vec::new();

        while let Some((input, indexed_values)) = stream.next() {
            for iv in indexed_values {
                let (index, automaton) = &automatons[iv.index];
                let distance = automaton.eval(input).to_u8();
                let is_exact = distance == 0 && input.len() == automaton.query_len();

                let doc_indexes = self.store.word_indexes(input)?;
                let doc_indexes = match doc_indexes {
                    Some(doc_indexes) => doc_indexes,
                    None => continue,
                };

                for di in doc_indexes.as_slice() {
                    if self.searchable_attrs.as_ref().map_or(true, |r| r.contains(&di.attribute)) {
                        let match_ = Match {
                            query_index: *index as u32,
                            distance,
                            attribute: di.attribute,
                            word_index: di.word_index,
                            is_exact,
                            char_index: di.char_index,
                            char_length: di.char_length,
                        };
                        matches.push((di.document_id, match_));
                    }
                }
            }
        }

        let total_matches = matches.len();
        let raw_documents = raw_documents_from_matches(matches);

        info!("{} total documents to classify", raw_documents.len());
        info!("{} total matches to classify", total_matches);

        Ok(raw_documents)
    }
}

impl<'c, S, FI> QueryBuilder<'c, S, FI>
where S: Store,
      FI: Fn(DocumentId) -> bool,
{
    pub fn query(self, query: &str, range: Range<usize>) -> Result<Vec<Document>, S::Error> {
        // We delegate the filter work to the distinct query builder,
        // specifying a distinct rule that has no effect.
        if self.filter.is_some() {
            let builder = self.with_distinct(|_| None as Option<()>, 1);
            return builder.query(query, range);
        }

        let start = Instant::now();
        let mut documents = self.query_all(query)?;
        info!("query_all took {:.2?}", start.elapsed());

        let mut groups = vec![documents.as_mut_slice()];

        'criteria: for criterion in self.criteria.as_ref() {
            let tmp_groups = mem::replace(&mut groups, Vec::new());
            let mut documents_seen = 0;

            for group in tmp_groups {
                // if this group does not overlap with the requested range,
                // push it without sorting and splitting it
                if documents_seen + group.len() < range.start {
                    documents_seen += group.len();
                    groups.push(group);
                    continue;
                }

                let start = Instant::now();
                group.par_sort_unstable_by(|a, b| criterion.evaluate(a, b));
                info!("criterion {} sort took {:.2?}", criterion.name(), start.elapsed());

                for group in group.binary_group_by_mut(|a, b| criterion.eq(a, b)) {
                    info!("criterion {} produced a group of size {}", criterion.name(), group.len());

                    documents_seen += group.len();
                    groups.push(group);

                    // we have sort enough documents if the last document sorted is after
                    // the end of the requested range, we can continue to the next criterion
                    if documents_seen >= range.end { continue 'criteria }
                }
            }
        }

        let offset = cmp::min(documents.len(), range.start);
        let iter = documents.into_iter().skip(offset).take(range.len());
        Ok(iter.map(|d| Document::from_raw(&d)).collect())
    }
}

pub struct DistinctQueryBuilder<'c, I, FI, FD> {
    inner: QueryBuilder<'c, I, FI>,
    function: FD,
    size: usize,
}

impl<'c, I, FI, FD> DistinctQueryBuilder<'c, I, FI, FD>
{
    pub fn with_filter<F>(self, function: F) -> DistinctQueryBuilder<'c, I, F, FD>
    where F: Fn(DocumentId) -> bool,
    {
        DistinctQueryBuilder {
            inner: self.inner.with_filter(function),
            function: self.function,
            size: self.size
        }
    }

    pub fn add_searchable_attribute(&mut self, attribute: u16) {
        self.inner.add_searchable_attribute(attribute);
    }
}

impl<'c, S, FI, FD, K> DistinctQueryBuilder<'c, S, FI, FD>
where S: Store,
      FI: Fn(DocumentId) -> bool,
      FD: Fn(DocumentId) -> Option<K>,
      K: Hash + Eq,
{
    pub fn query(self, query: &str, range: Range<usize>) -> Result<Vec<Document>, S::Error> {
        let start = Instant::now();
        let mut documents = self.inner.query_all(query)?;
        info!("query_all took {:.2?}", start.elapsed());

        let mut groups = vec![documents.as_mut_slice()];
        let mut key_cache = HashMap::new();

        let mut filter_map = HashMap::new();
        // these two variables informs on the current distinct map and
        // on the raw offset of the start of the group where the
        // range.start bound is located according to the distinct function
        let mut distinct_map = DistinctMap::new(self.size);
        let mut distinct_raw_offset = 0;

        'criteria: for criterion in self.inner.criteria.as_ref() {
            let tmp_groups = mem::replace(&mut groups, Vec::new());
            let mut buf_distinct = BufferedDistinctMap::new(&mut distinct_map);
            let mut documents_seen = 0;

            for group in tmp_groups {
                // if this group does not overlap with the requested range,
                // push it without sorting and splitting it
                if documents_seen + group.len() < distinct_raw_offset {
                    documents_seen += group.len();
                    groups.push(group);
                    continue;
                }

                let start = Instant::now();
                group.par_sort_unstable_by(|a, b| criterion.evaluate(a, b));
                info!("criterion {} sort took {:.2?}", criterion.name(), start.elapsed());

                for group in group.binary_group_by_mut(|a, b| criterion.eq(a, b)) {
                    // we must compute the real distinguished len of this sub-group
                    for document in group.iter() {
                        let filter_accepted = match &self.inner.filter {
                            Some(filter) => {
                                let entry = filter_map.entry(document.id);
                                *entry.or_insert_with(|| (filter)(document.id))
                            },
                            None => true,
                        };

                        if filter_accepted {
                            let entry = key_cache.entry(document.id);
                            let key = entry.or_insert_with(|| (self.function)(document.id).map(Rc::new));

                            match key.clone() {
                                Some(key) => buf_distinct.register(key),
                                None => buf_distinct.register_without_key(),
                            };
                        }

                        // the requested range end is reached: stop computing distinct
                        if buf_distinct.len() >= range.end { break }
                    }

                    info!("criterion {} produced a group of size {}", criterion.name(), group.len());

                    documents_seen += group.len();
                    groups.push(group);

                    // if this sub-group does not overlap with the requested range
                    // we must update the distinct map and its start index
                    if buf_distinct.len() < range.start {
                        buf_distinct.transfert_to_internal();
                        distinct_raw_offset = documents_seen;
                    }

                    // we have sort enough documents if the last document sorted is after
                    // the end of the requested range, we can continue to the next criterion
                    if buf_distinct.len() >= range.end { continue 'criteria }
                }
            }
        }

        let mut out_documents = Vec::with_capacity(range.len());
        let mut seen = BufferedDistinctMap::new(&mut distinct_map);

        for document in documents.into_iter().skip(distinct_raw_offset) {
            let filter_accepted = match &self.inner.filter {
                Some(_) => filter_map.remove(&document.id).expect("BUG: filtered not found"),
                None => true,
            };

            if filter_accepted {
                let key = key_cache.remove(&document.id).expect("BUG: cached key not found");
                let distinct_accepted = match key {
                    Some(key) => seen.register(key),
                    None => seen.register_without_key(),
                };

                if distinct_accepted && seen.len() > range.start {
                    out_documents.push(Document::from_raw(&document));
                    if out_documents.len() == range.len() { break }
                }
            }
        }

        Ok(out_documents)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    use std::collections::{BTreeSet, HashMap};
    use std::iter::FromIterator;

    use sdset::SetBuf;
    use fst::Set;

    use crate::DocIndex;
    use crate::store::Store;

    #[derive(Default)]
    struct InMemorySetStore {
        set: Set,
        indexes: HashMap<Vec<u8>, SetBuf<DocIndex>>,
    }

    impl Store for InMemorySetStore {
        type Error = std::io::Error;

        fn words(&self) -> Result<&Set, Self::Error> {
            Ok(&self.set)
        }

        fn word_indexes(&self, word: &[u8]) -> Result<Option<SetBuf<DocIndex>>, Self::Error> {
            Ok(self.indexes.get(word).cloned())
        }
    }

    impl<'a> FromIterator<(&'a [u8], &'a [DocIndex])> for InMemorySetStore {
        fn from_iter<I: IntoIterator<Item=(&'a [u8], &'a [DocIndex])>>(iter: I) -> Self {
            let mut tree = BTreeSet::new();
            let mut map = HashMap::new();

            for (word, indexes) in iter {
                tree.insert(word);
                map.insert(word.to_vec(), SetBuf::from_dirty(indexes.to_vec()));
            }

            InMemorySetStore {
                set: Set::from_iter(tree).unwrap(),
                indexes: map,
            }
        }
    }

    const fn doc_index(document_id: u64, word_index: u16) -> DocIndex {
        DocIndex {
            document_id: DocumentId(document_id),
            attribute: 0,
            word_index,
            char_index: 0,
            char_length: 0,
        }
    }

    #[test]
    fn simple_synonymes() {
        let store = InMemorySetStore::from_iter(vec![
            (&b"hello"[..], &[doc_index(0, 0)][..]),
        ]);

        let builder = QueryBuilder::new(&store);
        let results = builder.query("hello", 0..20).unwrap();
        let mut iter = results.into_iter();

        assert_matches!(iter.next(), Some(Document { id: DocumentId(0), matches }) => {
            assert_eq!(matches.len(), 1);
            let match_ = matches[0];
            assert_eq!(match_.query_index, 0);
            assert_eq!(match_.word_index, 0);
        });
        assert_matches!(iter.next(), None);

        let builder = QueryBuilder::new(&store);
        let results = builder.query("bonjour", 0..20).unwrap();
        let mut iter = results.into_iter();

        assert_matches!(iter.next(), Some(Document { id: DocumentId(0), matches }) => {
            assert_eq!(matches.len(), 1);
            let match_ = matches[0];
            assert_eq!(match_.query_index, 0);
            assert_eq!(match_.word_index, 0);
        });
        assert_matches!(iter.next(), None);
    }

    #[test]
    fn harder_synonymes() {
        let store = InMemorySetStore::from_iter(vec![
            (&b"hello"[..],     &[doc_index(0, 0)][..]),
            (&b"bonjour"[..],   &[doc_index(1, 3)]),
            (&b"salut"[..],     &[doc_index(2, 5)]),
        ]);

        let builder = QueryBuilder::new(&store);
        let results = builder.query("hello", 0..20).unwrap();
        let mut iter = results.into_iter();

        assert_matches!(iter.next(), Some(Document { id: DocumentId(0), matches }) => {
            assert_eq!(matches.len(), 1);
            let match_ = matches[0];
            assert_eq!(match_.query_index, 0);
            assert_eq!(match_.word_index, 0);
        });
        assert_matches!(iter.next(), Some(Document { id: DocumentId(1), matches }) => {
            assert_eq!(matches.len(), 1);
            let match_ = matches[0];
            assert_eq!(match_.query_index, 0);
            assert_eq!(match_.word_index, 3);
        });
        assert_matches!(iter.next(), Some(Document { id: DocumentId(2), matches }) => {
            assert_eq!(matches.len(), 1);
            let match_ = matches[0];
            assert_eq!(match_.query_index, 0);
            assert_eq!(match_.word_index, 5);
        });
        assert_matches!(iter.next(), None);

        let builder = QueryBuilder::new(&store);
        let results = builder.query("bonjour", 0..20).unwrap();
        let mut iter = results.into_iter();

        assert_matches!(iter.next(), Some(Document { id: DocumentId(0), matches }) => {
            assert_eq!(matches.len(), 1);
            let match_ = matches[0];
            assert_eq!(match_.query_index, 0);
            assert_eq!(match_.word_index, 0);
        });
        assert_matches!(iter.next(), Some(Document { id: DocumentId(1), matches }) => {
            assert_eq!(matches.len(), 1);
            let match_ = matches[0];
            assert_eq!(match_.query_index, 0);
            assert_eq!(match_.word_index, 3);
        });
        assert_matches!(iter.next(), Some(Document { id: DocumentId(2), matches }) => {
            assert_eq!(matches.len(), 1);
            let match_ = matches[0];
            assert_eq!(match_.query_index, 0);
            assert_eq!(match_.word_index, 5);
        });
        assert_matches!(iter.next(), None);

        let builder = QueryBuilder::new(&store);
        let results = builder.query("salut", 0..20).unwrap();
        let mut iter = results.into_iter();

        assert_matches!(iter.next(), Some(Document { id: DocumentId(0), matches }) => {
            assert_eq!(matches.len(), 1);
            let match_ = matches[0];
            assert_eq!(match_.query_index, 0);
            assert_eq!(match_.word_index, 0);
        });
        assert_matches!(iter.next(), Some(Document { id: DocumentId(1), matches }) => {
            assert_eq!(matches.len(), 1);
            let match_ = matches[0];
            assert_eq!(match_.query_index, 0);
            assert_eq!(match_.word_index, 3);
        });
        assert_matches!(iter.next(), Some(Document { id: DocumentId(2), matches }) => {
            assert_eq!(matches.len(), 1);
            let match_ = matches[0];
            assert_eq!(match_.query_index, 0);
            assert_eq!(match_.word_index, 5);
        });
        assert_matches!(iter.next(), None);
    }
}