feat: Introduce the Criterion trait

This commit is contained in:
Clément Renault 2018-10-11 14:04:41 +02:00
parent c56c35b45b
commit 8cd07462aa
11 changed files with 209 additions and 218 deletions

View File

@ -3,11 +3,10 @@ use std::io::{self, Write};
use structopt::StructOpt; use structopt::StructOpt;
use std::path::PathBuf; use std::path::PathBuf;
use fst::Streamer;
use elapsed::measure_time; use elapsed::measure_time;
use rocksdb::{DB, DBOptions, IngestExternalFileOptions}; use rocksdb::{DB, DBOptions, IngestExternalFileOptions};
use raptor::{automaton, Metadata, CommonWords}; use raptor::{automaton, Metadata, CommonWords};
use raptor::rank; use raptor::rank::{criterion, RankedStreamBuilder};
#[derive(Debug, StructOpt)] #[derive(Debug, StructOpt)]
pub struct CommandConsole { pub struct CommandConsole {
@ -70,15 +69,13 @@ fn search(metadata: &Metadata, database: &DB, common_words: &CommonWords, query:
automatons.push(lev); automatons.push(lev);
} }
let config = rank::Config { let mut builder = RankedStreamBuilder::new(metadata, automatons);
criteria: rank::criterion::default(), builder.criteria(criterion::default());
metadata: &metadata,
automatons: automatons,
limit: 20,
};
let mut stream = rank::RankedStream::new(config); let mut stream = builder.build();
while let Some(document) = stream.next() { let documents = stream.retrieve_documents(20);
for document in documents {
let id_key = format!("{}-id", document.id); let id_key = format!("{}-id", document.id);
let id = database.get(id_key.as_bytes()).unwrap().unwrap(); let id = database.get(id_key.as_bytes()).unwrap().unwrap();
let id = unsafe { from_utf8_unchecked(&id) }; let id = unsafe { from_utf8_unchecked(&id) };

View File

@ -7,10 +7,9 @@ use std::path::PathBuf;
use std::error::Error; use std::error::Error;
use std::sync::Arc; use std::sync::Arc;
use raptor::rank; use raptor::rank::{criterion, RankedStreamBuilder};
use raptor::{automaton, Metadata, CommonWords}; use raptor::{automaton, Metadata, CommonWords};
use rocksdb::{DB, DBOptions, IngestExternalFileOptions}; use rocksdb::{DB, DBOptions, IngestExternalFileOptions};
use fst::Streamer;
use warp::Filter; use warp::Filter;
use structopt::StructOpt; use structopt::StructOpt;
@ -100,19 +99,17 @@ where M: AsRef<Metadata>,
automatons.push(lev); automatons.push(lev);
} }
let config = rank::Config { let mut builder = RankedStreamBuilder::new(metadata.as_ref(), automatons);
criteria: rank::criterion::default(), builder.criteria(criterion::default());
metadata: metadata.as_ref(),
automatons: automatons, let mut stream = builder.build();
limit: 20, let documents = stream.retrieve_documents(20);
};
let mut stream = rank::RankedStream::new(config);
let mut body = Vec::new(); let mut body = Vec::new();
write!(&mut body, "[")?; write!(&mut body, "[")?;
let mut first = true; let mut first = true;
while let Some(document) = stream.next() { for document in documents {
let title_key = format!("{}-title", document.id); let title_key = format!("{}-title", document.id);
let title = database.as_ref().get(title_key.as_bytes()).unwrap().unwrap(); let title = database.as_ref().get(title_key.as_bytes()).unwrap().unwrap();
let title = unsafe { from_utf8_unchecked(&title) }; let title = unsafe { from_utf8_unchecked(&title) };

View File

@ -2,6 +2,7 @@ use std::cmp::Ordering;
use group_by::GroupBy; use group_by::GroupBy;
use crate::Match; use crate::Match;
use crate::rank::{match_query_index, Document}; use crate::rank::{match_query_index, Document};
use crate::rank::criterion::Criterion;
#[inline] #[inline]
fn contains_exact(matches: &[Match]) -> bool { fn contains_exact(matches: &[Match]) -> bool {
@ -13,10 +14,14 @@ fn number_exact_matches(matches: &[Match]) -> usize {
GroupBy::new(matches, match_query_index).map(contains_exact).count() GroupBy::new(matches, match_query_index).map(contains_exact).count()
} }
#[inline] #[derive(Debug, Clone, Copy)]
pub fn exact(lhs: &Document, rhs: &Document) -> Ordering { pub struct Exact;
impl Criterion for Exact {
fn evaluate(&self, lhs: &Document, rhs: &Document) -> Ordering {
let lhs = number_exact_matches(&lhs.matches); let lhs = number_exact_matches(&lhs.matches);
let rhs = number_exact_matches(&rhs.matches); let rhs = number_exact_matches(&rhs.matches);
lhs.cmp(&rhs).reverse() lhs.cmp(&rhs).reverse()
}
} }

View File

@ -7,65 +7,64 @@ mod exact;
use std::vec; use std::vec;
use std::cmp::Ordering; use std::cmp::Ordering;
use std::ops::Deref;
use crate::rank::Document; use crate::rank::Document;
pub use self::{ pub use self::{
sum_of_typos::sum_of_typos, sum_of_typos::SumOfTypos,
number_of_words::number_of_words, number_of_words::NumberOfWords,
words_proximity::words_proximity, words_proximity::WordsProximity,
sum_of_words_attribute::sum_of_words_attribute, sum_of_words_attribute::SumOfWordsAttribute,
sum_of_words_position::sum_of_words_position, sum_of_words_position::SumOfWordsPosition,
exact::exact, exact::Exact,
}; };
#[inline] pub trait Criterion {
pub fn document_id(lhs: &Document, rhs: &Document) -> Ordering { #[inline]
fn evaluate(&self, lhs: &Document, rhs: &Document) -> Ordering;
#[inline]
fn eq(&self, lhs: &Document, rhs: &Document) -> bool {
self.evaluate(lhs, rhs) == Ordering::Equal
}
}
impl<'a, T: Criterion + ?Sized> Criterion for &'a T {
fn evaluate(&self, lhs: &Document, rhs: &Document) -> Ordering {
self.deref().evaluate(lhs, rhs)
}
fn eq(&self, lhs: &Document, rhs: &Document) -> bool {
self.deref().eq(lhs, rhs)
}
}
impl<T: Criterion + ?Sized> Criterion for Box<T> {
fn evaluate(&self, lhs: &Document, rhs: &Document) -> Ordering {
self.deref().evaluate(lhs, rhs)
}
fn eq(&self, lhs: &Document, rhs: &Document) -> bool {
self.deref().eq(lhs, rhs)
}
}
#[derive(Debug, Clone, Copy)]
pub struct DocumentId;
impl Criterion for DocumentId {
fn evaluate(&self, lhs: &Document, rhs: &Document) -> Ordering {
lhs.id.cmp(&rhs.id) lhs.id.cmp(&rhs.id)
}
#[derive(Debug)]
pub struct Criteria<F>(Vec<F>);
impl<F> Criteria<F> {
pub fn new() -> Self {
Criteria(Vec::new())
}
pub fn with_capacity(cap: usize) -> Self {
Criteria(Vec::with_capacity(cap))
}
pub fn push(&mut self, criterion: F) {
self.0.push(criterion)
}
pub fn add(mut self, criterion: F) -> Self {
self.push(criterion);
self
} }
} }
impl<F> IntoIterator for Criteria<F> { pub fn default() -> Vec<Box<dyn Criterion>> {
type Item = F; vec![
type IntoIter = vec::IntoIter<Self::Item>; Box::new(SumOfTypos),
Box::new(NumberOfWords),
fn into_iter(self) -> Self::IntoIter { Box::new(WordsProximity),
self.0.into_iter() Box::new(SumOfWordsAttribute),
} Box::new(SumOfWordsPosition),
} Box::new(Exact),
]
pub fn default() -> Criteria<impl Fn(&Document, &Document) -> Ordering + Copy> {
let functions = &[
sum_of_typos,
number_of_words,
words_proximity,
sum_of_words_attribute,
sum_of_words_position,
exact,
document_id,
];
let mut criteria = Criteria::with_capacity(functions.len());
for f in functions { criteria.push(f) }
criteria
} }

View File

@ -2,16 +2,21 @@ use std::cmp::Ordering;
use group_by::GroupBy; use group_by::GroupBy;
use crate::Match; use crate::Match;
use crate::rank::{match_query_index, Document}; use crate::rank::{match_query_index, Document};
use crate::rank::criterion::Criterion;
#[inline] #[inline]
fn number_of_query_words(matches: &[Match]) -> usize { fn number_of_query_words(matches: &[Match]) -> usize {
GroupBy::new(matches, match_query_index).count() GroupBy::new(matches, match_query_index).count()
} }
#[inline] #[derive(Debug, Clone, Copy)]
pub fn number_of_words(lhs: &Document, rhs: &Document) -> Ordering { pub struct NumberOfWords;
impl Criterion for NumberOfWords {
fn evaluate(&self, lhs: &Document, rhs: &Document) -> Ordering {
let lhs = number_of_query_words(&lhs.matches); let lhs = number_of_query_words(&lhs.matches);
let rhs = number_of_query_words(&rhs.matches); let rhs = number_of_query_words(&rhs.matches);
lhs.cmp(&rhs).reverse() lhs.cmp(&rhs).reverse()
}
} }

View File

@ -2,6 +2,7 @@ use std::cmp::Ordering;
use group_by::GroupBy; use group_by::GroupBy;
use crate::Match; use crate::Match;
use crate::rank::{match_query_index, Document}; use crate::rank::{match_query_index, Document};
use crate::rank::criterion::Criterion;
#[inline] #[inline]
fn sum_matches_typos(matches: &[Match]) -> i8 { fn sum_matches_typos(matches: &[Match]) -> i8 {
@ -18,14 +19,19 @@ fn sum_matches_typos(matches: &[Match]) -> i8 {
sum_typos - number_words sum_typos - number_words
} }
#[inline] #[derive(Debug, Clone, Copy)]
pub fn sum_of_typos(lhs: &Document, rhs: &Document) -> Ordering { pub struct SumOfTypos;
impl Criterion for SumOfTypos {
fn evaluate(&self, lhs: &Document, rhs: &Document) -> Ordering {
let lhs = sum_matches_typos(&lhs.matches); let lhs = sum_matches_typos(&lhs.matches);
let rhs = sum_matches_typos(&rhs.matches); let rhs = sum_matches_typos(&rhs.matches);
lhs.cmp(&rhs) lhs.cmp(&rhs)
}
} }
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
@ -42,7 +48,7 @@ mod tests {
Match { query_index: 1, distance: 0, attribute: 0, attribute_index: 2, is_exact: false }, Match { query_index: 1, distance: 0, attribute: 0, attribute_index: 2, is_exact: false },
]; ];
Document { Document {
document_id: 0, id: 0,
matches: matches, matches: matches,
} }
}; };
@ -53,12 +59,12 @@ mod tests {
Match { query_index: 1, distance: 0, attribute: 0, attribute_index: 2, is_exact: false }, Match { query_index: 1, distance: 0, attribute: 0, attribute_index: 2, is_exact: false },
]; ];
Document { Document {
document_id: 1, id: 1,
matches: matches, matches: matches,
} }
}; };
assert_eq!(sum_of_typos(&doc0, &doc1), Ordering::Less); assert_eq!(SumOfTypos.evaluate(&doc0, &doc1), Ordering::Less);
} }
// typing: "bouton manchette" // typing: "bouton manchette"
@ -73,7 +79,7 @@ mod tests {
Match { query_index: 1, distance: 0, attribute: 0, attribute_index: 1, is_exact: false }, Match { query_index: 1, distance: 0, attribute: 0, attribute_index: 1, is_exact: false },
]; ];
Document { Document {
document_id: 0, id: 0,
matches: matches, matches: matches,
} }
}; };
@ -83,12 +89,12 @@ mod tests {
Match { query_index: 0, distance: 0, attribute: 0, attribute_index: 0, is_exact: false }, Match { query_index: 0, distance: 0, attribute: 0, attribute_index: 0, is_exact: false },
]; ];
Document { Document {
document_id: 1, id: 1,
matches: matches, matches: matches,
} }
}; };
assert_eq!(sum_of_typos(&doc0, &doc1), Ordering::Less); assert_eq!(SumOfTypos.evaluate(&doc0, &doc1), Ordering::Less);
} }
// typing: "bouton manchztte" // typing: "bouton manchztte"
@ -103,7 +109,7 @@ mod tests {
Match { query_index: 1, distance: 1, attribute: 0, attribute_index: 1, is_exact: false }, Match { query_index: 1, distance: 1, attribute: 0, attribute_index: 1, is_exact: false },
]; ];
Document { Document {
document_id: 0, id: 0,
matches: matches, matches: matches,
} }
}; };
@ -113,11 +119,11 @@ mod tests {
Match { query_index: 0, distance: 0, attribute: 0, attribute_index: 0, is_exact: false }, Match { query_index: 0, distance: 0, attribute: 0, attribute_index: 0, is_exact: false },
]; ];
Document { Document {
document_id: 1, id: 1,
matches: matches, matches: matches,
} }
}; };
assert_eq!(sum_of_typos(&doc0, &doc1), Ordering::Equal); assert_eq!(SumOfTypos.evaluate(&doc0, &doc1), Ordering::Equal);
} }
} }

