From b1dc10e771a757826fe400280c8bac84976ce95b Mon Sep 17 00:00:00 2001
From: Tamo <tamo@meilisearch.com>
Date: Tue, 1 Oct 2024 17:45:49 +0200
Subject: [PATCH] uses the new cancellation method in arroy

---
 milli/src/update/index_documents/mod.rs | 3 ++-
 milli/src/vector/mod.rs                 | 5 +++--
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/milli/src/update/index_documents/mod.rs b/milli/src/update/index_documents/mod.rs
index e164a0817..88d20fff0 100644
--- a/milli/src/update/index_documents/mod.rs
+++ b/milli/src/update/index_documents/mod.rs
@@ -699,6 +699,7 @@ where
         for (embedder_name, dimension) in dimension {
             let wtxn = &mut *self.wtxn;
             let vector_arroy = self.index.vector_arroy;
+            let cancel = &self.should_abort;
 
             let embedder_index = self.index.embedder_category_id.get(wtxn, &embedder_name)?.ok_or(
                 InternalError::DatabaseMissingEntry { db_name: "embedder_category_id", key: None },
@@ -713,7 +714,7 @@ where
 
             pool.install(|| {
                 let mut writer = ArroyWrapper::new(vector_arroy, embedder_index, was_quantized);
-                writer.build_and_quantize(wtxn, &mut rng, dimension, is_quantizing)?;
+                writer.build_and_quantize(wtxn, &mut rng, dimension, is_quantizing, cancel)?;
                 Result::Ok(())
             })
             .map_err(InternalError::from)??;
diff --git a/milli/src/vector/mod.rs b/milli/src/vector/mod.rs
index 097e93ad2..571c02c8c 100644
--- a/milli/src/vector/mod.rs
+++ b/milli/src/vector/mod.rs
@@ -82,6 +82,7 @@ impl ArroyWrapper {
         rng: &mut R,
         dimension: usize,
         quantizing: bool,
+        cancel: &(impl Fn() -> bool + Sync + Send),
     ) -> Result<(), arroy::Error> {
         for index in arroy_db_range_for_embedder(self.embedder_index) {
             if self.quantized {
@@ -100,9 +101,9 @@ impl ArroyWrapper {
                 // sensitive.
                 if quantizing && !self.quantized {
                     let writer = writer.prepare_changing_distance::<BinaryQuantizedCosine>(wtxn)?;
-                    writer.builder(rng).build(wtxn)?;
+                    writer.builder(rng).cancel(cancel).build(wtxn)?;
                 } else if writer.need_build(wtxn)? {
-                    writer.builder(rng).build(wtxn)?;
+                    writer.builder(rng).cancel(cancel).build(wtxn)?;
                 } else if writer.is_empty(wtxn)? {
                     break;
                 }