Fix a bug where we were skipping most of the prefix pairs

This commit is contained in:
Clément Renault 2022-01-12 15:22:06 +01:00
parent 9f2ff71581
commit e3c34684c6
No known key found for this signature in database
GPG Key ID: 92ADA4E935E71FA4

View File

@ -64,28 +64,26 @@ impl<'t, 'u, 'i> WordPrefixPairProximityDocids<'t, 'u, 'i> {
); );
let prefix_fst = self.index.words_prefixes_fst(self.wtxn)?; let prefix_fst = self.index.words_prefixes_fst(self.wtxn)?;
let prefix_fst_keys = prefix_fst.into_stream().into_bytes(); let prefix_fst_keys = prefix_fst.into_stream().into_strs()?;
let prefix_fst_keys: Vec<_> = prefix_fst_keys let prefix_fst_keys: Vec<_> =
.as_slice() prefix_fst_keys.as_slice().linear_group_by_key(|x| x.chars().nth(0).unwrap()).collect();
.linear_group_by_key(|x| std::str::from_utf8(&x).unwrap().chars().nth(0).unwrap())
.collect();
let mut db = let mut db =
self.index.word_pair_proximity_docids.remap_data_type::<ByteSlice>().iter(self.wtxn)?; self.index.word_pair_proximity_docids.remap_data_type::<ByteSlice>().iter(self.wtxn)?;
let mut buffer = Vec::new(); let mut buffer = Vec::new();
let mut current_prefixes: Option<&&[Vec<u8>]> = None; let mut current_prefixes: Option<&&[String]> = None;
let mut prefixes_cache = HashMap::new(); let mut prefixes_cache = HashMap::new();
while let Some(((w1, w2, prox), data)) = db.next().transpose()? { while let Some(((w1, w2, prox), data)) = db.next().transpose()? {
current_prefixes = match current_prefixes.take() { current_prefixes = match current_prefixes.take() {
Some(prefixes) if w2.as_bytes().starts_with(&prefixes[0]) => Some(prefixes), Some(prefixes) if w2.starts_with(&prefixes[0]) => Some(prefixes),
_otherwise => { _otherwise => {
write_prefixes_in_sorter( write_prefixes_in_sorter(
&mut prefixes_cache, &mut prefixes_cache,
&mut word_prefix_pair_proximity_docids_sorter, &mut word_prefix_pair_proximity_docids_sorter,
self.threshold, self.threshold,
)?; )?;
prefix_fst_keys.iter().find(|prefixes| w2.as_bytes().starts_with(&prefixes[0])) prefix_fst_keys.iter().find(|prefixes| w2.starts_with(&prefixes[0]))
} }
}; };
@ -93,9 +91,9 @@ impl<'t, 'u, 'i> WordPrefixPairProximityDocids<'t, 'u, 'i> {
buffer.clear(); buffer.clear();
buffer.extend_from_slice(w1.as_bytes()); buffer.extend_from_slice(w1.as_bytes());
buffer.push(0); buffer.push(0);
for prefix in prefixes.iter().filter(|prefix| w2.as_bytes().starts_with(prefix)) { for prefix in prefixes.iter().filter(|prefix| w2.starts_with(prefix.as_str())) {
buffer.truncate(w1.len() + 1); buffer.truncate(w1.len() + 1);
buffer.extend_from_slice(prefix); buffer.extend_from_slice(prefix.as_bytes());
buffer.push(prox); buffer.push(prox);
match prefixes_cache.get_mut(&buffer) { match prefixes_cache.get_mut(&buffer) {
@ -135,17 +133,13 @@ fn write_prefixes_in_sorter(
sorter: &mut grenad::Sorter<MergeFn>, sorter: &mut grenad::Sorter<MergeFn>,
min_word_per_prefix: u32, min_word_per_prefix: u32,
) -> Result<()> { ) -> Result<()> {
for (i, (key, data_slices)) in prefixes.drain().enumerate() { for (key, data_slices) in prefixes.drain() {
// if the number of words prefixed by the prefix is higher than the threshold, // if the number of words prefixed by the prefix is higher than the threshold,
// we insert it in the sorter. // we insert it in the sorter.
if data_slices.len() > min_word_per_prefix as usize { if data_slices.len() > min_word_per_prefix as usize {
for data in data_slices { for data in data_slices {
sorter.insert(&key, data)?; sorter.insert(&key, data)?;
} }
// if the first prefix isn't elligible for insertion,
// then the other prefixes can't be elligible.
} else if i == 0 {
break;
} }
} }