Support swapped word pairs in new proximity ranking rule impl

This commit is contained in:
Loïc Lecrenier 2023-02-21 15:10:22 +01:00
parent a61495d660
commit dd12d44134
3 changed files with 30 additions and 0 deletions

View File

@ -11,6 +11,8 @@ pub struct DatabaseCache<'transaction> {
pub word_pair_proximity_docids: FxHashMap<(u8, String, String), Option<&'transaction [u8]>>,
pub word_prefix_pair_proximity_docids:
FxHashMap<(u8, String, String), Option<&'transaction [u8]>>,
pub prefix_word_pair_proximity_docids:
FxHashMap<(u8, String, String), Option<&'transaction [u8]>>,
pub word_docids: FxHashMap<String, Option<&'transaction [u8]>>,
pub exact_word_docids: FxHashMap<String, Option<&'transaction [u8]>>,
pub word_prefix_docids: FxHashMap<String, Option<&'transaction [u8]>>,
@ -115,4 +117,25 @@ impl<'transaction> DatabaseCache<'transaction> {
}
}
}
pub fn get_prefix_word_pair_proximity_docids(
&mut self,
index: &Index,
txn: &'transaction RoTxn,
word1: &str,
prefix2: &str,
proximity: u8,
) -> Result<Option<&'transaction [u8]>> {
let key = (proximity, prefix2.to_owned(), word1.to_owned());
match self.prefix_word_pair_proximity_docids.entry(key) {
Entry::Occupied(bitmap_ptr) => Ok(*bitmap_ptr.get()),
Entry::Vacant(entry) => {
let bitmap_ptr = index
.prefix_word_pair_proximity_docids
.remap_data_type::<ByteSlice>()
.get(txn, &(proximity, prefix2, word1))?;
entry.insert(bitmap_ptr);
Ok(bitmap_ptr)
}
}
}
}

View File

@ -20,6 +20,11 @@ pub fn compute_docids<'transaction>(
}
WordPair::WordPrefix { left, right_prefix } => db_cache
.get_word_prefix_pair_proximity_docids(index, txn, left, right_prefix, *proximity),
WordPair::WordsSwapped { left, right } => {
db_cache.get_word_pair_proximity_docids(index, txn, left, right, *proximity)
}
WordPair::WordPrefixSwapped { left, right_prefix } => db_cache
.get_prefix_word_pair_proximity_docids(index, txn, left, right_prefix, *proximity),
}?;
let bitmap =
bytes.map(CboRoaringBitmapCodec::deserialize_from).transpose()?.unwrap_or_default();

View File

@ -13,7 +13,9 @@ use crate::{Index, Result};
pub enum WordPair {
// TODO: add WordsSwapped and WordPrefixSwapped case
Words { left: String, right: String },
WordsSwapped { left: String, right: String },
WordPrefix { left: String, right_prefix: String },
WordPrefixSwapped { left: String, right_prefix: String },
}
pub struct ProximityEdge {