feat: Introduce a way to distinct documents

src/lib.rs

@@ -24,7 +24,7 @@ pub type DocumentId = u64;
 #[repr(C)]
 pub struct DocIndex {
     /// The document identifier where the word was found.
-    pub document: DocumentId,
+    pub document_id: DocumentId,
 
     /// The attribute identifier in the document
     /// where the word was found.

src/rank/criterion/mod.rs

@@ -7,7 +7,6 @@ mod exact;
 
 use std::vec;
 use std::cmp::Ordering;
-use std::ops::Deref;
 use crate::rank::Document;
 
 pub use self::{
@@ -31,21 +30,21 @@ pub trait Criterion {
 
 impl<'a, T: Criterion + ?Sized> Criterion for &'a T {
     fn evaluate(&self, lhs: &Document, rhs: &Document) -> Ordering {
-        self.deref().evaluate(lhs, rhs)
+        (**self).evaluate(lhs, rhs)
     }
 
     fn eq(&self, lhs: &Document, rhs: &Document) -> bool {
-        self.deref().eq(lhs, rhs)
+        (**self).eq(lhs, rhs)
     }
 }
 
 impl<T: Criterion + ?Sized> Criterion for Box<T> {
     fn evaluate(&self, lhs: &Document, rhs: &Document) -> Ordering {
-        self.deref().evaluate(lhs, rhs)
+        (**self).evaluate(lhs, rhs)
     }
 
     fn eq(&self, lhs: &Document, rhs: &Document) -> bool {
-        self.deref().eq(lhs, rhs)
+        (**self).eq(lhs, rhs)
     }
 }
 

src/rank/mod.rs

@@ -3,7 +3,7 @@ mod ranked_stream;
 
 use crate::{Match, DocumentId};
 
-pub use self::ranked_stream::{RankedStreamBuilder, RankedStream};
+pub use self::ranked_stream::{Config, RankedStream};
 
 #[inline]
 fn match_query_index(a: &Match, b: &Match) -> bool {

src/rank/ranked_stream.rs

@@ -1,6 +1,8 @@
+use std::collections::HashMap;
+use std::hash::Hash;
 use std::ops::Range;
 use std::rc::Rc;
-use std::{mem, vec, cmp};
+use std::{mem, vec};
 
 use fnv::FnvHashMap;
 use fst::Streamer;
@@ -9,52 +11,41 @@ use group_by::GroupByMut;
 use crate::automaton::{DfaExt, AutomatonExt};
 use crate::metadata::Metadata;
 use crate::metadata::ops::OpBuilder;
-use crate::rank::criterion::Criterion;
+use crate::rank::criterion::{self, Criterion};
 use crate::rank::Document;
-use crate::Match;
+use crate::{Match, DocumentId};
 
-#[derive(Clone)]
-pub struct RankedStreamBuilder<'m, C> {
-    metadata: &'m Metadata,
-    automatons: Vec<Rc<DfaExt>>,
-    criteria: Vec<C>,
+pub struct Config<'m, C, F> {
+    pub metadata: &'m Metadata,
+    pub automatons: Vec<DfaExt>,
+    pub criteria: Vec<C>,
+    pub distinct: (F, usize),
 }
 
-impl<'m, C> RankedStreamBuilder<'m, C> {
-    pub fn new(metadata: &'m Metadata, automatons: Vec<DfaExt>) -> Self {
-        RankedStreamBuilder {
-            metadata: metadata,
-            automatons: automatons.into_iter().map(Rc::new).collect(),
-            criteria: Vec::new(), // hummm...  prefer the criterion::default() ones !
-        }
-    }
+pub struct RankedStream<'m, C, F> {
+    stream: crate::metadata::ops::Union<'m>,
+    automatons: Vec<Rc<DfaExt>>,
+    criteria: Vec<C>,
+    distinct: (F, usize),
+}
 
-    pub fn criteria(&mut self, criteria: Vec<C>) {
-        self.criteria = criteria;
-    }
-
-    pub fn build(&self) -> RankedStream<C> {
-        let mut builder = OpBuilder::with_automatons(self.automatons.clone());
-        builder.push(self.metadata);
+impl<'m, C, F> RankedStream<'m, C, F> {
+    pub fn new(config: Config<'m, C, F>) -> Self {
+        let automatons: Vec<_> = config.automatons.into_iter().map(Rc::new).collect();
+        let mut builder = OpBuilder::with_automatons(automatons.clone());
+        builder.push(config.metadata);
 
         RankedStream {
             stream: builder.union(),
-            automatons: &self.automatons,
-            criteria: &self.criteria,
+            automatons: automatons,
+            criteria: config.criteria,
+            distinct: config.distinct,
         }
     }
 }
 
-pub struct RankedStream<'a, 'm, C> {
-    stream: crate::metadata::ops::Union<'m>,
-    automatons: &'a [Rc<DfaExt>],
-    criteria: &'a [C],
-}
-
-impl<'a, 'm, C> RankedStream<'a, 'm, C> {
-    pub fn retrieve_documents(&mut self, range: Range<usize>) -> Vec<Document>
-    where C: Criterion
-    {
+impl<'m, C, F> RankedStream<'m, C, F> {
+    fn retrieve_all_documents(&mut self) -> Vec<Document> {
         let mut matches = FnvHashMap::default();
 
         while let Some((string, indexed_values)) = self.stream.next() {
@@ -63,55 +54,133 @@ impl<'a, 'm, C> RankedStream<'a, 'm, C> {
                 let distance = automaton.eval(string).to_u8();
                 let is_exact = distance == 0 && string.len() == automaton.query_len();
 
-                for di in iv.doc_indexes.as_slice() {
+                for doc_index in iv.doc_indexes.as_slice() {
                     let match_ = Match {
                         query_index: iv.index as u32,
                         distance: distance,
-                        attribute: di.attribute,
-                        attribute_index: di.attribute_index,
+                        attribute: doc_index.attribute,
+                        attribute_index: doc_index.attribute_index,
                         is_exact: is_exact,
                     };
-                    matches.entry(di.document).or_insert_with(Vec::new).push(match_);
+                    matches.entry(doc_index.document_id).or_insert_with(Vec::new).push(match_);
                 }
             }
         }
 
-        // collect matches from an HashMap into a Vec
-        let mut documents: Vec<_> = matches.into_iter().map(|(id, mut matches)| {
+        matches.into_iter().map(|(id, mut matches)| {
             matches.sort_unstable();
             unsafe { Document::from_sorted_matches(id, matches) }
-        }).collect();
+        }).collect()
+    }
+}
 
+impl<'a, C, F> RankedStream<'a, C, F>
+where C: Criterion
+{
+    pub fn retrieve_documents(mut self, range: Range<usize>) -> Vec<Document> {
+        let mut documents = self.retrieve_all_documents();
         let mut groups = vec![documents.as_mut_slice()];
 
         for criterion in self.criteria {
             let tmp_groups = mem::replace(&mut groups, Vec::new());
-            let mut current_range = Range { start: 0, end: 0 };
 
-            'grp: for group in tmp_groups {
-                current_range.end += group.len();
-
-                // if a part of the current group is in the range returned
-                // we must sort it and emit the sub-groups
-                if current_range.contains(&range.start) {
-                    group.sort_unstable_by(|a, b| criterion.evaluate(a, b));
-                    for group in GroupByMut::new(group, |a, b| criterion.eq(a, b)) {
-                        groups.push(group);
-                        if current_range.end >= range.end { break 'grp }
-                    }
-                } else {
-                    groups.push(group)
+            for group in tmp_groups {
+                group.sort_unstable_by(|a, b| criterion.evaluate(a, b));
+                for group in GroupByMut::new(group, |a, b| criterion.eq(a, b)) {
+                    groups.push(group);
                 }
-
-                current_range.start = current_range.end;
             }
         }
 
-        // TODO find a better algorithm, here we allocate for too many documents
-        //      and we do a useless allocation, we should reuse the documents Vec
-        let start = cmp::min(range.start, documents.len());
-        let mut documents = documents.split_off(start);
-        documents.truncate(range.len());
-        documents
+        documents[range].to_vec()
+    }
+
+    pub fn retrieve_distinct_documents<K>(mut self, range: Range<usize>) -> Vec<Document>
+    where F: Fn(&DocumentId) -> K,
+          K: Hash + Eq,
+    {
+        let mut documents = self.retrieve_all_documents();
+        let mut groups = vec![documents.as_mut_slice()];
+
+        for criterion in self.criteria {
+            let tmp_groups = mem::replace(&mut groups, Vec::new());
+
+            for group in tmp_groups {
+                group.sort_unstable_by(|a, b| criterion.evaluate(a, b));
+                for group in GroupByMut::new(group, |a, b| criterion.eq(a, b)) {
+                    groups.push(group);
+                }
+            }
+        }
+
+        let mut out_documents = Vec::with_capacity(range.len());
+        let (distinct, limit) = self.distinct;
+        let mut seen = DistinctMap::new(limit);
+
+        for document in documents {
+            let key = distinct(&document.id);
+            let accepted = seen.digest(key);
+
+            if accepted {
+                if seen.len() == range.end { break }
+                if seen.len() >= range.start {
+                    out_documents.push(document);
+                }
+            }
+        }
+
+        out_documents
+    }
+}
+
+pub struct DistinctMap<K> {
+    inner: HashMap<K, usize>,
+    limit: usize,
+    len: usize,
+}
+
+impl<K: Hash + Eq> DistinctMap<K> {
+    pub fn new(limit: usize) -> Self {
+        DistinctMap {
+            inner: HashMap::new(),
+            limit: limit,
+            len: 0,
+        }
+    }
+
+    pub fn digest(&mut self, key: K) -> bool {
+        let seen = self.inner.entry(key).or_insert(0);
+        if *seen < self.limit { *seen += 1; self.len += 1; true } else { false }
+    }
+
+    pub fn len(&self) -> usize {
+        self.len
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn easy_distinct_map() {
+        let mut map = DistinctMap::new(2);
+        for x in &[1, 1, 1, 2, 3, 4, 5, 6, 6, 6, 6, 6] {
+            map.digest(x);
+        }
+        assert_eq!(map.len(), 8);
+
+        let mut map = DistinctMap::new(2);
+        assert_eq!(map.digest(1), true);
+        assert_eq!(map.digest(1), true);
+        assert_eq!(map.digest(1), false);
+        assert_eq!(map.digest(1), false);
+
+        assert_eq!(map.digest(2), true);
+        assert_eq!(map.digest(3), true);
+        assert_eq!(map.digest(2), true);
+        assert_eq!(map.digest(2), false);
+
+        assert_eq!(map.len(), 5);
     }
 }