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

5
Cargo.lock generated
View File

@ -386,8 +386,9 @@ checksum = "96d30a06541fbafbc7f82ed10c06164cfbd2c401138f6addd8404629c4b16711"
[[package]]
name = "arroy"
version = "0.4.0"
source = "git+https://github.com/meilisearch/arroy/?rev=2386594dfb009ce08821a925ccc89fb8e30bf73d#2386594dfb009ce08821a925ccc89fb8e30bf73d"
version = "0.5.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "dfc5f272f38fa063bbff0a7ab5219404e221493de005e2b4078c62d626ef567e"
dependencies = [
"bytemuck",
"byteorder",

View File

@ -40,7 +40,7 @@ ureq = "2.10.0"
uuid = { version = "1.10.0", features = ["serde", "v4"] }
[dev-dependencies]
arroy = { git = "https://github.com/meilisearch/arroy/", rev = "2386594dfb009ce08821a925ccc89fb8e30bf73d" }
arroy = "0.5.0"
big_s = "1.0.2"
crossbeam = "0.8.4"
insta = { version = "1.39.0", features = ["json", "redactions"] }

View File

@ -80,7 +80,7 @@ hf-hub = { git = "https://github.com/dureuill/hf-hub.git", branch = "rust_tls",
tiktoken-rs = "0.5.9"
liquid = "0.26.6"
rhai = { git = "https://github.com/rhaiscript/rhai", rev = "ef3df63121d27aacd838f366f2b83fd65f20a1e4", features = ["serde", "no_module", "no_custom_syntax", "no_time", "sync"] }
arroy = { git = "https://github.com/meilisearch/arroy/", rev = "2386594dfb009ce08821a925ccc89fb8e30bf73d" }
arroy = "0.5.0"
rand = "0.8.5"
tracing = "0.1.40"
ureq = { version = "2.10.0", features = ["json"] }

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()
}
}