From ef9d9f8481d6d8d6808ab3ea3ae0941f4ab56860 Mon Sep 17 00:00:00 2001
From: Tamo <tamo@meilisearch.com>
Date: Thu, 13 Mar 2025 11:05:25 +0100
Subject: [PATCH] set the memory in arroy

---
 Cargo.lock                                     |  3 ++-
 crates/milli/src/update/index_documents/mod.rs | 10 +++++++++-
 crates/milli/src/update/new/indexer/mod.rs     |  3 +++
 crates/milli/src/update/new/indexer/write.rs   | 10 +++++++++-
 crates/milli/src/vector/mod.rs                 | 13 +++++++++++--
 5 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 79e56fb7c..8bcaac8d8 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -394,7 +394,7 @@ checksum = "96d30a06541fbafbc7f82ed10c06164cfbd2c401138f6addd8404629c4b16711"
 [[package]]
 name = "arroy"
 version = "0.6.0"
-source = "git+https://github.com/meilisearch/arroy?branch=main#80a7f1ba60bd7d88d55ce958a7579d664fc769ce"
+source = "git+https://github.com/meilisearch/arroy?branch=main#3350696381a4e29a838209663f39c1c58e9bc7b6"
 dependencies = [
  "bytemuck",
  "byteorder",
@@ -403,6 +403,7 @@ dependencies = [
  "memmap2",
  "nohash",
  "ordered-float",
+ "page_size",
  "rand",
  "rayon",
  "roaring",
diff --git a/crates/milli/src/update/index_documents/mod.rs b/crates/milli/src/update/index_documents/mod.rs
index ae082284a..5dc55030b 100644
--- a/crates/milli/src/update/index_documents/mod.rs
+++ b/crates/milli/src/update/index_documents/mod.rs
@@ -520,7 +520,15 @@ where
 
             pool.install(|| {
                 let mut writer = ArroyWrapper::new(vector_arroy, embedder_index, was_quantized);
-                writer.build_and_quantize(wtxn, &mut rng, dimension, is_quantizing, cancel)?;
+                writer.build_and_quantize(
+                    wtxn,
+                    &mut rng,
+                    dimension,
+                    is_quantizing,
+                    // Arroy should only use 50% of the memory
+                    self.indexer_config.max_memory.map(|mm| mm / 2),
+                    cancel,
+                )?;
                 Result::Ok(())
             })
             .map_err(InternalError::from)??;
diff --git a/crates/milli/src/update/new/indexer/mod.rs b/crates/milli/src/update/new/indexer/mod.rs
index 1cd227139..beda7c1ee 100644
--- a/crates/milli/src/update/new/indexer/mod.rs
+++ b/crates/milli/src/update/new/indexer/mod.rs
@@ -70,6 +70,8 @@ where
         max_memory: grenad_parameters.max_memory.map(|mm| mm * 5 / 100),
         ..grenad_parameters
     };
+    // Arroy should use 50% of the grenad memory instead of 5%
+    let arroy_memory = grenad_parameters.max_memory.map(|mm| mm * 10);
 
     // 5% percent of the allocated memory for the extractors, or min 100MiB
     // 5% percent of the allocated memory for the bbqueues, or min 50MiB
@@ -200,6 +202,7 @@ where
                 index,
                 wtxn,
                 index_embeddings,
+                arroy_memory,
                 &mut arroy_writers,
                 &indexing_context.must_stop_processing,
             )
diff --git a/crates/milli/src/update/new/indexer/write.rs b/crates/milli/src/update/new/indexer/write.rs
index a8bd3217f..d0951a30a 100644
--- a/crates/milli/src/update/new/indexer/write.rs
+++ b/crates/milli/src/update/new/indexer/write.rs
@@ -101,6 +101,7 @@ pub fn build_vectors<MSP>(
     index: &Index,
     wtxn: &mut RwTxn<'_>,
     index_embeddings: Vec<IndexEmbeddingConfig>,
+    arroy_memory: Option<usize>,
     arroy_writers: &mut HashMap<u8, (&str, &Embedder, ArroyWrapper, usize)>,
     must_stop_processing: &MSP,
 ) -> Result<()>
@@ -114,7 +115,14 @@ where
     let mut rng = rand::rngs::StdRng::seed_from_u64(42);
     for (_index, (_embedder_name, _embedder, writer, dimensions)) in arroy_writers {
         let dimensions = *dimensions;
-        writer.build_and_quantize(wtxn, &mut rng, dimensions, false, must_stop_processing)?;
+        writer.build_and_quantize(
+            wtxn,
+            &mut rng,
+            dimensions,
+            false,
+            arroy_memory,
+            must_stop_processing,
+        )?;
     }
 
     index.put_embedding_configs(wtxn, index_embeddings)?;
diff --git a/crates/milli/src/vector/mod.rs b/crates/milli/src/vector/mod.rs
index f67912b89..80efc210d 100644
--- a/crates/milli/src/vector/mod.rs
+++ b/crates/milli/src/vector/mod.rs
@@ -86,6 +86,7 @@ impl ArroyWrapper {
         rng: &mut R,
         dimension: usize,
         quantizing: bool,
+        arroy_memory: Option<usize>,
         cancel: &(impl Fn() -> bool + Sync + Send),
     ) -> Result<(), arroy::Error> {
         for index in arroy_db_range_for_embedder(self.embedder_index) {
@@ -105,9 +106,17 @@ impl ArroyWrapper {
                 // sensitive.
                 if quantizing && !self.quantized {
                     let writer = writer.prepare_changing_distance::<BinaryQuantizedCosine>(wtxn)?;
-                    writer.builder(rng).cancel(cancel).build(wtxn)?;
+                    writer
+                        .builder(rng)
+                        .available_memory(arroy_memory.unwrap_or(usize::MAX))
+                        .cancel(cancel)
+                        .build(wtxn)?;
                 } else if writer.need_build(wtxn)? {
-                    writer.builder(rng).cancel(cancel).build(wtxn)?;
+                    writer
+                        .builder(rng)
+                        .available_memory(arroy_memory.unwrap_or(usize::MAX))
+                        .cancel(cancel)
+                        .build(wtxn)?;
                 } else if writer.is_empty(wtxn)? {
                     break;
                 }