fix the tests and the arroy_readers method

This commit is contained in:
Tamo 2024-09-19 10:32:17 +02:00
parent cc45e264ca
commit 79f29eed3c
13 changed files with 39 additions and 40 deletions

View file

@ -1612,11 +1612,20 @@ impl Index {
pub fn arroy_readers<'a>(
&'a self,
rtxn: &'a RoTxn<'a>,
embedder_id: u8,
quantized: bool,
) -> impl Iterator<Item = ArroyReader> + 'a {
crate::vector::arroy_db_range_for_embedder(embedder_id)
.map_while(move |k| Some(ArroyReader::new(self.vector_arroy, k, quantized)))
) -> impl Iterator<Item = Result<ArroyReader>> + 'a {
crate::vector::arroy_db_range_for_embedder(embedder_id).map_while(move |k| {
let reader = ArroyReader::new(self.vector_arroy, k, quantized);
// Here we don't care about the dimensions, but we want to know if we can read
// in the database or if its medata are missing.
match reader.dimensions(rtxn) {
Ok(_) => Some(Ok(reader)),
Err(arroy::Error::MissingMetadata(_)) => None,
Err(e) => Some(Err(e.into())),
}
})
}
pub(crate) fn put_search_cutoff(&self, wtxn: &mut RwTxn<'_>, cutoff: u64) -> heed::Result<()> {

View file

@ -52,13 +52,13 @@ impl<Q: RankingRuleQueryTrait> VectorSort<Q> {
ctx: &mut SearchContext<'_>,
vector_candidates: &RoaringBitmap,
) -> Result<()> {
let readers: Vec<_> =
ctx.index.arroy_readers(self.embedder_index, self.quantized).collect();
let readers: Result<Vec<_>> =
ctx.index.arroy_readers(ctx.txn, self.embedder_index, self.quantized).collect();
let target = &self.target;
let mut results = Vec::new();
for reader in readers.iter() {
for reader in readers?.iter() {
let nns_by_vector =
reader.nns_by_vector(ctx.txn, target, self.limit, Some(vector_candidates))?;
results.extend(nns_by_vector.into_iter());

View file

@ -70,7 +70,9 @@ impl<'a> Similar<'a> {
.get(self.rtxn, &self.embedder_name)?
.ok_or_else(|| crate::UserError::InvalidEmbedder(self.embedder_name.to_owned()))?;
let readers: Vec<_> = self.index.arroy_readers(embedder_index, self.quantized).collect();
let readers: Result<Vec<_>> =
self.index.arroy_readers(self.rtxn, embedder_index, self.quantized).collect();
let readers = readers?;
let mut results = Vec::new();

View file

@ -997,9 +997,14 @@ impl<'a, 'i> Transform<'a, 'i> {
if let Some(WriteBackToDocuments { embedder_id, user_provided }) =
action.write_back()
{
let readers: Vec<_> =
self.index.arroy_readers(*embedder_id, action.was_quantized).collect();
Some(Ok((name.as_str(), (readers, user_provided))))
let readers: Result<Vec<_>> = self
.index
.arroy_readers(wtxn, *embedder_id, action.was_quantized)
.collect();
match readers {
Ok(readers) => Some(Ok((name.as_str(), (readers, user_provided)))),
Err(error) => Some(Err(error)),
}
} else {
None
}