Lazily compute the FSTs during indexing

Cargo.lock

@@ -394,7 +394,8 @@ checksum = "96d30a06541fbafbc7f82ed10c06164cfbd2c401138f6addd8404629c4b16711"
 [[package]]
 name = "arroy"
 version = "0.6.1"
-source = "git+https://github.com/meilisearch/arroy.git?rev=5b748bac2c69c65a97980901b02067a3a545e357#5b748bac2c69c65a97980901b02067a3a545e357"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "08e6111f351d004bd13e95ab540721272136fd3218b39d3ec95a2ea1c4e6a0a6"
 dependencies = [
  "bytemuck",
  "byteorder",

crates/index-scheduler/src/lib.rs

@@ -625,8 +625,8 @@ impl IndexScheduler {
         task_id: Option<TaskId>,
         dry_run: bool,
     ) -> Result<Task> {
-        // if the task doesn't delete or cancel anything and 40% of the task queue is full, we must refuse to enqueue the incoming task
-        if !matches!(&kind, KindWithContent::TaskDeletion { tasks, .. } | KindWithContent::TaskCancelation { tasks, .. } if !tasks.is_empty())
+        // if the task doesn't delete anything and 50% of the task queue is full, we must refuse to enqueue the incomming task
+        if !matches!(&kind, KindWithContent::TaskDeletion { tasks, .. } if !tasks.is_empty())
             && (self.env.non_free_pages_size()? * 100) / self.env.info().map_size as u64 > 40
         {
             return Err(Error::NoSpaceLeftInTaskQueue);

crates/index-scheduler/src/queue/mod.rs

@@ -292,6 +292,8 @@ impl Queue {
             return Ok(task);
         }
 
+        // Get rid of the mutability.
+        let task = task;
         self.tasks.register(wtxn, &task)?;
 
         Ok(task)

crates/index-scheduler/src/queue/test.rs

@@ -364,7 +364,7 @@ fn test_task_queue_is_full() {
     // we won't be able to test this error in an integration test thus as a best effort test I still ensure the error return the expected error code
     snapshot!(format!("{:?}", result.error_code()), @"NoSpaceLeftOnDevice");
 
-    // Even the task deletion and cancelation that don't delete anything should be refused
+    // Even the task deletion that doesn't delete anything shouldn't be accepted
     let result = index_scheduler
         .register(
             KindWithContent::TaskDeletion { query: S("test"), tasks: RoaringBitmap::new() },
@@ -373,39 +373,10 @@ fn test_task_queue_is_full() {
         )
         .unwrap_err();
     snapshot!(result, @"Meilisearch cannot receive write operations because the limit of the task database has been reached. Please delete tasks to continue performing write operations.");
-    let result = index_scheduler
-        .register(
-            KindWithContent::TaskCancelation { query: S("test"), tasks: RoaringBitmap::new() },
-            None,
-            false,
-        )
-        .unwrap_err();
-    snapshot!(result, @"Meilisearch cannot receive write operations because the limit of the task database has been reached. Please delete tasks to continue performing write operations.");
-
     // we won't be able to test this error in an integration test thus as a best effort test I still ensure the error return the expected error code
     snapshot!(format!("{:?}", result.error_code()), @"NoSpaceLeftOnDevice");
 
-    // But a task cancelation that cancel something should work
-    index_scheduler
-        .register(
-            KindWithContent::TaskCancelation { query: S("test"), tasks: (0..100).collect() },
-            None,
-            false,
-        )
-        .unwrap();
-    handle.advance_one_successful_batch();
-
-    // But we should still be forbidden from enqueuing new tasks
-    let result = index_scheduler
-        .register(
-            KindWithContent::IndexCreation { index_uid: S("doggo"), primary_key: None },
-            None,
-            false,
-        )
-        .unwrap_err();
-    snapshot!(result, @"Meilisearch cannot receive write operations because the limit of the task database has been reached. Please delete tasks to continue performing write operations.");
-
-    // And a task deletion that delete something should works
+    // But a task deletion that delete something should works
     index_scheduler
         .register(
             KindWithContent::TaskDeletion { query: S("test"), tasks: (0..100).collect() },

crates/milli/Cargo.toml

@@ -87,8 +87,7 @@ rhai = { git = "https://github.com/rhaiscript/rhai", rev = "ef3df63121d27aacd838
     "no_time",
     "sync",
 ] }
-# arroy = "0.6.1"
-arroy = { git = "https://github.com/meilisearch/arroy.git", rev = "5b748bac2c69c65a97980901b02067a3a545e357" } # incremental update
+arroy = "0.6.1"
 rand = "0.8.5"
 tracing = "0.1.41"
 ureq = { version = "2.12.1", features = ["json"] }

crates/milli/src/progress.rs

@@ -266,15 +266,12 @@ impl Step for arroy::MainStep {
                 "writing the descendants and metadata"
             }
             arroy::MainStep::RetrieveTheUpdatedItems => "retrieve the updated items",
+            arroy::MainStep::RetrievingTheTreeAndItemNodes => "retrieving the tree and item nodes",
+            arroy::MainStep::UpdatingTheTrees => "updating the trees",
+            arroy::MainStep::CreateNewTrees => "create new trees",
+            arroy::MainStep::WritingNodesToDatabase => "writing nodes to database",
+            arroy::MainStep::DeleteExtraneousTrees => "delete extraneous trees",
             arroy::MainStep::WriteTheMetadata => "write the metadata",
-            arroy::MainStep::RetrievingTheItemsIds => "retrieving the items ids",
-            arroy::MainStep::RetrievingTheUsedTreeNodes => "retrieving the used tree nodes",
-            arroy::MainStep::DeletingExtraTrees => "deleting extra trees",
-            arroy::MainStep::RemoveItemsFromExistingTrees => "remove items from existing trees",
-            arroy::MainStep::InsertItemsInCurrentTrees => "insert items in current trees",
-            arroy::MainStep::IncrementalIndexLargeDescendants => {
-                "incremental index large descendants"
-            }
         }
         .into()
     }

crates/milli/src/update/new/facet_search_builder.rs

@@ -144,7 +144,7 @@ impl<'indexer> FacetSearchBuilder<'indexer> {
         let mut merger_iter = builder.build().into_stream_merger_iter()?;
         let mut current_field_id = None;
         let mut fst;
-        let mut fst_merger_builder: Option<FstMergerBuilder> = None;
+        let mut fst_merger_builder: Option<FstMergerBuilder<_>> = None;
         while let Some((key, deladd)) = merger_iter.next()? {
             let (field_id, normalized_facet_string) =
                 BEU16StrCodec::bytes_decode(key).map_err(heed::Error::Encoding)?;
@@ -153,12 +153,13 @@ impl<'indexer> FacetSearchBuilder<'indexer> {
                 if let (Some(current_field_id), Some(fst_merger_builder)) =
                     (current_field_id, fst_merger_builder)
                 {
-                    let mmap = fst_merger_builder.build(&mut callback)?;
-                    index.facet_id_string_fst.remap_data_type::<Bytes>().put(
-                        wtxn,
-                        &current_field_id,
-                        &mmap,
-                    )?;
+                    if let Some(mmap) = fst_merger_builder.build(&mut callback)? {
+                        index.facet_id_string_fst.remap_data_type::<Bytes>().put(
+                            wtxn,
+                            &current_field_id,
+                            &mmap,
+                        )?;
+                    }
                 }
 
                 fst = index.facet_id_string_fst.get(rtxn, &field_id)?;
@@ -209,8 +210,9 @@ impl<'indexer> FacetSearchBuilder<'indexer> {
         }
 
         if let (Some(field_id), Some(fst_merger_builder)) = (current_field_id, fst_merger_builder) {
-            let mmap = fst_merger_builder.build(&mut callback)?;
-            index.facet_id_string_fst.remap_data_type::<Bytes>().put(wtxn, &field_id, &mmap)?;
+            if let Some(mmap) = fst_merger_builder.build(&mut callback)? {
+                index.facet_id_string_fst.remap_data_type::<Bytes>().put(wtxn, &field_id, &mmap)?;
+            }
         }
 
         Ok(())

crates/milli/src/update/new/fst_merger_builder.rs

@@ -1,25 +1,27 @@
 use std::fs::File;
 use std::io::BufWriter;
 
-use fst::{Set, SetBuilder, Streamer};
+use fst::{IntoStreamer, Set, SetBuilder, Streamer};
 use memmap2::Mmap;
 use tempfile::tempfile;
 
 use crate::update::del_add::DelAdd;
 use crate::{InternalError, Result};
 
-pub struct FstMergerBuilder<'a> {
+pub struct FstMergerBuilder<'a, D: AsRef<[u8]>> {
+    fst: Option<&'a Set<D>>,
     stream: Option<fst::set::Stream<'a>>,
-    fst_builder: SetBuilder<BufWriter<File>>,
+    fst_builder: Option<SetBuilder<BufWriter<File>>>,
     last: Option<Vec<u8>>,
     inserted_words: usize,
 }
 
-impl<'a> FstMergerBuilder<'a> {
-    pub fn new<D: AsRef<[u8]>>(fst: Option<&'a Set<D>>) -> Result<Self> {
+impl<'a, D: AsRef<[u8]>> FstMergerBuilder<'a, D> {
+    pub fn new(fst: Option<&'a Set<D>>) -> Result<Self> {
         Ok(Self {
+            fst,
             stream: fst.map(|fst| fst.stream()),
-            fst_builder: SetBuilder::new(BufWriter::new(tempfile()?))?,
+            fst_builder: None,
             last: None,
             inserted_words: 0,
         })
@@ -110,11 +112,17 @@ impl<'a> FstMergerBuilder<'a> {
         is_modified: bool,
         insertion_callback: &mut impl FnMut(&[u8], DelAdd, bool) -> Result<()>,
     ) -> Result<()> {
-        // Addition: We insert the word
-        // Deletion: We delete the word by not inserting it
-        if deladd == DelAdd::Addition {
-            self.inserted_words += 1;
-            self.fst_builder.insert(bytes)?;
+        if is_modified && self.fst_builder.is_none() {
+            self.build_new_fst(bytes)?;
+        }
+
+        if let Some(fst_builder) = self.fst_builder.as_mut() {
+            // Addition: We insert the word
+            // Deletion: We delete the word by not inserting it
+            if deladd == DelAdd::Addition {
+                self.inserted_words += 1;
+                fst_builder.insert(bytes)?;
+            }
         }
 
         insertion_callback(bytes, deladd, is_modified)?;
@@ -122,6 +130,19 @@ impl<'a> FstMergerBuilder<'a> {
         Ok(())
     }
 
+    // Lazily build the new fst
+    fn build_new_fst(&mut self, bytes: &[u8]) -> Result<()> {
+        let mut fst_builder = SetBuilder::new(BufWriter::new(tempfile()?))?;
+
+        if let Some(fst) = self.fst {
+            fst_builder.extend_stream(fst.range().lt(bytes).into_stream())?;
+        }
+
+        self.fst_builder = Some(fst_builder);
+
+        Ok(())
+    }
+
     fn drain_stream(
         &mut self,
         insertion_callback: &mut impl FnMut(&[u8], DelAdd, bool) -> Result<()>,
@@ -142,16 +163,20 @@ impl<'a> FstMergerBuilder<'a> {
     pub fn build(
         mut self,
         insertion_callback: &mut impl FnMut(&[u8], DelAdd, bool) -> Result<()>,
-    ) -> Result<Mmap> {
+    ) -> Result<Option<Mmap>> {
         self.drain_stream(insertion_callback)?;
 
-        let fst_file = self
-            .fst_builder
-            .into_inner()?
-            .into_inner()
-            .map_err(|_| InternalError::IndexingMergingKeys { process: "building-fst" })?;
-        let fst_mmap = unsafe { Mmap::map(&fst_file)? };
+        match self.fst_builder {
+            Some(fst_builder) => {
+                let fst_file = fst_builder
+                    .into_inner()?
+                    .into_inner()
+                    .map_err(|_| InternalError::IndexingMergingKeys { process: "building-fst" })?;
+                let fst_mmap = unsafe { Mmap::map(&fst_file)? };
 
-        Ok(fst_mmap)
+                Ok(Some(fst_mmap))
+            }
+            None => Ok(None),
+        }
     }
 }

crates/milli/src/update/new/indexer/post_processing.rs

@@ -118,7 +118,9 @@ fn compute_word_fst(
     }
 
     let (word_fst_mmap, prefix_data) = word_fst_builder.build(index, &rtxn)?;
-    index.main.remap_types::<Str, Bytes>().put(wtxn, WORDS_FST_KEY, &word_fst_mmap)?;
+    if let Some(word_fst_mmap) = word_fst_mmap {
+        index.main.remap_types::<Str, Bytes>().put(wtxn, WORDS_FST_KEY, &word_fst_mmap)?;
+    }
     if let Some(PrefixData { prefixes_fst_mmap, prefix_delta }) = prefix_data {
         index.main.remap_types::<Str, Bytes>().put(
             wtxn,

crates/milli/src/update/new/word_fst_builder.rs

@@ -10,14 +10,14 @@ use crate::index::PrefixSettings;
 use crate::update::del_add::DelAdd;
 use crate::{InternalError, Prefix, Result};
 
-pub struct WordFstBuilder<'a> {
-    word_fst_builder: FstMergerBuilder<'a>,
+pub struct WordFstBuilder<'a, D: AsRef<[u8]>> {
+    word_fst_builder: FstMergerBuilder<'a, D>,
     prefix_fst_builder: Option<PrefixFstBuilder>,
     registered_words: usize,
 }
 
-impl<'a> WordFstBuilder<'a> {
-    pub fn new(words_fst: &'a Set<std::borrow::Cow<'a, [u8]>>) -> Result<Self> {
+impl<'a, D: AsRef<[u8]>> WordFstBuilder<'a, D> {
+    pub fn new(words_fst: &'a Set<D>) -> Result<Self> {
         Ok(Self {
             word_fst_builder: FstMergerBuilder::new(Some(words_fst))?,
             prefix_fst_builder: None,
@@ -50,7 +50,7 @@ impl<'a> WordFstBuilder<'a> {
         mut self,
         index: &crate::Index,
         rtxn: &heed::RoTxn,
-    ) -> Result<(Mmap, Option<PrefixData>)> {
+    ) -> Result<(Option<Mmap>, Option<PrefixData>)> {
         let words_fst_mmap = self.word_fst_builder.build(&mut |bytes, deladd, is_modified| {
             if let Some(prefix_fst_builder) = &mut self.prefix_fst_builder {
                 prefix_fst_builder.insert_word(bytes, deladd, is_modified)