From 858589dc6beee337f31b1d499ba5024e52b462bc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cl=C3=A9ment=20Renault?= <renault.cle@gmail.com>
Date: Tue, 5 Mar 2019 16:34:29 +0100
Subject: [PATCH] feat: Limit the QueryBuilder to search only into some
 attributes

---
 src/database/schema.rs    |  2 +-
 src/rank/query_builder.rs | 39 ++++++++++++++++++++++++++-------------
 2 files changed, 27 insertions(+), 14 deletions(-)

diff --git a/src/database/schema.rs b/src/database/schema.rs
index 7ae4289cf..fc64ffccc 100644
--- a/src/database/schema.rs
+++ b/src/database/schema.rs
@@ -195,7 +195,7 @@ impl Schema {
 
 #[derive(Serialize, Deserialize)]
 #[derive(Debug, Copy, Clone, PartialOrd, Ord, PartialEq, Eq, Hash)]
-pub struct SchemaAttr(pub(crate) u16);
+pub struct SchemaAttr(pub u16);
 
 impl SchemaAttr {
     pub fn new(value: u16) -> SchemaAttr {
diff --git a/src/rank/query_builder.rs b/src/rank/query_builder.rs
index 810d77a26..6b145b493 100644
--- a/src/rank/query_builder.rs
+++ b/src/rank/query_builder.rs
@@ -6,7 +6,7 @@ use std::rc::Rc;
 
 use rayon::slice::ParallelSliceMut;
 use slice_group_by::{GroupByMut, LinearStrGroupBy};
-use hashbrown::HashMap;
+use hashbrown::{HashMap, HashSet};
 use fst::Streamer;
 use log::info;
 
@@ -66,6 +66,7 @@ pub type FilterFunc = fn(DocumentId) -> bool;
 pub struct QueryBuilder<'i, 'c, FI> {
     index: &'i Index,
     criteria: Criteria<'c>,
+    searchable_attrs: Option<HashSet<u16>>,
     filter: Option<FI>,
 }
 
@@ -75,7 +76,7 @@ impl<'i, 'c> QueryBuilder<'i, 'c, FilterFunc> {
     }
 
     pub fn with_criteria(index: &'i Index, criteria: Criteria<'c>) -> Self {
-        QueryBuilder { index, criteria, filter: None }
+        QueryBuilder { index, criteria, searchable_attrs: None, filter: None }
     }
 }
 
@@ -87,6 +88,7 @@ impl<'i, 'c, FI> QueryBuilder<'i, 'c, FI>
         QueryBuilder {
             index: self.index,
             criteria: self.criteria,
+            searchable_attrs: self.searchable_attrs,
             filter: Some(function)
         }
     }
@@ -102,6 +104,11 @@ impl<'i, 'c, FI> QueryBuilder<'i, 'c, FI>
         }
     }
 
+    pub fn add_searchable_attribute(&mut self, attribute: u16) {
+        let attributes = self.searchable_attrs.get_or_insert_with(HashSet::new);
+        attributes.insert(attribute);
+    }
+
     fn query_all(&self, query: &str) -> Vec<RawDocument> {
         let automatons = split_whitespace_automatons(query);
 
@@ -125,17 +132,19 @@ impl<'i, 'c, FI> QueryBuilder<'i, 'c, FI>
                 let doc_indexes = &self.index.indexes;
                 let doc_indexes = &doc_indexes[iv.value as usize];
 
-                for doc_index in doc_indexes {
-                    let match_ = Match {
-                        query_index: iv.index as u32,
-                        distance: distance,
-                        attribute: doc_index.attribute,
-                        word_index: doc_index.word_index,
-                        is_exact: is_exact,
-                        char_index: doc_index.char_index,
-                        char_length: doc_index.char_length,
-                    };
-                    matches.push((doc_index.document_id, match_));
+                for di in doc_indexes {
+                    if self.searchable_attrs.as_ref().map_or(true, |r| r.contains(&di.attribute)) {
+                        let match_ = Match {
+                            query_index: iv.index as u32,
+                            distance: distance,
+                            attribute: di.attribute,
+                            word_index: di.word_index,
+                            is_exact: is_exact,
+                            char_index: di.char_index,
+                            char_length: di.char_length,
+                        };
+                        matches.push((di.document_id, match_));
+                    }
                 }
             }
         }
@@ -220,6 +229,10 @@ impl<'i, 'c, FI, FD> DistinctQueryBuilder<'i, 'c, FI, FD>
             size: self.size
         }
     }
+
+    pub fn add_searchable_attribute(&mut self, attribute: u16) {
+        self.inner.add_searchable_attribute(attribute);
+    }
 }
 
 impl<'i, 'c, FI, FD, K> DistinctQueryBuilder<'i, 'c, FI, FD>