Fix the computation of the newly added and common prefix words

milli/src/update/index_documents/helpers/mod.rs

@@ -46,7 +46,7 @@ pub fn read_u32_ne_bytes(bytes: &[u8]) -> impl Iterator<Item = u32> + '_ {
     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<Vec<u8>>
 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<String>
+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
+}

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};

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<A>,
     ) -> 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()))
                 }