View File

@ -2,6 +2,7 @@ use std::cmp::Ordering;
use group_by::GroupBy; use group_by::GroupBy;
use crate::Match; use crate::Match;
use crate::rank::{match_query_index, Document}; use crate::rank::{match_query_index, Document};
use crate::rank::criterion::Criterion;
#[inline] #[inline]
fn sum_matches_attributes(matches: &[Match]) -> u8 { fn sum_matches_attributes(matches: &[Match]) -> u8 {
@ -12,10 +13,14 @@ fn sum_matches_attributes(matches: &[Match]) -> u8 {
}).sum() }).sum()
} }
#[inline] #[derive(Debug, Clone, Copy)]
pub fn sum_of_words_attribute(lhs: &Document, rhs: &Document) -> Ordering { pub struct SumOfWordsAttribute;
impl Criterion for SumOfWordsAttribute {
fn evaluate(&self, lhs: &Document, rhs: &Document) -> Ordering {
let lhs = sum_matches_attributes(&lhs.matches); let lhs = sum_matches_attributes(&lhs.matches);
let rhs = sum_matches_attributes(&rhs.matches); let rhs = sum_matches_attributes(&rhs.matches);
lhs.cmp(&rhs) lhs.cmp(&rhs)
}
} }

View File

@ -2,6 +2,7 @@ use std::cmp::Ordering;
use group_by::GroupBy; use group_by::GroupBy;
use crate::Match; use crate::Match;
use crate::rank::{match_query_index, Document}; use crate::rank::{match_query_index, Document};
use crate::rank::criterion::Criterion;
#[inline] #[inline]
fn sum_matches_attribute_index(matches: &[Match]) -> u32 { fn sum_matches_attribute_index(matches: &[Match]) -> u32 {
@ -12,10 +13,14 @@ fn sum_matches_attribute_index(matches: &[Match]) -> u32 {
}).sum() }).sum()
} }
#[inline] #[derive(Debug, Clone, Copy)]
pub fn sum_of_words_position(lhs: &Document, rhs: &Document) -> Ordering { pub struct SumOfWordsPosition;
impl Criterion for SumOfWordsPosition {
fn evaluate(&self, lhs: &Document, rhs: &Document) -> Ordering {
let lhs = sum_matches_attribute_index(&lhs.matches); let lhs = sum_matches_attribute_index(&lhs.matches);
let rhs = sum_matches_attribute_index(&rhs.matches); let rhs = sum_matches_attribute_index(&rhs.matches);
lhs.cmp(&rhs) lhs.cmp(&rhs)
}
} }

