update arroy

This commit is contained in:
Tamo 2024-09-30 13:12:01 +02:00
parent 71b364286b
commit 4b598fa648
5 changed files with 27 additions and 16 deletions

View file

@ -297,6 +297,7 @@ impl From<arroy::Error> for Error {
arroy::Error::InvalidVecDimension { expected, received } => {
Error::UserError(UserError::InvalidVectorDimensions { expected, found: received })
}
arroy::Error::BuildCancelled => Error::InternalError(InternalError::AbortedIndexation),
arroy::Error::DatabaseFull
| arroy::Error::InvalidItemAppend
| arroy::Error::UnmatchingDistance { .. }

View file

@ -1,7 +1,7 @@
use std::collections::HashMap;
use std::sync::Arc;
use arroy::distances::{Angular, BinaryQuantizedAngular};
use arroy::distances::{BinaryQuantizedCosine, Cosine};
use arroy::ItemId;
use deserr::{DeserializeError, Deserr};
use heed::{RoTxn, RwTxn, Unspecified};
@ -87,7 +87,7 @@ impl ArroyWrapper {
if self.quantized {
let writer = arroy::Writer::new(self.quantized_db(), index, dimension);
if writer.need_build(wtxn)? {
writer.build(wtxn, rng, None)?
writer.builder(rng).build(wtxn)?
} else if writer.is_empty(wtxn)? {
break;
}
@ -99,11 +99,10 @@ impl ArroyWrapper {
// only happens once in the life of an embedder, it's not very performances
// sensitive.
if quantizing && !self.quantized {
let writer =
writer.prepare_changing_distance::<BinaryQuantizedAngular>(wtxn)?;
writer.build(wtxn, rng, None)?
let writer = writer.prepare_changing_distance::<BinaryQuantizedCosine>(wtxn)?;
writer.builder(rng).build(wtxn)?;
} else if writer.need_build(wtxn)? {
writer.build(wtxn, rng, None)?
writer.builder(rng).build(wtxn)?;
} else if writer.is_empty(wtxn)? {
break;
}
@ -323,8 +322,13 @@ impl ArroyWrapper {
let mut results = Vec::new();
for reader in self.readers(rtxn, db) {
let ret = reader?.nns_by_item(rtxn, item, limit, None, None, filter)?;
if let Some(mut ret) = ret {
let reader = reader?;
let mut searcher = reader.nns(limit);
if let Some(filter) = filter {
searcher.candidates(filter);
}
if let Some(mut ret) = searcher.by_item(rtxn, item)? {
results.append(&mut ret);
} else {
break;
@ -359,8 +363,13 @@ impl ArroyWrapper {
let mut results = Vec::new();
for reader in self.readers(rtxn, db) {
let mut ret = reader?.nns_by_vector(rtxn, vector, limit, None, None, filter)?;
results.append(&mut ret);
let reader = reader?;
let mut searcher = reader.nns(limit);
if let Some(filter) = filter {
searcher.candidates(filter);
}
results.append(&mut searcher.by_vector(rtxn, vector)?);
}
results.sort_unstable_by_key(|(_, distance)| OrderedFloat(*distance));
@ -391,11 +400,11 @@ impl ArroyWrapper {
Ok(vectors)
}
fn angular_db(&self) -> arroy::Database<Angular> {
fn angular_db(&self) -> arroy::Database<Cosine> {
self.database.remap_data_type()
}
fn quantized_db(&self) -> arroy::Database<BinaryQuantizedAngular> {
fn quantized_db(&self) -> arroy::Database<BinaryQuantizedCosine> {
self.database.remap_data_type()
}
}