2018-05-27 15:23:43 +02:00
|
|
|
use std::cmp::{self, Ordering};
|
2018-06-25 22:26:49 +02:00
|
|
|
use std::collections::HashMap;
|
|
|
|
use std::{mem, vec, iter};
|
2018-05-27 15:23:43 +02:00
|
|
|
use DocIndexMap;
|
|
|
|
use fst;
|
2018-07-06 20:58:06 +02:00
|
|
|
use levenshtein::Levenshtein;
|
2018-05-27 15:23:43 +02:00
|
|
|
use map::{
|
|
|
|
OpWithStateBuilder, UnionWithState,
|
|
|
|
StreamWithStateBuilder,
|
|
|
|
Values,
|
|
|
|
};
|
|
|
|
use {Match, DocIndex, DocumentId};
|
2018-06-25 22:26:49 +02:00
|
|
|
use group_by::{GroupBy, GroupByMut};
|
2018-05-27 15:23:43 +02:00
|
|
|
|
2018-06-25 22:26:49 +02:00
|
|
|
const MAX_DISTANCE: u32 = 8;
|
2018-05-27 15:23:43 +02:00
|
|
|
|
2018-06-25 22:26:49 +02:00
|
|
|
#[inline]
|
|
|
|
fn match_query_index(a: &Match, b: &Match) -> bool {
|
|
|
|
a.query_index == b.query_index
|
|
|
|
}
|
|
|
|
|
|
|
|
#[derive(Debug, Clone)]
|
2018-05-27 15:23:43 +02:00
|
|
|
pub struct Document {
|
2018-07-06 22:05:51 +02:00
|
|
|
pub document_id: DocumentId,
|
|
|
|
pub matches: Vec<Match>,
|
2018-05-27 15:23:43 +02:00
|
|
|
}
|
|
|
|
|
|
|
|
impl Document {
|
|
|
|
pub fn new(doc: DocumentId, match_: Match) -> Self {
|
|
|
|
Self::from_sorted_matches(doc, vec![match_])
|
|
|
|
}
|
|
|
|
|
|
|
|
pub fn from_sorted_matches(doc: DocumentId, matches: Vec<Match>) -> Self {
|
|
|
|
Self {
|
|
|
|
document_id: doc,
|
|
|
|
matches: matches,
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2018-06-25 22:26:49 +02:00
|
|
|
fn sum_of_typos(lhs: &Document, rhs: &Document) -> Ordering {
|
2018-07-07 14:55:40 +02:00
|
|
|
let key = |matches: &[Match]| -> u8 {
|
|
|
|
GroupBy::new(matches, match_query_index).map(|m| m[0].distance).sum()
|
2018-07-05 20:22:08 +02:00
|
|
|
};
|
|
|
|
|
2018-07-07 14:55:40 +02:00
|
|
|
key(&lhs.matches).cmp(&key(&rhs.matches))
|
2018-05-27 15:23:43 +02:00
|
|
|
}
|
|
|
|
|
2018-06-25 22:26:49 +02:00
|
|
|
fn number_of_words(lhs: &Document, rhs: &Document) -> Ordering {
|
2018-07-07 14:55:40 +02:00
|
|
|
let key = |matches: &[Match]| -> usize {
|
|
|
|
GroupBy::new(matches, match_query_index).count()
|
2018-07-05 20:22:08 +02:00
|
|
|
};
|
|
|
|
|
2018-07-07 14:55:40 +02:00
|
|
|
key(&lhs.matches).cmp(&key(&rhs.matches)).reverse()
|
2018-05-27 15:23:43 +02:00
|
|
|
}
|
|
|
|
|
2018-07-05 20:22:08 +02:00
|
|
|
fn index_proximity(lhs: u32, rhs: u32) -> u32 {
|
|
|
|
if lhs < rhs {
|
|
|
|
cmp::min(rhs - lhs, MAX_DISTANCE)
|
|
|
|
} else {
|
|
|
|
cmp::min(lhs - rhs, MAX_DISTANCE) + 1
|
|
|
|
}
|
|
|
|
}
|
2018-05-27 15:23:43 +02:00
|
|
|
|
2018-07-05 20:22:08 +02:00
|
|
|
fn attribute_proximity(lhs: &Match, rhs: &Match) -> u32 {
|
|
|
|
if lhs.attribute != rhs.attribute { return MAX_DISTANCE }
|
|
|
|
index_proximity(lhs.attribute_index, rhs.attribute_index)
|
|
|
|
}
|
2018-05-27 15:23:43 +02:00
|
|
|
|
2018-07-07 14:55:40 +02:00
|
|
|
fn min_proximity(lhs: &[Match], rhs: &[Match]) -> u32 {
|
|
|
|
let mut min_prox = u32::max_value();
|
|
|
|
for a in lhs {
|
|
|
|
for b in rhs {
|
|
|
|
min_prox = cmp::min(min_prox, attribute_proximity(a, b));
|
2018-05-27 15:23:43 +02:00
|
|
|
}
|
2018-07-07 14:55:40 +02:00
|
|
|
}
|
|
|
|
min_prox
|
|
|
|
}
|
|
|
|
|
|
|
|
fn matches_proximity(matches: &[Match]) -> u32 {
|
|
|
|
let mut proximity = 0;
|
|
|
|
let mut iter = GroupBy::new(matches, match_query_index);
|
|
|
|
|
|
|
|
let mut last = iter.next();
|
|
|
|
while let (Some(lhs), Some(rhs)) = (last, iter.next()) {
|
|
|
|
proximity += min_proximity(lhs, rhs);
|
|
|
|
last = Some(rhs);
|
|
|
|
}
|
|
|
|
|
|
|
|
proximity
|
|
|
|
}
|
|
|
|
|
|
|
|
fn words_proximity(lhs: &Document, rhs: &Document) -> Ordering {
|
|
|
|
matches_proximity(&lhs.matches).cmp(&matches_proximity(&rhs.matches))
|
|
|
|
}
|
|
|
|
|
|
|
|
#[test]
|
|
|
|
fn easy_matches_proximity() {
|
|
|
|
|
|
|
|
// "soup" "of the" "the day"
|
|
|
|
//
|
|
|
|
// { id: 0, attr: 0, attr_index: 0 }
|
|
|
|
// { id: 1, attr: 1, attr_index: 0 }
|
|
|
|
// { id: 2, attr: 1, attr_index: 1 }
|
|
|
|
// { id: 2, attr: 2, attr_index: 0 }
|
|
|
|
// { id: 3, attr: 3, attr_index: 1 }
|
|
|
|
|
|
|
|
let matches = &[
|
|
|
|
Match { query_index: 0, attribute: 0, attribute_index: 0, ..Match::zero() },
|
|
|
|
Match { query_index: 1, attribute: 1, attribute_index: 0, ..Match::zero() },
|
|
|
|
Match { query_index: 2, attribute: 1, attribute_index: 1, ..Match::zero() },
|
|
|
|
Match { query_index: 2, attribute: 2, attribute_index: 0, ..Match::zero() },
|
|
|
|
Match { query_index: 3, attribute: 3, attribute_index: 1, ..Match::zero() },
|
|
|
|
];
|
|
|
|
|
|
|
|
// soup -> of = 8
|
|
|
|
// + of -> the = 1
|
|
|
|
// + the -> day = 8 (not 1)
|
|
|
|
assert_eq!(matches_proximity(matches), 17);
|
|
|
|
}
|
2018-07-05 20:22:08 +02:00
|
|
|
|
2018-07-07 14:55:40 +02:00
|
|
|
#[test]
|
|
|
|
fn another_matches_proximity() {
|
|
|
|
|
|
|
|
// "soup day" "soup of the day"
|
|
|
|
//
|
|
|
|
// { id: 0, attr: 0, attr_index: 0 }
|
|
|
|
// { id: 0, attr: 1, attr_index: 0 }
|
|
|
|
// { id: 1, attr: 1, attr_index: 1 }
|
|
|
|
// { id: 2, attr: 1, attr_index: 2 }
|
|
|
|
// { id: 3, attr: 0, attr_index: 1 }
|
|
|
|
// { id: 3, attr: 1, attr_index: 3 }
|
|
|
|
|
|
|
|
let matches = &[
|
|
|
|
Match { query_index: 0, attribute: 0, attribute_index: 0, ..Match::zero() },
|
|
|
|
Match { query_index: 0, attribute: 1, attribute_index: 0, ..Match::zero() },
|
|
|
|
Match { query_index: 1, attribute: 1, attribute_index: 1, ..Match::zero() },
|
|
|
|
Match { query_index: 2, attribute: 1, attribute_index: 2, ..Match::zero() },
|
|
|
|
Match { query_index: 3, attribute: 0, attribute_index: 1, ..Match::zero() },
|
|
|
|
Match { query_index: 3, attribute: 1, attribute_index: 3, ..Match::zero() },
|
|
|
|
];
|
|
|
|
|
|
|
|
// soup -> of = 1
|
|
|
|
// + of -> the = 1
|
|
|
|
// + the -> day = 1
|
|
|
|
assert_eq!(matches_proximity(matches), 3);
|
2018-05-27 15:23:43 +02:00
|
|
|
}
|
|
|
|
|
2018-06-25 22:26:49 +02:00
|
|
|
fn sum_of_words_attribute(lhs: &Document, rhs: &Document) -> Ordering {
|
2018-07-07 14:55:40 +02:00
|
|
|
let key = |matches: &[Match]| -> u8 {
|
|
|
|
GroupBy::new(matches, match_query_index).map(|m| m[0].attribute).sum()
|
2018-07-05 20:22:08 +02:00
|
|
|
};
|
|
|
|
|
2018-07-07 14:55:40 +02:00
|
|
|
key(&lhs.matches).cmp(&key(&rhs.matches))
|
2018-05-27 15:23:43 +02:00
|
|
|
}
|
|
|
|
|
2018-06-25 22:26:49 +02:00
|
|
|
fn sum_of_words_position(lhs: &Document, rhs: &Document) -> Ordering {
|
2018-07-07 14:55:40 +02:00
|
|
|
let key = |matches: &[Match]| -> u32 {
|
|
|
|
GroupBy::new(matches, match_query_index).map(|m| m[0].attribute_index).sum()
|
2018-07-05 20:22:08 +02:00
|
|
|
};
|
|
|
|
|
2018-07-07 14:55:40 +02:00
|
|
|
key(&lhs.matches).cmp(&key(&rhs.matches))
|
2018-05-27 15:23:43 +02:00
|
|
|
}
|
|
|
|
|
2018-06-25 22:26:49 +02:00
|
|
|
fn exact(lhs: &Document, rhs: &Document) -> Ordering {
|
2018-07-06 20:58:06 +02:00
|
|
|
let contains_exact = |matches: &[Match]| matches.iter().any(|m| m.is_exact);
|
2018-07-07 14:55:40 +02:00
|
|
|
let key = |matches: &[Match]| -> usize {
|
|
|
|
GroupBy::new(matches, match_query_index).map(contains_exact).filter(Clone::clone).count()
|
2018-07-06 20:58:06 +02:00
|
|
|
};
|
|
|
|
|
2018-07-07 14:55:40 +02:00
|
|
|
key(&lhs.matches).cmp(&key(&rhs.matches))
|
2018-05-27 15:23:43 +02:00
|
|
|
}
|
|
|
|
|
2018-06-25 22:26:49 +02:00
|
|
|
pub struct Pool {
|
|
|
|
documents: Vec<Document>,
|
|
|
|
limit: usize,
|
2018-05-27 15:23:43 +02:00
|
|
|
}
|
|
|
|
|
|
|
|
impl Pool {
|
2018-06-24 15:36:24 +02:00
|
|
|
pub fn new(query_size: usize, limit: usize) -> Self {
|
2018-05-27 15:23:43 +02:00
|
|
|
Self {
|
|
|
|
documents: Vec::new(),
|
2018-06-25 22:26:49 +02:00
|
|
|
limit: limit,
|
2018-05-27 15:23:43 +02:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2018-06-25 22:26:49 +02:00
|
|
|
// TODO remove the matches HashMap, not proud of it
|
|
|
|
pub fn extend(&mut self, matches: &mut HashMap<DocumentId, Vec<Match>>) {
|
2018-05-27 15:23:43 +02:00
|
|
|
for doc in self.documents.iter_mut() {
|
|
|
|
if let Some(matches) = matches.remove(&doc.document_id) {
|
|
|
|
doc.matches.extend(matches);
|
|
|
|
doc.matches.sort_unstable();
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2018-06-25 22:26:49 +02:00
|
|
|
for (id, mut matches) in matches.drain() {
|
|
|
|
// note that matches are already sorted we do that by security
|
|
|
|
// TODO remove this useless sort
|
2018-05-27 15:23:43 +02:00
|
|
|
matches.sort_unstable();
|
|
|
|
|
|
|
|
let document = Document::from_sorted_matches(id, matches);
|
|
|
|
self.documents.push(document);
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2018-06-25 22:26:49 +02:00
|
|
|
fn invert_sorts<F>(a: &Document, b: &Document, sorts: &[F]) -> bool
|
|
|
|
where F: Fn(&Document, &Document) -> Ordering,
|
|
|
|
{
|
|
|
|
sorts.iter().rev().all(|sort| sort(a, b) == Ordering::Equal)
|
|
|
|
}
|
|
|
|
|
2018-05-27 15:23:43 +02:00
|
|
|
impl IntoIterator for Pool {
|
|
|
|
type Item = Document;
|
|
|
|
type IntoIter = vec::IntoIter<Self::Item>;
|
|
|
|
|
|
|
|
fn into_iter(mut self) -> Self::IntoIter {
|
2018-06-25 22:26:49 +02:00
|
|
|
let sorts = &[
|
|
|
|
sum_of_typos,
|
|
|
|
number_of_words,
|
|
|
|
words_proximity,
|
|
|
|
sum_of_words_attribute,
|
|
|
|
sum_of_words_position,
|
2018-07-06 20:58:06 +02:00
|
|
|
exact,
|
2018-06-25 22:26:49 +02:00
|
|
|
];
|
|
|
|
|
|
|
|
for (i, sort) in sorts.iter().enumerate() {
|
|
|
|
let mut computed = 0;
|
|
|
|
for group in GroupByMut::new(&mut self.documents, |a, b| invert_sorts(a, b, &sorts[..i])) {
|
|
|
|
// TODO prefer using `sort_unstable_by_key` to allow reusing the key computation
|
|
|
|
// `number of words` needs to be reversed, we can use the `cmp::Reverse` struct to do that
|
|
|
|
group.sort_unstable_by(sort);
|
|
|
|
computed += group.len();
|
|
|
|
if computed >= self.limit { break }
|
|
|
|
}
|
|
|
|
}
|
2018-06-24 15:36:24 +02:00
|
|
|
|
2018-06-25 22:26:49 +02:00
|
|
|
self.documents.truncate(self.limit);
|
2018-06-24 15:36:24 +02:00
|
|
|
self.documents.into_iter()
|
2018-05-27 15:23:43 +02:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
pub enum RankedStream<'m, 'v> {
|
|
|
|
Fed {
|
|
|
|
inner: UnionWithState<'m, 'v, DocIndex, u32>,
|
2018-07-06 20:58:06 +02:00
|
|
|
automatons: Vec<Levenshtein>,
|
2018-05-27 15:23:43 +02:00
|
|
|
pool: Pool,
|
|
|
|
},
|
|
|
|
Pours {
|
|
|
|
inner: vec::IntoIter<Document>,
|
|
|
|
},
|
|
|
|
}
|
|
|
|
|
|
|
|
impl<'m, 'v> RankedStream<'m, 'v> {
|
2018-07-06 20:58:06 +02:00
|
|
|
pub fn new(map: &'m DocIndexMap, values: &'v Values<DocIndex>, automatons: Vec<Levenshtein>, limit: usize) -> Self {
|
2018-05-27 15:23:43 +02:00
|
|
|
let mut op = OpWithStateBuilder::new(values);
|
|
|
|
|
2018-07-06 20:58:06 +02:00
|
|
|
for automaton in automatons.iter().map(|l| l.dfa.clone()) {
|
2018-05-27 15:23:43 +02:00
|
|
|
let stream = map.as_map().search(automaton).with_state();
|
|
|
|
op.push(stream);
|
|
|
|
}
|
|
|
|
|
2018-06-24 15:36:24 +02:00
|
|
|
let pool = Pool::new(automatons.len(), limit);
|
2018-05-27 15:23:43 +02:00
|
|
|
|
|
|
|
RankedStream::Fed {
|
|
|
|
inner: op.union(),
|
|
|
|
automatons: automatons,
|
|
|
|
pool: pool,
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
impl<'m, 'v, 'a> fst::Streamer<'a> for RankedStream<'m, 'v> {
|
2018-07-06 22:05:51 +02:00
|
|
|
type Item = Document;
|
2018-05-27 15:23:43 +02:00
|
|
|
|
|
|
|
fn next(&'a mut self) -> Option<Self::Item> {
|
2018-06-25 22:26:49 +02:00
|
|
|
let mut matches = HashMap::new();
|
|
|
|
|
2018-05-27 15:23:43 +02:00
|
|
|
loop {
|
|
|
|
// TODO remove that when NLL are here !
|
|
|
|
let mut transfert_pool = None;
|
|
|
|
|
|
|
|
match self {
|
|
|
|
RankedStream::Fed { inner, automatons, pool } => {
|
|
|
|
match inner.next() {
|
2018-07-06 20:58:06 +02:00
|
|
|
Some((string, indexed_values)) => {
|
2018-05-27 15:23:43 +02:00
|
|
|
for iv in indexed_values {
|
|
|
|
|
2018-06-25 22:26:49 +02:00
|
|
|
// TODO extend documents matches by batch of query_index
|
|
|
|
// that way it will be possible to discard matches that
|
|
|
|
// have an invalid distance *before* adding them
|
|
|
|
// to the matches of the documents and, that way, avoid a sort
|
|
|
|
|
2018-07-06 20:58:06 +02:00
|
|
|
let automaton = &automatons[iv.index];
|
|
|
|
let distance = automaton.dfa.distance(iv.state).to_u8();
|
2018-05-27 15:23:43 +02:00
|
|
|
|
|
|
|
// TODO remove the Pool system !
|
|
|
|
// this is an internal Pool rule but
|
|
|
|
// it is more efficient to test that here
|
2018-06-25 22:26:49 +02:00
|
|
|
// if pool.limitation.is_reached() && distance != 0 { continue }
|
2018-05-27 15:23:43 +02:00
|
|
|
|
|
|
|
for di in iv.values {
|
|
|
|
let match_ = Match {
|
|
|
|
query_index: iv.index as u32,
|
|
|
|
distance: distance,
|
|
|
|
attribute: di.attribute,
|
|
|
|
attribute_index: di.attribute_index,
|
2018-07-06 20:58:06 +02:00
|
|
|
is_exact: string.len() == automaton.query_len,
|
2018-05-27 15:23:43 +02:00
|
|
|
};
|
|
|
|
matches.entry(di.document)
|
2018-06-25 22:26:49 +02:00
|
|
|
.and_modify(|ms: &mut Vec<_>| ms.push(match_))
|
2018-05-27 15:23:43 +02:00
|
|
|
.or_insert_with(|| vec![match_]);
|
|
|
|
}
|
2018-06-25 22:26:49 +02:00
|
|
|
pool.extend(&mut matches);
|
2018-05-27 15:23:43 +02:00
|
|
|
}
|
|
|
|
},
|
|
|
|
None => {
|
2018-07-06 20:58:06 +02:00
|
|
|
// TODO remove this when NLL are here !
|
2018-06-24 15:36:24 +02:00
|
|
|
transfert_pool = Some(mem::replace(pool, Pool::new(1, 1)));
|
2018-05-27 15:23:43 +02:00
|
|
|
},
|
|
|
|
}
|
|
|
|
},
|
|
|
|
RankedStream::Pours { inner } => {
|
2018-07-06 22:05:51 +02:00
|
|
|
return inner.next()
|
2018-05-27 15:23:43 +02:00
|
|
|
},
|
|
|
|
}
|
|
|
|
|
|
|
|
// transform the `RankedStream` into a `Pours`
|
|
|
|
if let Some(pool) = transfert_pool {
|
|
|
|
*self = RankedStream::Pours {
|
|
|
|
inner: pool.into_iter(),
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|