diff --git a/examples/query-database.rs b/examples/query-database.rs index 804741751..97bf49bf7 100644 --- a/examples/query-database.rs +++ b/examples/query-database.rs @@ -46,7 +46,7 @@ fn main() -> Result<(), Box> { 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()); diff --git a/src/rank/distinct_map.rs b/src/rank/distinct_map.rs index b7a5edee3..81afe1cf0 100644 --- a/src/rank/distinct_map.rs +++ b/src/rank/distinct_map.rs @@ -17,10 +17,33 @@ impl DistinctMap { } } + pub fn len(&self) -> usize { + self.len + } +} + +pub struct BufferedDistinctMap<'a, K> { + internal: &'a mut DistinctMap, + inner: HashMap, + len: usize, +} + +impl<'a, K: Hash + Eq> BufferedDistinctMap<'a, K> { + pub fn new(internal: &'a mut DistinctMap) -> 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 DistinctMap { 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); } } diff --git a/src/rank/query_builder.rs b/src/rank/query_builder.rs index 3f8579b2a..1dff4f2f0 100644 --- a/src/rank/query_builder.rs +++ b/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 impl<'a, D> QueryBuilder<'a, D> where D: Deref, { - pub fn query(&self, query: &str, limit: usize) -> Vec { + pub fn query(&self, query: &str, range: Range) -> Vec { 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, F: Fn(DocumentId, &DatabaseView) -> Option, K: Hash + Eq, { - pub fn query(&self, query: &str, limit: usize) -> Vec { + pub fn query(&self, query: &str, range: Range) -> Vec { 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() { - let tmp_groups = mem::replace(&mut groups, Vec::new()); - let mut seen = DistinctMap::new(self.size); + // 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; - '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(), - }; - } + '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); - 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(), + 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 } } }