MeiliSearch/milli/src/search/new/query_term.rs

607 lines
22 KiB
Rust
Raw Normal View History

use std::mem;
use std::ops::RangeInclusive;
use charabia::normalizer::NormalizedTokenIter;
use charabia::{SeparatorKind, TokenKind};
use fst::automaton::Str;
use fst::{Automaton, IntoStreamer, Streamer};
use heed::types::DecodeIgnore;
use heed::RoTxn;
use itertools::Itertools;
2023-03-08 09:55:53 +01:00
use super::interner::{Interned, Interner};
use super::SearchContext;
use crate::search::fst_utils::{Complement, Intersection, StartsWith, Union};
use crate::search::{build_dfa, get_first};
use crate::{CboRoaringBitmapLenCodec, Index, Result, MAX_WORD_LENGTH};
/// A phrase in the user's search query, consisting of several words
/// that must appear side-by-side in the search results.
#[derive(Default, Clone, PartialEq, Eq, Hash)]
pub struct Phrase {
pub words: Vec<Option<Interned<String>>>,
}
impl Phrase {
pub fn description(&self, interner: &Interner<String>) -> String {
self.words.iter().flatten().map(|w| interner.get(*w)).join(" ")
}
}
/// A structure storing all the different ways to match
/// a term in the user's search query.
2023-03-09 11:12:31 +01:00
#[derive(Clone, PartialEq, Eq, Hash)]
pub struct WordDerivations {
/// The original terms, for debugging purposes
pub original: Interned<String>,
pub is_prefix: bool,
/// A single word equivalent to the original one, with zero typos
pub zero_typo: Option<Interned<String>>,
/// All the words that contain the original word as prefix
pub prefix_of: Box<[Interned<String>]>,
/// All the synonyms of the original word
pub synonyms: Box<[Interned<Phrase>]>,
/// The original word split into multiple consecutive words
pub split_words: Option<Interned<Phrase>>,
/// Words that are 1 typo away from the original word
pub one_typo: Box<[Interned<String>]>,
/// Words that are 2 typos away from the original word
pub two_typos: Box<[Interned<String>]>,
/// A prefix in the prefix databases matching the original word
pub use_prefix_db: Option<Interned<String>>,
}
impl WordDerivations {
pub fn empty(word_interner: &mut Interner<String>, original: &str) -> Self {
Self {
original: word_interner.insert(original.to_owned()),
is_prefix: false,
zero_typo: None,
prefix_of: Box::new([]),
synonyms: Box::new([]),
split_words: None,
one_typo: Box::new([]),
two_typos: Box::new([]),
use_prefix_db: None,
}
}
/// Return an iterator over all the single words derived from the original word.
///
/// This excludes synonyms, split words, and words stored in the prefix databases.
2023-03-09 11:12:31 +01:00
pub fn all_single_word_derivations_except_prefix_db(
&'_ self,
) -> impl Iterator<Item = Interned<String>> + Clone + '_ {
self.zero_typo
.iter()
.chain(self.prefix_of.iter())
.chain(self.one_typo.iter())
.chain(self.two_typos.iter())
.copied()
}
/// Return an iterator over all the single words derived from the original word.
///
/// This excludes synonyms, split words, and words stored in the prefix databases.
pub fn all_phrase_derivations(&'_ self) -> impl Iterator<Item = Interned<Phrase>> + Clone + '_ {
self.split_words.iter().chain(self.synonyms.iter()).copied()
}
2023-03-09 11:12:31 +01:00
pub fn is_empty(&self) -> bool {
self.zero_typo.is_none()
&& self.one_typo.is_empty()
&& self.two_typos.is_empty()
&& self.prefix_of.is_empty()
&& self.synonyms.is_empty()
&& self.split_words.is_none()
&& self.use_prefix_db.is_none()
}
}
/// Compute the word derivations for the given word
pub fn word_derivations(
ctx: &mut SearchContext,
word: &str,
max_typo: u8,
is_prefix: bool,
) -> Result<WordDerivations> {
if word.len() > MAX_WORD_LENGTH {
return Ok(WordDerivations::empty(&mut ctx.word_interner, word));
}
let fst = ctx.index.words_fst(ctx.txn)?;
let word_interned = ctx.word_interner.insert(word.to_owned());
let use_prefix_db = is_prefix
&& ctx
.index
.word_prefix_docids
.remap_data_type::<DecodeIgnore>()
.get(ctx.txn, word)?
.is_some();
let use_prefix_db = if use_prefix_db { Some(word_interned) } else { None };
let mut zero_typo = None;
let mut prefix_of = vec![];
let mut one_typo = vec![];
let mut two_typos = vec![];
if fst.contains(word) {
zero_typo = Some(word_interned);
}
if max_typo == 0 {
if is_prefix && use_prefix_db.is_none() {
let prefix = Str::new(word).starts_with();
let mut stream = fst.search(prefix).into_stream();
2023-03-09 11:12:31 +01:00
while let Some(derived_word) = stream.next() {
let derived_word = std::str::from_utf8(derived_word)?.to_owned();
let derived_word_interned = ctx.word_interner.insert(derived_word);
if derived_word_interned != word_interned {
prefix_of.push(derived_word_interned);
}
}
}
} else if max_typo == 1 {
let dfa = build_dfa(word, 1, is_prefix);
let starts = StartsWith(Str::new(get_first(word)));
let mut stream = fst.search_with_state(Intersection(starts, &dfa)).into_stream();
2023-03-09 11:12:31 +01:00
// TODO: There may be wayyy too many matches (e.g. in the thousands), how to reduce them?
while let Some((derived_word, state)) = stream.next() {
let derived_word = std::str::from_utf8(derived_word)?;
let d = dfa.distance(state.1);
2023-03-09 11:12:31 +01:00
let derived_word_interned = ctx.word_interner.insert(derived_word.to_owned());
match d.to_u8() {
0 => {
if derived_word_interned != word_interned {
prefix_of.push(derived_word_interned);
}
}
1 => {
2023-03-09 11:12:31 +01:00
one_typo.push(derived_word_interned);
}
_ => panic!(),
}
}
} else {
let starts = StartsWith(Str::new(get_first(word)));
let first = Intersection(build_dfa(word, 1, is_prefix), Complement(&starts));
let second_dfa = build_dfa(word, 2, is_prefix);
let second = Intersection(&second_dfa, &starts);
let automaton = Union(first, &second);
let mut stream = fst.search_with_state(automaton).into_stream();
2023-03-09 11:12:31 +01:00
// TODO: There may be wayyy too many matches (e.g. in the thousands), how to reduce them?
2023-03-09 11:12:31 +01:00
while let Some((derived_word, state)) = stream.next() {
let derived_word = std::str::from_utf8(derived_word)?;
let derived_word_interned = ctx.word_interner.insert(derived_word.to_owned());
// in the case the typo is on the first letter, we know the number of typo
// is two
2023-03-09 11:12:31 +01:00
if get_first(derived_word) != get_first(word) {
two_typos.push(derived_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 => {
if derived_word_interned != word_interned {
prefix_of.push(derived_word_interned);
}
}
1 => {
2023-03-09 11:12:31 +01:00
one_typo.push(derived_word_interned);
}
2 => {
2023-03-09 11:12:31 +01:00
two_typos.push(derived_word_interned);
}
_ => panic!(),
}
}
}
}
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 = synonyms
.get(&vec![word.to_owned()])
.cloned()
.unwrap_or_default()
.into_iter()
.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_interned,
is_prefix,
zero_typo,
prefix_of: prefix_of.into_boxed_slice(),
synonyms,
split_words,
one_typo: one_typo.into_boxed_slice(),
two_typos: two_typos.into_boxed_slice(),
use_prefix_db,
})
}
/// Split the original word into the two words that appear the
/// most next to each other in the index.
///
/// Return `None` if the original word cannot be split.
fn split_best_frequency(
index: &Index,
txn: &RoTxn,
original: &str,
) -> Result<Option<(String, String)>> {
let chars = original.char_indices().skip(1);
let mut best = None;
for (i, _) in chars {
let (left, right) = original.split_at(i);
let key = (1, left, right);
let frequency = index
.word_pair_proximity_docids
.remap_data_type::<CboRoaringBitmapLenCodec>()
.get(txn, &key)?
.unwrap_or(0);
if frequency != 0 && best.map_or(true, |(old, _, _)| frequency > old) {
best = Some((frequency, left, right));
}
}
Ok(best.map(|(_, left, right)| (left.to_owned(), right.to_owned())))
}
2023-03-09 11:12:31 +01:00
#[derive(Clone, PartialEq, Eq, Hash)]
pub enum QueryTerm {
Phrase { phrase: Interned<Phrase> },
2023-03-09 11:12:31 +01:00
// TODO: change to `Interned<WordDerivations>`?
Word { derivations: Interned<WordDerivations> },
}
impl QueryTerm {
pub fn is_prefix(&self, derivations_interner: &Interner<WordDerivations>) -> bool {
match self {
QueryTerm::Phrase { .. } => false,
QueryTerm::Word { derivations } => derivations_interner.get(*derivations).is_prefix,
}
}
/// Return the original word from the given query term
pub fn original_single_word(
&self,
derivations_interner: &Interner<WordDerivations>,
) -> Option<Interned<String>> {
match self {
QueryTerm::Phrase { phrase: _ } => None,
QueryTerm::Word { derivations } => {
2023-03-09 11:12:31 +01:00
let derivations = derivations_interner.get(*derivations);
if derivations.is_empty() {
None
} else {
Some(derivations.original)
}
}
}
}
}
/// A query term term coupled with its position in the user's search query.
#[derive(Clone)]
pub struct LocatedQueryTerm {
pub value: QueryTerm,
pub positions: RangeInclusive<i8>,
}
impl LocatedQueryTerm {
/// Return `true` iff the word derivations within the query term are empty
2023-03-09 11:12:31 +01:00
pub fn is_empty(&self, interner: &Interner<WordDerivations>) -> bool {
match self.value {
// TODO: phrases should be greedily computed, so that they can be excluded from
// the query graph right from the start?
QueryTerm::Phrase { phrase: _ } => false,
2023-03-09 11:12:31 +01:00
QueryTerm::Word { derivations, .. } => interner.get(derivations).is_empty(),
}
}
}
/// Convert the tokenised search query into a list of located query terms.
2023-03-13 14:03:48 +01:00
pub fn located_query_terms_from_string<'ctx>(
ctx: &mut SearchContext<'ctx>,
query: NormalizedTokenIter<Vec<u8>>,
words_limit: Option<usize>,
) -> Result<Vec<LocatedQueryTerm>> {
let nbr_typos = number_of_typos_allowed(ctx)?;
let mut located_terms = Vec::new();
let mut phrase = Vec::new();
let mut quoted = false;
let parts_limit = words_limit.unwrap_or(usize::MAX);
let mut position = -1i8;
let mut phrase_start = -1i8;
let mut phrase_end = -1i8;
let mut peekable = query.peekable();
while let Some(token) = peekable.next() {
// early return if word limit is exceeded
if located_terms.len() >= parts_limit {
return Ok(located_terms);
}
match token.kind {
TokenKind::Word | TokenKind::StopWord => {
position += 1;
// 1. if the word is quoted we push it in a phrase-buffer waiting for the ending quote,
// 2. if the word is not the last token of the query and is not a stop_word we push it as a non-prefix word,
// 3. if the word is the last token of the query we push it as a prefix word.
if quoted {
phrase_end = position;
if phrase.is_empty() {
phrase_start = position;
}
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(word));
}
} else if peekable.peek().is_some() {
match token.kind {
TokenKind::Word => {
let word = token.lemma();
let derivations = word_derivations(ctx, word, nbr_typos(word), false)?;
let located_term = LocatedQueryTerm {
value: QueryTerm::Word {
derivations: ctx.derivations_interner.insert(derivations),
},
positions: position..=position,
};
located_terms.push(located_term);
}
TokenKind::StopWord | TokenKind::Separator(_) | TokenKind::Unknown => {}
}
} else {
let word = token.lemma();
let derivations = word_derivations(ctx, word, nbr_typos(word), true)?;
let located_term = LocatedQueryTerm {
value: QueryTerm::Word {
derivations: ctx.derivations_interner.insert(derivations),
},
positions: position..=position,
};
located_terms.push(located_term);
}
}
TokenKind::Separator(separator_kind) => {
match separator_kind {
SeparatorKind::Hard => {
position += 1;
}
SeparatorKind::Soft => {
position += 0;
}
}
let quote_count = token.lemma().chars().filter(|&s| s == '"').count();
// swap quoted state if we encounter a double quote
if quote_count % 2 != 0 {
quoted = !quoted;
}
// if there is a quote or a hard separator we close the phrase.
if !phrase.is_empty() && (quote_count > 0 || separator_kind == SeparatorKind::Hard)
{
let located_query_term = LocatedQueryTerm {
value: QueryTerm::Phrase {
phrase: ctx
.phrase_interner
.insert(Phrase { words: mem::take(&mut phrase) }),
},
positions: phrase_start..=phrase_end,
};
located_terms.push(located_query_term);
}
}
_ => (),
}
}
// 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: ctx.phrase_interner.insert(Phrase { words: mem::take(&mut phrase) }),
},
positions: phrase_start..=phrase_end,
};
located_terms.push(located_query_term);
}
Ok(located_terms)
}
pub fn number_of_typos_allowed<'ctx>(
ctx: &SearchContext<'ctx>,
) -> Result<impl Fn(&str) -> u8 + 'ctx> {
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)?;
// TODO: should `exact_words` also disable prefix search, ngrams, split words, or synonyms?
let exact_words = ctx.index.exact_words(ctx.txn)?;
Ok(Box::new(move |word: &str| {
if !authorize_typos
|| word.len() < min_len_one_typo as usize
|| exact_words.as_ref().map_or(false, |fst| fst.contains(word))
{
0
} else if word.len() < min_len_two_typos as usize {
1
} else {
2
}
}))
}
pub fn make_ngram(
ctx: &mut SearchContext,
terms: &[LocatedQueryTerm],
number_of_typos_allowed: &impl Fn(&str) -> u8,
) -> Result<Option<LocatedQueryTerm>> {
assert!(!terms.is_empty());
for ts in terms.windows(2) {
let [t1, t2] = ts else { panic!() };
if *t1.positions.end() != t2.positions.start() - 1 {
return Ok(None);
}
}
let mut words_interned = vec![];
for term in terms {
if let Some(original_term_word) = term.value.original_single_word(&ctx.derivations_interner)
{
words_interned.push(original_term_word);
} else {
return Ok(None);
}
}
let words =
words_interned.iter().map(|&i| ctx.word_interner.get(i).to_owned()).collect::<Vec<_>>();
let start = *terms.first().as_ref().unwrap().positions.start();
let end = *terms.last().as_ref().unwrap().positions.end();
let is_prefix = terms.last().as_ref().unwrap().value.is_prefix(&ctx.derivations_interner);
let ngram_str = words.join("");
if ngram_str.len() > MAX_WORD_LENGTH {
return Ok(None);
}
let mut derivations = word_derivations(
ctx,
&ngram_str,
number_of_typos_allowed(ngram_str.as_str()).saturating_sub(terms.len() as u8),
is_prefix,
)?;
derivations.original = ctx.word_interner.insert(words.join(" "));
// Now add the synonyms
let index_synonyms = ctx.index.synonyms(ctx.txn)?;
let mut derivations_synonyms = derivations.synonyms.to_vec();
derivations_synonyms.extend(
index_synonyms.get(&words).cloned().unwrap_or_default().into_iter().map(|words| {
let words = words.into_iter().map(|w| Some(ctx.word_interner.insert(w))).collect();
ctx.phrase_interner.insert(Phrase { words })
}),
);
derivations.synonyms = derivations_synonyms.into_boxed_slice();
if let Some(split_words) = derivations.split_words {
let split_words = ctx.phrase_interner.get(split_words);
if split_words.words == words_interned.iter().map(|&i| Some(i)).collect::<Vec<_>>() {
derivations.split_words = None;
}
}
if derivations.is_empty() {
return Ok(None);
}
let term = LocatedQueryTerm {
value: QueryTerm::Word { derivations: ctx.derivations_interner.insert(derivations) },
positions: start..=end,
};
Ok(Some(term))
}
// // TODO: return a word derivations instead?
// pub fn ngram2(
// ctx: &mut SearchContext,
// x: &LocatedQueryTerm,
// y: &LocatedQueryTerm,
// number_of_typos_allowed: impl Fn(&str) -> u8,
// ) -> Result<Option<LocatedQueryTerm>> {
// if *x.positions.end() != y.positions.start() - 1 {
// return Ok(None);
// }
// match (
// x.value.original_single_word(&ctx.word_interner, &ctx.derivations_interner),
// y.value.original_single_word(&ctx.word_interner, &ctx.derivations_interner),
// ) {
// (Some(w1), Some(w2)) => {
// let ngram2_str = format!("{w1}{w2}");
// let mut derivations = word_derivations(
// ctx,
// &ngram2_str,
// number_of_typos_allowed(ngram2_str.as_str()).saturating_sub(1),
// y.value.is_prefix(&ctx.derivations_interner),
// )?;
// // Now add the synonyms
// let index_synonyms = ctx.index.synonyms(ctx.txn)?;
// let mut derivations_synonyms = derivations.synonyms.to_vec();
// derivations_synonyms.extend(
// index_synonyms
// .get(&vec![w1.to_owned(), w2.to_owned()])
// .cloned()
// .unwrap_or_default()
// .into_iter()
// .map(|words| {
// let words =
// words.into_iter().map(|w| Some(ctx.word_interner.insert(w))).collect();
// ctx.phrase_interner.insert(Phrase { words })
// }),
// );
// let term = LocatedQueryTerm {
// value: QueryTerm::Word {
// derivations: ctx.derivations_interner.insert(derivations),
// },
// positions: *x.positions.start()..=*y.positions.end(),
// };
// Ok(Some(term))
// }
// _ => Ok(None),
// }
// }
// // TODO: return a word derivations instead?
// pub fn ngram3(
// ctx: &mut SearchContext,
// x: &LocatedQueryTerm,
// y: &LocatedQueryTerm,
// z: &LocatedQueryTerm,
// ) -> 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(&ctx.word_interner, &ctx.derivations_interner),
// &y.value.original_single_word(&ctx.word_interner, &ctx.derivations_interner),
// &z.value.original_single_word(&ctx.word_interner, &ctx.derivations_interner),
// ) {
// (Some(w1), Some(w2), Some(w3)) => {
// let term = (
// ctx.word_interner.insert(format!("{w1}{w2}{w3}")),
// *x.positions.start()..=*z.positions.end(),
// );
// Some(term)
// }
// _ => None,
// }
// }