Merge pull request #32 from Kerollmops/query-range

Improve the query interface

examples/query-database.rs

@@ -46,7 +46,7 @@ fn main() -> Result<(), Box<Error>> {
 
         let (elapsed, documents) = elapsed::measure_time(|| {
             let builder = view.query_builder().unwrap();
-            builder.query(&buffer, opt.number_results)
+            builder.query(&buffer, 0..opt.number_results)
         });
 
         let mut full_documents = Vec::with_capacity(documents.len());

src/rank/distinct_map.rs

@@ -17,10 +17,33 @@ impl<K: Hash + Eq> DistinctMap<K> {
         }
     }
 
+    pub fn len(&self) -> usize {
+        self.len
+    }
+}
+
+pub struct BufferedDistinctMap<'a, K> {
+    internal: &'a mut DistinctMap<K>,
+    inner: HashMap<K, usize>,
+    len: usize,
+}
+
+impl<'a, K: Hash + Eq> BufferedDistinctMap<'a, K> {
+    pub fn new(internal: &'a mut DistinctMap<K>) -> BufferedDistinctMap<'a, K> {
+        BufferedDistinctMap {
+            internal: internal,
+            inner: HashMap::new(),
+            len: 0,
+        }
+    }
+
     pub fn register(&mut self, key: K) -> bool {
-        let seen = self.inner.entry(key).or_insert(0);
-        if *seen < self.limit {
-            *seen += 1;
+        let internal_seen = self.internal.inner.get(&key).unwrap_or(&0);
+        let inner_seen = self.inner.entry(key).or_insert(0);
+        let seen = *internal_seen + *inner_seen;
+
+        if seen < self.internal.limit {
+            *inner_seen += 1;
             self.len += 1;
             true
         } else {
@@ -33,8 +56,18 @@ impl<K: Hash + Eq> DistinctMap<K> {
         true
     }
 
+    pub fn transfert_to_internal(&mut self) {
+        for (k, v) in self.inner.drain() {
+            let value = self.internal.inner.entry(k).or_insert(0);
+            *value += v;
+        }
+
+        self.internal.len += self.len;
+        self.len = 0;
+    }
+
     pub fn len(&self) -> usize {
-        self.len
+        self.internal.len() + self.len
     }
 }
 
@@ -45,22 +78,27 @@ mod tests {
     #[test]
     fn easy_distinct_map() {
         let mut map = DistinctMap::new(2);
+        let mut buffered = BufferedDistinctMap::new(&mut map);
+
         for x in &[1, 1, 1, 2, 3, 4, 5, 6, 6, 6, 6, 6] {
-            map.register(x);
+            buffered.register(x);
         }
+        buffered.transfert_to_internal();
         assert_eq!(map.len(), 8);
 
         let mut map = DistinctMap::new(2);
-        assert_eq!(map.register(1), true);
-        assert_eq!(map.register(1), true);
-        assert_eq!(map.register(1), false);
-        assert_eq!(map.register(1), false);
+        let mut buffered = BufferedDistinctMap::new(&mut map);
+        assert_eq!(buffered.register(1), true);
+        assert_eq!(buffered.register(1), true);
+        assert_eq!(buffered.register(1), false);
+        assert_eq!(buffered.register(1), false);
 
-        assert_eq!(map.register(2), true);
-        assert_eq!(map.register(3), true);
-        assert_eq!(map.register(2), true);
-        assert_eq!(map.register(2), false);
+        assert_eq!(buffered.register(2), true);
+        assert_eq!(buffered.register(3), true);
+        assert_eq!(buffered.register(2), true);
+        assert_eq!(buffered.register(2), false);
 
+        buffered.transfert_to_internal();
         assert_eq!(map.len(), 5);
     }
 }

src/rank/query_builder.rs

@@ -1,7 +1,8 @@
-use std::{mem, vec, str, char};
+use std::{cmp, mem, vec, str, char};
+use std::ops::{Deref, Range};
 use std::error::Error;
-use std::ops::Deref;
 use std::hash::Hash;
+use std::rc::Rc;
 
 use group_by::GroupByMut;
 use hashbrown::HashMap;
@@ -9,7 +10,7 @@ use fst::Streamer;
 use rocksdb::DB;
 
 use crate::automaton::{self, DfaExt, AutomatonExt};
-use crate::rank::distinct_map::DistinctMap;
+use crate::rank::distinct_map::{DistinctMap, BufferedDistinctMap};
 use crate::rank::criterion::Criteria;
 use crate::database::DatabaseView;
 use crate::{Match, DocumentId};
@@ -111,26 +112,42 @@ where D: Deref<Target=DB>
 impl<'a, D> QueryBuilder<'a, D>
 where D: Deref<Target=DB>,
 {
-    pub fn query(&self, query: &str, limit: usize) -> Vec<Document> {
+    pub fn query(&self, query: &str, range: Range<usize>) -> Vec<Document> {
         let mut documents = self.query_all(query);
         let mut groups = vec![documents.as_mut_slice()];
         let view = &self.view;
 
-        for criterion in self.criteria.as_ref() {
+        'criteria: for criterion in self.criteria.as_ref() {
             let tmp_groups = mem::replace(&mut groups, Vec::new());
-            let mut computed = 0;
+            let mut documents_seen = 0;
 
-            'group: for group in tmp_groups {
-                group.sort_unstable_by(|a, b| criterion.evaluate(a, b, view));
-                for group in GroupByMut::new(group, |a, b| criterion.eq(a, b, view)) {
-                    computed += group.len();
+            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);
-                    if computed >= limit { break 'group }
+                    continue;
+                }
+
+                group.sort_unstable_by(|a, b| criterion.evaluate(a, b, view));
+
+                for group in GroupByMut::new(group, |a, b| criterion.eq(a, b, view)) {
+                    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 }
                 }
             }
         }
 
-        documents.truncate(limit);
+        // `drain` removes the documents efficiently using `ptr::copy`
+        // TODO it could be more efficient to have a custom iterator
+        let offset = cmp::min(documents.len(), range.start);
+        documents.drain(0..offset);
+        documents.truncate(range.len());
         documents
     }
 }
@@ -148,42 +165,80 @@ where D: Deref<Target=DB>,
       F: Fn(DocumentId, &DatabaseView<D>) -> Option<K>,
       K: Hash + Eq,
 {
-    pub fn query(&self, query: &str, limit: usize) -> Vec<Document> {
+    pub fn query(&self, query: &str, range: Range<usize>) -> Vec<Document> {
         let mut documents = self.inner.query_all(query);
         let mut groups = vec![documents.as_mut_slice()];
+        let mut key_cache = HashMap::new();
         let view = &self.inner.view;
 
-        for criterion in self.inner.criteria.as_ref() {
+        // 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 seen = DistinctMap::new(self.size);
+            let mut buf_distinct = BufferedDistinctMap::new(&mut distinct_map);
+            let mut documents_seen = 0;
 
-            'group: for group in tmp_groups {
-                group.sort_unstable_by(|a, b| criterion.evaluate(a, b, view));
-                for group in GroupByMut::new(group, |a, b| criterion.eq(a, b, view)) {
-                    for document in group.iter() {
-                        match (self.function)(document.id, view) {
-                            Some(key) => seen.register(key),
-                            None => seen.register_without_key(),
-                        };
-                    }
+            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);
-                    if seen.len() >= limit { break 'group }
+                    continue;
+                }
+
+                group.sort_unstable_by(|a, b| criterion.evaluate(a, b, view));
+
+                for group in GroupByMut::new(group, |a, b| criterion.eq(a, b, view)) {
+                    // we must compute the real distinguished len of this sub-group
+                    for document in group.iter() {
+                        let entry = key_cache.entry(document.id);
+                        let key = entry.or_insert_with(|| (self.function)(document.id, view).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 }
+                    }
+
+                    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(limit);
-        let mut seen = DistinctMap::new(self.size);
+        let mut out_documents = Vec::with_capacity(range.len());
+        let mut seen = BufferedDistinctMap::new(&mut distinct_map);
 
-        for document in documents {
-            let accepted = match (self.function)(document.id, view) {
+        for document in documents.into_iter().skip(distinct_raw_offset) {
+            let key = key_cache.remove(&document.id).expect("BUG: cached key not found");
+
+            let accepted = match key {
                 Some(key) => seen.register(key),
                 None      => seen.register_without_key(),
             };
 
-            if accepted {
+            if accepted && seen.len() > range.start {
                 out_documents.push(document);
-                if out_documents.len() == limit { break }
+                if out_documents.len() == range.len() { break }
             }
         }