From d59e5593179accbfe71dda09b482bbed69c8cb1e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Cl=C3=A9ment=20Renault?= Date: Thu, 20 Jan 2022 17:55:52 +0100 Subject: [PATCH] Fix the computation of the newly added and common prefix words --- .../src/update/index_documents/helpers/mod.rs | 17 +++++++++++++++- milli/src/update/index_documents/mod.rs | 6 +++--- milli/src/update/word_prefix_docids.rs | 20 ++++++++++++------- 3 files changed, 32 insertions(+), 11 deletions(-) diff --git a/milli/src/update/index_documents/helpers/mod.rs b/milli/src/update/index_documents/helpers/mod.rs index 4086bfb7f..bbb2b9b95 100644 --- a/milli/src/update/index_documents/helpers/mod.rs +++ b/milli/src/update/index_documents/helpers/mod.rs @@ -46,7 +46,7 @@ pub fn read_u32_ne_bytes(bytes: &[u8]) -> impl Iterator + '_ { bytes.chunks_exact(4).flat_map(TryInto::try_into).map(u32::from_ne_bytes) } -/// Converts an fst Stream into an HashSet. +/// Converts an fst Stream into an HashSet of Strings. pub fn fst_stream_into_hashset<'f, I, S>(stream: I) -> HashSet> where I: for<'a> IntoStreamer<'a, Into = S, Item = &'a [u8]>, @@ -59,3 +59,18 @@ where } hashset } + +// Converts an fst Stream into a Vec of Strings. +pub fn fst_stream_into_vec<'f, I, S>(stream: I) -> Vec +where + I: for<'a> IntoStreamer<'a, Into = S, Item = &'a [u8]>, + S: 'f + for<'a> Streamer<'a, Item = &'a [u8]>, +{ + let mut strings = Vec::new(); + let mut stream = stream.into_stream(); + while let Some(word) = stream.next() { + let s = std::str::from_utf8(word).unwrap(); + strings.push(s.to_owned()); + } + strings +} diff --git a/milli/src/update/index_documents/mod.rs b/milli/src/update/index_documents/mod.rs index 77b761e6e..ad3f73d0d 100644 --- a/milli/src/update/index_documents/mod.rs +++ b/milli/src/update/index_documents/mod.rs @@ -15,9 +15,9 @@ use serde::{Deserialize, Serialize}; use typed_chunk::{write_typed_chunk_into_index, TypedChunk}; pub use self::helpers::{ - create_sorter, create_writer, fst_stream_into_hashset, merge_cbo_roaring_bitmaps, - merge_roaring_bitmaps, sorter_into_lmdb_database, write_into_lmdb_database, writer_into_reader, - ClonableMmap, MergeFn, + create_sorter, create_writer, fst_stream_into_hashset, fst_stream_into_vec, + merge_cbo_roaring_bitmaps, merge_roaring_bitmaps, sorter_into_lmdb_database, + write_into_lmdb_database, writer_into_reader, ClonableMmap, MergeFn, }; use self::helpers::{grenad_obkv_into_chunks, GrenadParameters}; pub use self::transform::{Transform, TransformOutput}; diff --git a/milli/src/update/word_prefix_docids.rs b/milli/src/update/word_prefix_docids.rs index 105083d87..1e2996c9b 100644 --- a/milli/src/update/word_prefix_docids.rs +++ b/milli/src/update/word_prefix_docids.rs @@ -1,12 +1,13 @@ use std::collections::HashMap; -use fst::IntoStreamer; +use fst::Streamer; use grenad::{CompressionType, MergerBuilder}; +use heed::types::ByteSlice; use slice_group_by::GroupBy; use crate::update::index_documents::{ - create_sorter, fst_stream_into_hashset, merge_roaring_bitmaps, sorter_into_lmdb_database, - CursorClonableMmap, MergeFn, WriteMethod, + create_sorter, fst_stream_into_hashset, fst_stream_into_vec, merge_roaring_bitmaps, + sorter_into_lmdb_database, CursorClonableMmap, MergeFn, WriteMethod, }; use crate::{Index, Result}; @@ -41,9 +42,14 @@ impl<'t, 'u, 'i> WordPrefixDocids<'t, 'u, 'i> { old_prefix_fst: &fst::Set, ) -> Result<()> { let prefix_fst = self.index.words_prefixes_fst(self.wtxn)?; - let prefix_fst_keys = prefix_fst.into_stream().into_strs()?; - let prefix_fst_keys: Vec<_> = - prefix_fst_keys.as_slice().linear_group_by_key(|x| x.chars().nth(0).unwrap()).collect(); + + // We retrieve the common words between the previous and new prefix word fst. + let common_prefix_fst_keys = + fst_stream_into_vec(old_prefix_fst.op().add(&prefix_fst).intersection()); + let common_prefix_fst_keys: Vec<_> = common_prefix_fst_keys + .as_slice() + .linear_group_by_key(|x| x.chars().nth(0).unwrap()) + .collect(); // We compute the set of prefixes that are no more part of the prefix fst. let suppr_pw = fst_stream_into_hashset(old_prefix_fst.op().add(&prefix_fst).difference()); @@ -69,7 +75,7 @@ impl<'t, 'u, 'i> WordPrefixDocids<'t, 'u, 'i> { Some(prefixes) if word.starts_with(&prefixes[0].as_bytes()) => Some(prefixes), _otherwise => { write_prefixes_in_sorter(&mut prefixes_cache, &mut prefix_docids_sorter)?; - prefix_fst_keys + common_prefix_fst_keys .iter() .find(|prefixes| word.starts_with(&prefixes[0].as_bytes())) }