set the memory in arroy

This commit is contained in:
Tamo 2025-03-13 11:05:25 +01:00 committed by Kerollmops
parent d3d22d8ed4
commit ef9d9f8481
No known key found for this signature in database
GPG key ID: F250A4C4E3AE5F5F
5 changed files with 34 additions and 5 deletions

View file

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