uses the new cancellation method in arroy

This commit is contained in:
Tamo 2024-10-01 17:45:49 +02:00
parent 4b598fa648
commit b1dc10e771
2 changed files with 5 additions and 3 deletions

View File

@ -699,6 +699,7 @@ where
for (embedder_name, dimension) in dimension { for (embedder_name, dimension) in dimension {
let wtxn = &mut *self.wtxn; let wtxn = &mut *self.wtxn;
let vector_arroy = self.index.vector_arroy; 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( let embedder_index = self.index.embedder_category_id.get(wtxn, &embedder_name)?.ok_or(
InternalError::DatabaseMissingEntry { db_name: "embedder_category_id", key: None }, InternalError::DatabaseMissingEntry { db_name: "embedder_category_id", key: None },
@ -713,7 +714,7 @@ where
pool.install(|| { pool.install(|| {
let mut writer = ArroyWrapper::new(vector_arroy, embedder_index, was_quantized); 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(()) Result::Ok(())
}) })
.map_err(InternalError::from)??; .map_err(InternalError::from)??;

View File

@ -82,6 +82,7 @@ impl ArroyWrapper {
rng: &mut R, rng: &mut R,
dimension: usize, dimension: usize,
quantizing: bool, quantizing: bool,
cancel: &(impl Fn() -> bool + Sync + Send),
) -> Result<(), arroy::Error> { ) -> Result<(), arroy::Error> {
for index in arroy_db_range_for_embedder(self.embedder_index) { for index in arroy_db_range_for_embedder(self.embedder_index) {
if self.quantized { if self.quantized {
@ -100,9 +101,9 @@ impl ArroyWrapper {
// sensitive. // sensitive.
if quantizing && !self.quantized { if quantizing && !self.quantized {
let writer = writer.prepare_changing_distance::<BinaryQuantizedCosine>(wtxn)?; 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)? { } else if writer.need_build(wtxn)? {
writer.builder(rng).build(wtxn)?; writer.builder(rng).cancel(cancel).build(wtxn)?;
} else if writer.is_empty(wtxn)? { } else if writer.is_empty(wtxn)? {
break; break;
} }