Intern all strings and phrases in the search logic

This commit is contained in:
Loïc Lecrenier
2023-03-06 19:21:55 +01:00
parent 3f1729a17f
commit e8c76cf7bf
19 changed files with 635 additions and 654 deletions

View File

@ -16,30 +16,35 @@ use crate::search::fst_utils::{Complement, Intersection, StartsWith, Union};
use crate::search::{build_dfa, get_first};
use crate::{CboRoaringBitmapLenCodec, Index, Result};
#[derive(Debug, Default, Clone)]
use super::interner::{Interned, Interner};
use super::SearchContext;
#[derive(Default, Clone, PartialEq, Eq, Hash)]
pub struct Phrase {
pub words: Vec<Option<String>>,
pub words: Vec<Option<Interned<String>>>,
}
impl Phrase {
pub fn description(&self) -> String {
self.words.iter().flatten().join(" ")
pub fn description(&self, interner: &Interner<String>) -> String {
self.words.iter().flatten().map(|w| interner.get(*w)).join(" ")
}
}
#[derive(Debug, Clone)]
#[derive(Clone)]
pub struct WordDerivations {
pub original: String,
pub original: Interned<String>,
// TODO: pub prefix_of: Vec<String>,
pub synonyms: Vec<Phrase>,
pub split_words: Option<(String, String)>,
pub zero_typo: Vec<String>,
pub one_typo: Vec<String>,
pub two_typos: Vec<String>,
pub synonyms: Box<[Interned<Phrase>]>,
pub split_words: Option<Interned<Phrase>>,
pub zero_typo: Box<[Interned<String>]>,
pub one_typo: Box<[Interned<String>]>,
pub two_typos: Box<[Interned<String>]>,
pub use_prefix_db: bool,
}
impl WordDerivations {
pub fn all_derivations_except_prefix_db(&self) -> impl Iterator<Item = &String> + Clone {
self.zero_typo.iter().chain(self.one_typo.iter()).chain(self.two_typos.iter())
pub fn all_derivations_except_prefix_db(
&'_ self,
) -> impl Iterator<Item = Interned<String>> + Clone + '_ {
self.zero_typo.iter().chain(self.one_typo.iter()).chain(self.two_typos.iter()).copied()
}
fn is_empty(&self) -> bool {
self.zero_typo.is_empty()
@ -50,15 +55,21 @@ impl WordDerivations {
}
pub fn word_derivations(
index: &Index,
txn: &RoTxn,
ctx: &mut SearchContext,
word: &str,
max_typo: u8,
is_prefix: bool,
fst: &fst::Set<Cow<[u8]>>,
) -> Result<WordDerivations> {
let word_interned = ctx.word_interner.insert(word.to_owned());
let use_prefix_db = is_prefix
&& index.word_prefix_docids.remap_data_type::<DecodeIgnore>().get(txn, word)?.is_some();
&& ctx
.index
.word_prefix_docids
.remap_data_type::<DecodeIgnore>()
.get(ctx.txn, word)?
.is_some();
let mut zero_typo = vec![];
let mut one_typo = vec![];
@ -70,11 +81,12 @@ pub fn word_derivations(
let mut stream = fst.search(prefix).into_stream();
while let Some(word) = stream.next() {
let word = std::str::from_utf8(word)?;
zero_typo.push(word.to_string());
let word = std::str::from_utf8(word)?.to_owned();
let word_interned = ctx.word_interner.insert(word);
zero_typo.push(word_interned);
}
} else if fst.contains(word) {
zero_typo.push(word.to_string());
zero_typo.push(word_interned);
}
} else if max_typo == 1 {
let dfa = build_dfa(word, 1, is_prefix);
@ -83,13 +95,14 @@ pub fn word_derivations(
while let Some((word, state)) = stream.next() {
let word = std::str::from_utf8(word)?;
let word_interned = ctx.word_interner.insert(word.to_owned());
let d = dfa.distance(state.1);
match d.to_u8() {
0 => {
zero_typo.push(word.to_string());
zero_typo.push(word_interned);
}
1 => {
one_typo.push(word.to_string());
one_typo.push(word_interned);
}
_ => panic!(),
}
@ -105,47 +118,56 @@ pub fn word_derivations(
while let Some((found_word, state)) = stream.next() {
let found_word = std::str::from_utf8(found_word)?;
let found_word_interned = ctx.word_interner.insert(found_word.to_owned());
// in the case the typo is on the first letter, we know the number of typo
// is two
if get_first(found_word) != get_first(word) {
two_typos.push(found_word.to_string());
two_typos.push(found_word_interned);
} else {
// Else, we know that it is the second dfa that matched and compute the
// correct distance
let d = second_dfa.distance((state.1).0);
match d.to_u8() {
0 => {
zero_typo.push(found_word.to_string());
zero_typo.push(found_word_interned);
}
1 => {
one_typo.push(found_word.to_string());
one_typo.push(found_word_interned);
}
2 => {
two_typos.push(found_word.to_string());
two_typos.push(found_word_interned);
}
_ => panic!(),
}
}
}
}
let split_words = split_best_frequency(index, txn, word)?;
let split_words = split_best_frequency(ctx.index, ctx.txn, word)?.map(|(l, r)| {
ctx.phrase_interner.insert(Phrase {
words: vec![Some(ctx.word_interner.insert(l)), Some(ctx.word_interner.insert(r))],
})
});
let synonyms = ctx.index.synonyms(ctx.txn)?;
let synonyms = index.synonyms(txn)?;
let synonyms = synonyms
.get(&vec![word.to_owned()])
.cloned()
.unwrap_or_default()
.into_iter()
.map(|words| Phrase { words: words.into_iter().map(Some).collect() })
.map(|words| {
let words = words.into_iter().map(|w| Some(ctx.word_interner.insert(w))).collect();
ctx.phrase_interner.insert(Phrase { words })
})
.collect();
Ok(WordDerivations {
original: word.to_owned(),
original: ctx.word_interner.insert(word.to_owned()),
synonyms,
split_words,
zero_typo,
one_typo,
two_typos,
zero_typo: zero_typo.into_boxed_slice(),
one_typo: one_typo.into_boxed_slice(),
two_typos: two_typos.into_boxed_slice(),
use_prefix_db,
})
}
@ -176,33 +198,36 @@ fn split_best_frequency(
Ok(best.map(|(_, left, right)| (left.to_owned(), right.to_owned())))
}
#[derive(Debug, Clone)]
#[derive(Clone)]
pub enum QueryTerm {
// TODO: should there be SplitWord, NGram2, and NGram3 variants?
// NGram2 can have 1 typo and synonyms
// NGram3 cannot have typos but can have synonyms
// SplitWords are a phrase
// Can NGrams be prefixes?
Phrase { phrase: Phrase },
Phrase { phrase: Interned<Phrase> },
Word { derivations: WordDerivations },
}
impl QueryTerm {
pub fn original_single_word(&self) -> Option<&str> {
pub fn original_single_word<'interner>(
&self,
word_interner: &'interner Interner<String>,
) -> Option<&'interner str> {
match self {
QueryTerm::Phrase { phrase: _ } => None,
QueryTerm::Word { derivations } => {
if derivations.is_empty() {
None
} else {
Some(derivations.original.as_str())
Some(word_interner.get(derivations.original))
}
}
}
}
}
#[derive(Debug, Clone)]
#[derive(Clone)]
pub struct LocatedQueryTerm {
pub value: QueryTerm,
pub positions: RangeInclusive<i8>,
@ -217,18 +242,17 @@ impl LocatedQueryTerm {
}
}
pub fn located_query_terms_from_string<'transaction>(
index: &Index,
txn: &'transaction RoTxn,
pub fn located_query_terms_from_string<'search>(
ctx: &mut SearchContext<'search>,
query: NormalizedTokenIter<Vec<u8>>,
words_limit: Option<usize>,
) -> Result<Vec<LocatedQueryTerm>> {
let authorize_typos = index.authorize_typos(txn)?;
let min_len_one_typo = index.min_word_len_one_typo(txn)?;
let min_len_two_typos = index.min_word_len_two_typos(txn)?;
let authorize_typos = ctx.index.authorize_typos(ctx.txn)?;
let min_len_one_typo = ctx.index.min_word_len_one_typo(ctx.txn)?;
let min_len_two_typos = ctx.index.min_word_len_two_typos(ctx.txn)?;
let exact_words = index.exact_words(txn)?;
let fst = index.words_fst(txn)?;
let exact_words = ctx.index.exact_words(ctx.txn)?;
let fst = ctx.index.words_fst(ctx.txn)?;
let nbr_typos = |word: &str| {
if !authorize_typos
@ -243,10 +267,6 @@ pub fn located_query_terms_from_string<'transaction>(
}
};
let derivations = |word: &str, is_prefix: bool| {
word_derivations(index, txn, word, nbr_typos(word), is_prefix, &fst)
};
let mut primitive_query = Vec::new();
let mut phrase = Vec::new();
@ -279,14 +299,17 @@ pub fn located_query_terms_from_string<'transaction>(
if let TokenKind::StopWord = token.kind {
phrase.push(None);
} else {
let word = ctx.word_interner.insert(token.lemma().to_string());
// TODO: in a phrase, check that every word exists
// otherwise return WordDerivations::Empty
phrase.push(Some(token.lemma().to_string()));
phrase.push(Some(word));
}
} else if peekable.peek().is_some() {
match token.kind {
TokenKind::Word => {
let derivations = derivations(token.lemma(), false)?;
let word = token.lemma();
let derivations =
word_derivations(ctx, word, nbr_typos(word), false, &fst)?;
let located_term = LocatedQueryTerm {
value: QueryTerm::Word { derivations },
positions: position..=position,
@ -296,7 +319,8 @@ pub fn located_query_terms_from_string<'transaction>(
TokenKind::StopWord | TokenKind::Separator(_) | TokenKind::Unknown => {}
}
} else {
let derivations = derivations(token.lemma(), true)?;
let word = token.lemma();
let derivations = word_derivations(ctx, word, nbr_typos(word), true, &fst)?;
let located_term = LocatedQueryTerm {
value: QueryTerm::Word { derivations },
positions: position..=position,
@ -323,7 +347,9 @@ pub fn located_query_terms_from_string<'transaction>(
{
let located_query_term = LocatedQueryTerm {
value: QueryTerm::Phrase {
phrase: Phrase { words: mem::take(&mut phrase) },
phrase: ctx
.phrase_interner
.insert(Phrase { words: mem::take(&mut phrase) }),
},
positions: phrase_start..=phrase_end,
};
@ -337,7 +363,9 @@ pub fn located_query_terms_from_string<'transaction>(
// If a quote is never closed, we consider all of the end of the query as a phrase.
if !phrase.is_empty() {
let located_query_term = LocatedQueryTerm {
value: QueryTerm::Phrase { phrase: Phrase { words: mem::take(&mut phrase) } },
value: QueryTerm::Phrase {
phrase: ctx.phrase_interner.insert(Phrase { words: mem::take(&mut phrase) }),
},
positions: phrase_start..=phrase_end,
};
primitive_query.push(located_query_term);
@ -347,35 +375,49 @@ pub fn located_query_terms_from_string<'transaction>(
}
// TODO: return a word derivations instead?
pub fn ngram2(x: &LocatedQueryTerm, y: &LocatedQueryTerm) -> Option<(String, RangeInclusive<i8>)> {
pub fn ngram2(
ctx: &mut SearchContext,
x: &LocatedQueryTerm,
y: &LocatedQueryTerm,
) -> Option<(Interned<String>, RangeInclusive<i8>)> {
if *x.positions.end() != y.positions.start() - 1 {
return None;
}
match (&x.value.original_single_word(), &y.value.original_single_word()) {
match (
&x.value.original_single_word(&ctx.word_interner),
&y.value.original_single_word(&ctx.word_interner),
) {
(Some(w1), Some(w2)) => {
let term = (format!("{w1}{w2}"), *x.positions.start()..=*y.positions.end());
let term = (
ctx.word_interner.insert(format!("{w1}{w2}")),
*x.positions.start()..=*y.positions.end(),
);
Some(term)
}
_ => None,
}
}
pub fn ngram3(
ctx: &mut SearchContext,
x: &LocatedQueryTerm,
y: &LocatedQueryTerm,
z: &LocatedQueryTerm,
) -> Option<(String, RangeInclusive<i8>)> {
) -> Option<(Interned<String>, RangeInclusive<i8>)> {
if *x.positions.end() != y.positions.start() - 1
|| *y.positions.end() != z.positions.start() - 1
{
return None;
}
match (
&x.value.original_single_word(),
&y.value.original_single_word(),
&z.value.original_single_word(),
&x.value.original_single_word(&ctx.word_interner),
&y.value.original_single_word(&ctx.word_interner),
&z.value.original_single_word(&ctx.word_interner),
) {
(Some(w1), Some(w2), Some(w3)) => {
let term = (format!("{w1}{w2}{w3}"), *x.positions.start()..=*z.positions.end());
let term = (
ctx.word_interner.insert(format!("{w1}{w2}{w3}")),
*x.positions.start()..=*z.positions.end(),
);
Some(term)
}
_ => None,