Merge pull request #32 from Kerollmops/query-range

Improve the query interface
This commit is contained in:
Clément Renault 2018-12-18 16:51:48 +01:00 committed by GitHub
commit 264718c1b9
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 140 additions and 47 deletions

View File

@ -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());

View File

@ -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);
}
}

View File

@ -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 }
}
}