Rework the QueryBuilder to make it easier to construct and use

meilidb-core/examples/from_file.rs

@@ -268,9 +268,10 @@ fn crop_text(
 fn search_command(command: SearchCommand, database: Database) -> Result<(), Box<dyn Error>> {
     let env = &database.env;
     let index = database.open_index(INDEX_NAME).expect("Could not find index");
-    let reader = env.read_txn().unwrap();
 
+    let reader = env.read_txn().unwrap();
     let schema = index.main.schema(&reader)?;
+    reader.abort();
     let schema = schema.ok_or(meilidb_core::Error::SchemaMissing)?;
 
     let fields = command.displayed_fields.iter().map(String::as_str);
@@ -285,15 +286,16 @@ fn search_command(command: SearchCommand, database: Database) -> Result<(), Box<
             Ok(query) => {
                 let start_total = Instant::now();
 
-                let builder = index.query_builder();
+                let reader = env.read_txn().unwrap();
-                let builder = if let Some(timeout) = command.fetch_timeout_ms {
+                let ref_index = &index;
-                    builder.with_fetch_timeout(Duration::from_millis(timeout))
+                let ref_reader = &reader;
-                } else {
-                    builder
-                };
 
-                let documents = match command.filter {
+                let mut builder = index.query_builder();
-                    Some(ref filter) => {
+                if let Some(timeout) = command.fetch_timeout_ms {
+                    builder.with_fetch_timeout(Duration::from_millis(timeout));
+                }
+
+                if let Some(ref filter) = command.filter {
                     let filter = filter.as_str();
                     let (positive, filter) = if filter.chars().next() == Some('!') {
                         (false, &filter[1..])
@@ -303,17 +305,13 @@ fn search_command(command: SearchCommand, database: Database) -> Result<(), Box<
 
                     let attr = schema.attribute(&filter).expect("Could not find filtered attribute");
 
-                        let builder = builder.with_filter(|document_id| {
+                    builder.with_filter(move |document_id| {
-                            let string: String = index.document_attribute(&reader, document_id, attr).unwrap().unwrap();
+                        let string: String = ref_index.document_attribute(ref_reader, document_id, attr).unwrap().unwrap();
                         (string == "true") == positive
                     });
-
-                        builder.query(&reader, &query, 0..command.number_results)?
-                    },
-                    None => {
-                        builder.query(&reader, &query, 0..command.number_results)?
                 }
-                };
+
+                let documents = builder.query(ref_reader, &query, 0..command.number_results)?;
 
                 let mut retrieve_duration = Duration::default();
 

meilidb-core/src/query_builder.rs

@@ -1,5 +1,4 @@
 use hashbrown::HashMap;
-use std::hash::Hash;
 use std::mem;
 use std::ops::Range;
 use std::rc::Rc;
@@ -15,10 +14,11 @@ use crate::raw_document::{RawDocument, raw_documents_from};
 use crate::{Document, DocumentId, Highlight, TmpMatch, criterion::Criteria};
 use crate::{store, MResult, reordered_attrs::ReorderedAttrs};
 
-pub struct QueryBuilder<'c, FI = fn(DocumentId) -> bool> {
+pub struct QueryBuilder<'c, 'f, 'd> {
     criteria: Criteria<'c>,
     searchable_attrs: Option<ReorderedAttrs>,
-    filter: Option<FI>,
+    filter: Option<Box<dyn Fn(DocumentId) -> bool + 'f>>,
+    distinct: Option<(Box<dyn Fn(DocumentId) -> Option<u64> + 'd>, usize)>,
     timeout: Option<Duration>,
     main_store: store::Main,
     postings_lists_store: store::PostingsLists,
@@ -204,13 +204,13 @@ fn fetch_raw_documents(
     Ok(raw_documents_from(matches, highlights, fields_counts))
 }
 
-impl<'c> QueryBuilder<'c> {
+impl<'c, 'f, 'd> QueryBuilder<'c, 'f, 'd> {
     pub fn new(
         main: store::Main,
         postings_lists: store::PostingsLists,
         documents_fields_counts: store::DocumentsFieldsCounts,
         synonyms: store::Synonyms,
-    ) -> QueryBuilder<'c>
+    ) -> QueryBuilder<'c, 'f, 'd>
     {
         QueryBuilder::with_criteria(
             main,
@@ -227,12 +227,13 @@ impl<'c> QueryBuilder<'c> {
         documents_fields_counts: store::DocumentsFieldsCounts,
         synonyms: store::Synonyms,
         criteria: Criteria<'c>,
-    ) -> QueryBuilder<'c>
+    ) -> QueryBuilder<'c, 'f, 'd>
     {
         QueryBuilder {
             criteria,
             searchable_attrs: None,
             filter: None,
+            distinct: None,
             timeout: None,
             main_store: main,
             postings_lists_store: postings_lists,
@@ -242,52 +243,111 @@ impl<'c> QueryBuilder<'c> {
     }
 }
 
-impl<'c, FI> QueryBuilder<'c, FI> {
+impl<'c, 'f, 'd> QueryBuilder<'c, 'f, 'd> {
-    pub fn with_filter<F>(self, function: F) -> QueryBuilder<'c, F>
+    pub fn with_filter<F>(&mut self, function: F)
-    where F: Fn(DocumentId) -> bool,
+    where F: Fn(DocumentId) -> bool + 'f,
     {
-        QueryBuilder {
+        self.filter = Some(Box::new(function))
-            criteria: self.criteria,
-            searchable_attrs: self.searchable_attrs,
-            filter: Some(function),
-            timeout: self.timeout,
-            main_store: self.main_store,
-            postings_lists_store: self.postings_lists_store,
-            documents_fields_counts_store: self.documents_fields_counts_store,
-            synonyms_store: self.synonyms_store,
-        }
     }
 
-    pub fn with_fetch_timeout(self, timeout: Duration) -> QueryBuilder<'c, FI> {
+    pub fn with_fetch_timeout(&mut self, timeout: Duration) {
-        QueryBuilder { timeout: Some(timeout), ..self }
+        self.timeout = Some(timeout)
     }
 
-    pub fn with_distinct<F, K>(self, function: F, size: usize) -> DistinctQueryBuilder<'c, FI, F>
+    pub fn with_distinct<F, K>(&mut self, function: F, size: usize)
-    where F: Fn(DocumentId) -> Option<K>,
+    where F: Fn(DocumentId) -> Option<u64> + 'd,
-          K: Hash + Eq,
     {
-        DistinctQueryBuilder { inner: self, function, size }
+        self.distinct = Some((Box::new(function), size))
     }
 
     pub fn add_searchable_attribute(&mut self, attribute: u16) {
         let reorders = self.searchable_attrs.get_or_insert_with(ReorderedAttrs::new);
         reorders.insert_attribute(attribute);
     }
-}
 
-impl<FI> QueryBuilder<'_, FI> where FI: Fn(DocumentId) -> bool {
     pub fn query(
         self,
         reader: &zlmdb::RoTxn,
         query: &str,
         range: Range<usize>,
     ) -> MResult<Vec<Document>>
+    {
+        match self.distinct {
+            Some((distinct, distinct_size)) => {
+                raw_query_with_distinct(
+                    reader,
+                    query,
+                    range,
+                    self.filter,
+                    distinct,
+                    distinct_size,
+                    self.timeout,
+                    self.criteria,
+                    self.searchable_attrs,
+                    self.main_store,
+                    self.postings_lists_store,
+                    self.documents_fields_counts_store,
+                    self.synonyms_store,
+                )
+            },
+            None => {
+                raw_query(
+                    reader,
+                    query,
+                    range,
+                    self.filter,
+                    self.timeout,
+                    self.criteria,
+                    self.searchable_attrs,
+                    self.main_store,
+                    self.postings_lists_store,
+                    self.documents_fields_counts_store,
+                    self.synonyms_store,
+                )
+            }
+        }
+    }
+}
+
+fn raw_query<'c, FI>(
+    reader: &zlmdb::RoTxn,
+
+    query: &str,
+    range: Range<usize>,
+
+    filter: Option<FI>,
+    timeout: Option<Duration>,
+
+    criteria: Criteria<'c>,
+    searchable_attrs: Option<ReorderedAttrs>,
+
+    main_store: store::Main,
+    postings_lists_store: store::PostingsLists,
+    documents_fields_counts_store: store::DocumentsFieldsCounts,
+    synonyms_store: store::Synonyms,
+) -> MResult<Vec<Document>>
+where FI: Fn(DocumentId) -> bool,
 {
     // We delegate the filter work to the distinct query builder,
     // specifying a distinct rule that has no effect.
-        if self.filter.is_some() {
+    if filter.is_some() {
-            let builder = self.with_distinct(|_| None as Option<()>, 1);
+        let distinct = |_| None;
-            return builder.query(reader, query, range);
+        let distinct_size = 1;
+        return raw_query_with_distinct(
+            reader,
+            query,
+            range,
+            filter,
+            distinct,
+            distinct_size,
+            timeout,
+            criteria,
+            searchable_attrs,
+            main_store,
+            postings_lists_store,
+            documents_fields_counts_store,
+            synonyms_store,
+        )
     }
 
     let start_processing = Instant::now();
@@ -296,8 +356,8 @@ impl<FI> QueryBuilder<'_, FI> where FI: Fn(DocumentId) -> bool {
     let (automaton_producer, query_enhancer) = AutomatonProducer::new(
         reader,
         query,
-            self.main_store,
+        main_store,
-            self.synonyms_store,
+        synonyms_store,
     )?;
 
     let mut automaton_producer = automaton_producer.into_iter();
@@ -313,14 +373,14 @@ impl<FI> QueryBuilder<'_, FI> where FI: Fn(DocumentId) -> bool {
             reader,
             &automatons,
             &query_enhancer,
-                self.searchable_attrs.as_ref(),
+            searchable_attrs.as_ref(),
-                &self.main_store,
+            &main_store,
-                &self.postings_lists_store,
+            &postings_lists_store,
-                &self.documents_fields_counts_store,
+            &documents_fields_counts_store,
         )?;
 
         // stop processing when time is running out
-            if let Some(timeout) = self.timeout {
+        if let Some(timeout) = timeout {
             if !raw_documents_processed.is_empty() && start_processing.elapsed() > timeout {
                 break
             }
@@ -328,7 +388,7 @@ impl<FI> QueryBuilder<'_, FI> where FI: Fn(DocumentId) -> bool {
 
         let mut groups = vec![raw_documents.as_mut_slice()];
 
-            'criteria: for criterion in self.criteria.as_ref() {
+        'criteria: for criterion in criteria.as_ref() {
             let tmp_groups = mem::replace(&mut groups, Vec::new());
             let mut documents_seen = 0;
 
@@ -361,7 +421,7 @@ impl<FI> QueryBuilder<'_, FI> where FI: Fn(DocumentId) -> bool {
         raw_documents_processed.extend(iter);
 
         // stop processing when time is running out
-            if let Some(timeout) = self.timeout {
+        if let Some(timeout) = timeout {
             if start_processing.elapsed() > timeout { break }
         }
     }
@@ -375,49 +435,29 @@ impl<FI> QueryBuilder<'_, FI> where FI: Fn(DocumentId) -> bool {
 
     Ok(documents)
 }
-}
 
-pub struct DistinctQueryBuilder<'c, FI, FD> {
+fn raw_query_with_distinct<'c, FI, FD>(
-    inner: QueryBuilder<'c, FI>,
-    function: FD,
-    size: usize,
-}
-
-impl<'c, FI, FD> DistinctQueryBuilder<'c, FI, FD> {
-    pub fn with_filter<F>(self, function: F) -> DistinctQueryBuilder<'c, F, FD>
-    where F: Fn(DocumentId) -> bool,
-    {
-        DistinctQueryBuilder {
-            inner: self.inner.with_filter(function),
-            function: self.function,
-            size: self.size,
-        }
-    }
-
-    pub fn with_fetch_timeout(self, timeout: Duration) -> DistinctQueryBuilder<'c, FI, FD> {
-        DistinctQueryBuilder {
-            inner: self.inner.with_fetch_timeout(timeout),
-            function: self.function,
-            size: self.size,
-        }
-    }
-
-    pub fn add_searchable_attribute(&mut self, attribute: u16) {
-        self.inner.add_searchable_attribute(attribute);
-    }
-}
-
-impl<'c, FI, FD, K> DistinctQueryBuilder<'c, FI, FD>
-where FI: Fn(DocumentId) -> bool,
-      FD: Fn(DocumentId) -> Option<K>,
-      K: Hash + Eq,
-{
-    pub fn query(
-        self,
     reader: &zlmdb::RoTxn,
+
     query: &str,
     range: Range<usize>,
+
+    filter: Option<FI>,
+
+    distinct: FD,
+    distinct_size: usize,
+    timeout: Option<Duration>,
+
+    criteria: Criteria<'c>,
+    searchable_attrs: Option<ReorderedAttrs>,
+
+    main_store: store::Main,
+    postings_lists_store: store::PostingsLists,
+    documents_fields_counts_store: store::DocumentsFieldsCounts,
+    synonyms_store: store::Synonyms,
 ) -> MResult<Vec<Document>>
+where FI: Fn(DocumentId) -> bool,
+      FD: Fn(DocumentId) -> Option<u64>,
 {
     let start_processing = Instant::now();
     let mut raw_documents_processed = Vec::new();
@@ -425,8 +465,8 @@ where FI: Fn(DocumentId) -> bool,
     let (automaton_producer, query_enhancer) = AutomatonProducer::new(
         reader,
         query,
-            self.inner.main_store,
+        main_store,
-            self.inner.synonyms_store,
+        synonyms_store,
     )?;
 
     let mut automaton_producer = automaton_producer.into_iter();
@@ -442,14 +482,14 @@ where FI: Fn(DocumentId) -> bool,
             reader,
             &automatons,
             &query_enhancer,
-                self.inner.searchable_attrs.as_ref(),
+            searchable_attrs.as_ref(),
-                &self.inner.main_store,
+            &main_store,
-                &self.inner.postings_lists_store,
+            &postings_lists_store,
-                &self.inner.documents_fields_counts_store,
+            &documents_fields_counts_store,
         )?;
 
         // stop processing when time is running out
-            if let Some(timeout) = self.inner.timeout {
+        if let Some(timeout) = timeout {
             if !raw_documents_processed.is_empty() && start_processing.elapsed() > timeout {
                 break
             }
@@ -462,10 +502,10 @@ where FI: Fn(DocumentId) -> bool,
         // 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_map = DistinctMap::new(distinct_size);
         let mut distinct_raw_offset = 0;
 
-            'criteria: for criterion in self.inner.criteria.as_ref() {
+        'criteria: for criterion in 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;
@@ -484,7 +524,7 @@ where FI: Fn(DocumentId) -> bool,
                 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 {
+                        let filter_accepted = match &filter {
                             Some(filter) => {
                                 let entry = filter_map.entry(document.id);
                                 *entry.or_insert_with(|| (filter)(document.id))
@@ -494,7 +534,7 @@ where FI: Fn(DocumentId) -> bool,
 
                         if filter_accepted {
                             let entry = key_cache.entry(document.id);
-                                let key = entry.or_insert_with(|| (self.function)(document.id).map(Rc::new));
+                            let key = entry.or_insert_with(|| (distinct)(document.id).map(Rc::new));
 
                             match key.clone() {
                                 Some(key) => buf_distinct.register(key),
@@ -529,7 +569,7 @@ where FI: Fn(DocumentId) -> bool,
         raw_documents_processed.clear();
 
         for document in raw_documents.into_iter().skip(distinct_raw_offset) {
-                let filter_accepted = match &self.inner.filter {
+            let filter_accepted = match &filter {
                 Some(_) => filter_map.remove(&document.id).unwrap(),
                 None => true,
             };
@@ -549,7 +589,7 @@ where FI: Fn(DocumentId) -> bool,
         }
 
         // stop processing when time is running out
-            if let Some(timeout) = self.inner.timeout {
+        if let Some(timeout) = timeout {
             if start_processing.elapsed() > timeout { break }
         }
     }
@@ -563,7 +603,6 @@ where FI: Fn(DocumentId) -> bool,
 
     Ok(documents)
 }
-}
 
 #[cfg(test)]
 mod tests {

meilidb-core/src/store/mod.rs

@@ -202,7 +202,11 @@ impl Index {
         )
     }
 
-    pub fn query_builder_with_criteria<'c>(&self, criteria: Criteria<'c>) -> QueryBuilder<'c> {
+    pub fn query_builder_with_criteria<'c, 'f, 'd>(
+        &self,
+        criteria: Criteria<'c>,
+    ) -> QueryBuilder<'c, 'f, 'd>
+    {
         QueryBuilder::with_criteria(
             self.main,
             self.postings_lists,