Simplify some search function by reducing the number of parameters

This commit is contained in:
Kerollmops 2020-09-29 15:33:49 +02:00
parent 68f4af7d2e
commit f277ea134f
No known key found for this signature in database
GPG Key ID: 92ADA4E935E71FA4

View File

@ -80,8 +80,7 @@ impl<'a> Search<'a> {
/// Fetch the words from the given FST related to the given DFAs along with
/// the associated documents ids.
fn fetch_words_docids(
rtxn: &heed::RoTxn,
index: &Index,
&self,
fst: &fst::Set<&[u8]>,
dfas: Vec<(String, bool, DFA)>,
) -> anyhow::Result<Vec<(HashMap<String, (u8, RoaringBitmap)>, RoaringBitmap)>>
@ -98,7 +97,7 @@ impl<'a> Search<'a> {
while let Some((word, state)) = stream.next() {
let word = std::str::from_utf8(word)?;
let docids = index.word_docids.get(rtxn, word)?.unwrap();
let docids = self.index.word_docids.get(self.rtxn, word)?.unwrap();
let distance = dfa.distance(state);
unions_docids.union_with(&docids);
acc_derived_words.insert(word.to_string(), (distance.to_u8(), docids));
@ -134,8 +133,7 @@ impl<'a> Search<'a> {
}
fn fecth_keywords(
rtxn: &heed::RoTxn,
index: &Index,
&self,
derived_words: &[(HashMap<String, (u8, RoaringBitmap)>, RoaringBitmap)],
candidate: DocumentId,
) -> anyhow::Result<Vec<IntoIter>>
@ -148,7 +146,7 @@ impl<'a> Search<'a> {
for (word, (_distance, docids)) in words {
if !docids.contains(candidate) { continue; }
if let Some(positions) = index.docid_word_positions.get(rtxn, &(candidate, word))? {
if let Some(positions) = self.index.docid_word_positions.get(self.rtxn, &(candidate, word))? {
union_positions.union_with(&positions);
}
}
@ -158,43 +156,10 @@ impl<'a> Search<'a> {
Ok(keywords)
}
pub fn execute(&self) -> anyhow::Result<SearchResult> {
let rtxn = self.rtxn;
let index = self.index;
let limit = self.limit;
let fst = match index.fst(rtxn)? {
Some(fst) => fst,
None => return Ok(Default::default()),
};
// Construct the DFAs related to the query words.
// TODO do a placeholder search when query string isn't present.
let dfas = match &self.query {
Some(q) => Self::generate_query_dfas(q),
None => return Ok(Default::default()),
};
if dfas.is_empty() {
return Ok(Default::default());
}
let derived_words = Self::fetch_words_docids(rtxn, index, &fst, dfas)?;
let candidates = Self::compute_candidates(&derived_words);
debug!("candidates: {:?}", candidates);
// If there is only one query word, no need to compute the best proximities.
if derived_words.len() == 1 || candidates.is_empty() {
let found_words = derived_words.into_iter().flat_map(|(w, _)| w).map(|(w, _)| w).collect();
let documents_ids = candidates.iter().take(limit).collect();
return Ok(SearchResult { found_words, documents_ids });
}
fn words_pair_combinations<'a>(
w1: &'a HashMap<String, (u8, RoaringBitmap)>,
w2: &'a HashMap<String, (u8, RoaringBitmap)>,
) -> Vec<(&'a str, &'a str)>
fn words_pair_combinations<'h>(
w1: &'h HashMap<String, (u8, RoaringBitmap)>,
w2: &'h HashMap<String, (u8, RoaringBitmap)>,
) -> Vec<(&'h str, &'h str)>
{
let mut pairs = Vec::new();
for (w1, (_typos, docids1)) in w1 {
@ -208,8 +173,7 @@ impl<'a> Search<'a> {
}
fn depth_first_search(
index: &Index,
rtxn: &heed::RoTxn,
&self,
words: &[(HashMap<String, (u8, RoaringBitmap)>, RoaringBitmap)],
candidates: &RoaringBitmap,
parent_docids: Option<&RoaringBitmap>,
@ -217,7 +181,7 @@ impl<'a> Search<'a> {
) -> anyhow::Result<Option<RoaringBitmap>>
{
let (words1, words2) = (&words[0].0, &words[1].0);
let pairs = words_pair_combinations(words1, words2);
let pairs = Self::words_pair_combinations(words1, words2);
for proximity in 1..=8 {
let mut docids = match union_cache.entry((words.len(), proximity)) {
@ -229,7 +193,7 @@ impl<'a> Search<'a> {
} else {
for (w1, w2) in pairs.iter().cloned() {
let key = (w1, w2, proximity);
if let Some(di) = index.word_pair_proximity_docids.get(rtxn, &key)? {
if let Some(di) = self.index.word_pair_proximity_docids.get(self.rtxn, &key)? {
docids.union_with(&di);
}
}
@ -246,7 +210,7 @@ impl<'a> Search<'a> {
let words = &words[1..];
// We are the last word.
if words.len() < 2 { return Ok(Some(docids)) }
if let Some(di) = depth_first_search(index, rtxn, words, candidates, Some(&docids), union_cache)? {
if let Some(di) = self.depth_first_search(words, candidates, Some(&docids), union_cache)? {
return Ok(Some(di))
}
}
@ -255,9 +219,38 @@ impl<'a> Search<'a> {
Ok(None)
}
let mut union_cache = HashMap::new();
let answer = depth_first_search(index, rtxn, &derived_words, &candidates, None, &mut union_cache)?;
pub fn execute(&self) -> anyhow::Result<SearchResult> {
let limit = self.limit;
let fst = match self.index.fst(self.rtxn)? {
Some(fst) => fst,
None => return Ok(Default::default()),
};
// Construct the DFAs related to the query words.
// TODO do a placeholder search when query string isn't present.
let dfas = match &self.query {
Some(q) => Self::generate_query_dfas(q),
None => return Ok(Default::default()),
};
if dfas.is_empty() {
return Ok(Default::default());
}
let derived_words = self.fetch_words_docids(&fst, dfas)?;
let mut candidates = Self::compute_candidates(&derived_words);
debug!("candidates: {:?}", candidates);
// If there is only one query word, no need to compute the best proximities.
if derived_words.len() == 1 || candidates.is_empty() {
let found_words = derived_words.into_iter().flat_map(|(w, _)| w).map(|(w, _)| w).collect();
let documents_ids = candidates.iter().take(limit).collect();
return Ok(SearchResult { found_words, documents_ids });
}
let mut union_cache = HashMap::new();
let mut documents = Vec::new();
if let Some(answer) = answer {
documents.push(answer);