View File

@ -2,6 +2,7 @@ use std::cmp::{self, Ordering};
use group_by::GroupBy; use group_by::GroupBy;
use crate::Match; use crate::Match;
use crate::rank::{match_query_index, Document}; use crate::rank::{match_query_index, Document};
use crate::rank::criterion::Criterion;
const MAX_DISTANCE: u32 = 8; const MAX_DISTANCE: u32 = 8;
@ -42,10 +43,19 @@ fn matches_proximity(matches: &[Match]) -> u32 {
proximity proximity
} }
pub fn words_proximity(lhs: &Document, rhs: &Document) -> Ordering { #[derive(Debug, Clone, Copy)]
matches_proximity(&lhs.matches).cmp(&matches_proximity(&rhs.matches)) pub struct WordsProximity;
impl Criterion for WordsProximity {
fn evaluate(&self, lhs: &Document, rhs: &Document) -> Ordering {
let lhs = matches_proximity(&lhs.matches);
let rhs = matches_proximity(&rhs.matches);
lhs.cmp(&rhs)
}
} }
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;

View File

@ -3,7 +3,7 @@ mod ranked_stream;
use crate::{Match, DocumentId}; use crate::{Match, DocumentId};
pub use self::ranked_stream::{RankedStream, Config}; pub use self::ranked_stream::{RankedStreamBuilder, RankedStream};
#[inline] #[inline]
fn match_query_index(a: &Match, b: &Match) -> bool { fn match_query_index(a: &Match, b: &Match) -> bool {
@ -18,10 +18,10 @@ pub struct Document {
impl Document { impl Document {
pub fn new(doc: DocumentId, match_: Match) -> Self { pub fn new(doc: DocumentId, match_: Match) -> Self {
Self::from_sorted_matches(doc, vec![match_]) unsafe { Self::from_sorted_matches(doc, vec![match_]) }
} }
pub fn from_sorted_matches(id: DocumentId, matches: Vec<Match>) -> Self { pub unsafe fn from_sorted_matches(id: DocumentId, matches: Vec<Match>) -> Self {
Self { id, matches } Self { id, matches }
} }
} }

View File

@ -1,4 +1,3 @@
use std::cmp::Ordering;
use std::rc::Rc; use std::rc::Rc;
use std::{mem, vec}; use std::{mem, vec};
@ -8,76 +7,60 @@ use group_by::GroupByMut;
use crate::automaton::{DfaExt, AutomatonExt}; use crate::automaton::{DfaExt, AutomatonExt};
use crate::metadata::Metadata; use crate::metadata::Metadata;
use crate::metadata::ops::{OpBuilder, Union}; use crate::metadata::ops::OpBuilder;
use crate::rank::criterion::Criteria; use crate::rank::criterion::Criterion;
use crate::rank::Document; use crate::rank::Document;
use crate::{Match, DocumentId}; use crate::Match;
pub struct Config<'m, F> { #[derive(Clone)]
pub criteria: Criteria<F>, pub struct RankedStreamBuilder<'m, C> {
pub metadata: &'m Metadata, metadata: &'m Metadata,
pub automatons: Vec<DfaExt>,
pub limit: usize,
}
pub struct RankedStream<'m, F>(RankedStreamInner<'m, F>);
impl<'m, F> RankedStream<'m, F> {
pub fn new(config: Config<'m, 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);
let inner = RankedStreamInner::Fed {
inner: builder.union(),
automatons: automatons,
criteria: config.criteria,
limit: config.limit,
matches: FnvHashMap::default(),
};
RankedStream(inner)
}
}
impl<'m, 'a, F> fst::Streamer<'a> for RankedStream<'m, F>
where F: Fn(&Document, &Document) -> Ordering + Copy,
{
type Item = Document;
fn next(&'a mut self) -> Option<Self::Item> {
self.0.next()
}
}
enum RankedStreamInner<'m, F> {
Fed {
inner: Union<'m>,
automatons: Vec<Rc<DfaExt>>, automatons: Vec<Rc<DfaExt>>,
criteria: Criteria<F>, criteria: Vec<C>,
limit: usize,
matches: FnvHashMap<DocumentId, Vec<Match>>,
},
Pours {
inner: vec::IntoIter<Document>,
},
} }
impl<'m, 'a, F> fst::Streamer<'a> for RankedStreamInner<'m, F> impl<'m, C> RankedStreamBuilder<'m, C> {
where F: Fn(&Document, &Document) -> Ordering + Copy, pub fn new(metadata: &'m Metadata, automatons: Vec<DfaExt>) -> Self {
{ RankedStreamBuilder {
type Item = Document; metadata: metadata,
automatons: automatons.into_iter().map(Rc::new).collect(),
criteria: Vec::new(), // hummm... prefer the criterion::default() ones !
}
}
fn next(&'a mut self) -> Option<Self::Item> { pub fn criteria(&mut self, criteria: Vec<C>) {
loop { self.criteria = criteria;
match self { }
RankedStreamInner::Fed { inner, automatons, criteria, limit, matches } => {
match inner.next() { pub fn build(&self) -> RankedStream<C> {
Some((string, indexed_values)) => { let mut builder = OpBuilder::with_automatons(self.automatons.clone());
builder.push(self.metadata);
RankedStream {
stream: builder.union(),
automatons: &self.automatons,
criteria: &self.criteria,
}
}
}
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, limit: usize) -> Vec<Document>
where C: Criterion
{
let mut matches = FnvHashMap::default();
while let Some((string, indexed_values)) = self.stream.next() {
for iv in indexed_values { for iv in indexed_values {
let automaton = &automatons[iv.index]; let automaton = &self.automatons[iv.index];
let distance = automaton.eval(string).to_u8(); let distance = automaton.eval(string).to_u8();
let same_length = string.len() == automaton.query_len(); let is_exact = distance == 0 && string.len() == automaton.query_len();
for di in iv.doc_indexes.as_slice() { for di in iv.doc_indexes.as_slice() {
let match_ = Match { let match_ = Match {
@ -85,50 +68,28 @@ where F: Fn(&Document, &Document) -> Ordering + Copy,
distance: distance, distance: distance,
attribute: di.attribute, attribute: di.attribute,
attribute_index: di.attribute_index, attribute_index: di.attribute_index,
is_exact: distance == 0 && same_length, is_exact: is_exact,
}; };
matches.entry(di.document) matches.entry(di.document).or_insert_with(Vec::new).push(match_);
.or_insert_with(Vec::new)
.push(match_);
}
}
},
None => {
let matches = mem::replace(matches, FnvHashMap::default());
let criteria = mem::replace(criteria, Criteria::new());
*self = RankedStreamInner::Pours {
inner: matches_into_iter(matches, criteria, *limit).into_iter()
};
},
}
},
RankedStreamInner::Pours { inner } => {
return inner.next()
},
} }
} }
} }
}
fn matches_into_iter<F>(matches: FnvHashMap<DocumentId, Vec<Match>>, // collect matches from an HashMap into a Vec
criteria: Criteria<F>,
limit: usize) -> vec::IntoIter<Document>
where F: Fn(&Document, &Document) -> Ordering + Copy,
{
let mut documents: Vec<_> = matches.into_iter().map(|(id, mut matches)| { let mut documents: Vec<_> = matches.into_iter().map(|(id, mut matches)| {
matches.sort_unstable(); matches.sort_unstable();
Document::from_sorted_matches(id, matches) unsafe { Document::from_sorted_matches(id, matches) }
}).collect(); }).collect();
let mut groups = vec![documents.as_mut_slice()]; let mut groups = vec![documents.as_mut_slice()];
for sort in criteria { for criterion in self.criteria {
let temp = mem::replace(&mut groups, Vec::new()); let temp = mem::replace(&mut groups, Vec::new());
let mut computed = 0; let mut computed = 0;
'grp: for group in temp { 'grp: for group in temp {
group.sort_unstable_by(sort); group.sort_unstable_by(|a, b| criterion.evaluate(a, b));
for group in GroupByMut::new(group, |a, b| sort(a, b) == Ordering::Equal) { for group in GroupByMut::new(group, |a, b| criterion.eq(a, b)) {
computed += group.len(); computed += group.len();
groups.push(group); groups.push(group);
if computed >= limit { break 'grp } if computed >= limit { break 'grp }
@ -137,5 +98,6 @@ where F: Fn(&Document, &Document) -> Ordering + Copy,
} }
documents.truncate(limit); documents.truncate(limit);
documents.into_iter() documents
}
} }