Move back to the hnsw crate

This reverts commit 7a4b6c065482f988b01298642f4c18775503f92f.
This commit is contained in:
Kerollmops
2023-06-14 14:20:05 +02:00
committed by Clément Renault
parent aca305bb77
commit c79e82c62a
9 changed files with 101 additions and 60 deletions

View File

@ -28,6 +28,7 @@ use db_cache::DatabaseCache;
use exact_attribute::ExactAttribute;
use graph_based_ranking_rule::{Exactness, Fid, Position, Proximity, Typo};
use heed::RoTxn;
use hnsw::Searcher;
use interner::{DedupInterner, Interner};
pub use logger::visual::VisualSearchLogger;
pub use logger::{DefaultSearchLogger, SearchLogger};
@ -39,7 +40,7 @@ use ranking_rules::{
use resolve_query_graph::{compute_query_graph_docids, PhraseDocIdsCache};
use roaring::RoaringBitmap;
use sort::Sort;
use space::{KnnMap, Neighbor};
use space::Neighbor;
use self::geo_sort::GeoSort;
pub use self::geo_sort::Strategy as GeoSortStrategy;
@ -47,7 +48,9 @@ use self::graph_based_ranking_rule::Words;
use self::interner::Interned;
use crate::score_details::{ScoreDetails, ScoringStrategy};
use crate::search::new::distinct::apply_distinct_rule;
use crate::{AscDesc, DocumentId, Filter, Index, Member, Result, TermsMatchingStrategy, UserError};
use crate::{
AscDesc, DocumentId, Filter, Index, Member, Result, TermsMatchingStrategy, UserError, BEU32,
};
/// A structure used throughout the execution of a search query.
pub struct SearchContext<'ctx> {
@ -447,15 +450,26 @@ pub fn execute_search(
let docids = match vector {
Some(vector) => {
// return the nearest documents that are also part of the candidates.
let hgg = ctx.index.vector_hgg(ctx.txn)?.unwrap_or_default();
hgg.knn_values(&vector, 100)
.into_iter()
.filter(|(Neighbor { distance, .. }, docid)| {
dbg!(distance, f32::from_bits(*distance));
universe.contains(**docid)
})
.map(|(_, docid)| *docid)
.collect()
let mut searcher = Searcher::new();
let hnsw = ctx.index.vector_hnsw(ctx.txn)?.unwrap_or_default();
let ef = hnsw.len().min(100);
let mut dest = vec![Neighbor { index: 0, distance: 0 }; ef];
let neighbors = hnsw.nearest(&vector, ef, &mut searcher, &mut dest[..]);
let mut docids = Vec::new();
for Neighbor { index, distance } in neighbors.iter() {
let index = BEU32::new(*index as u32);
let docid = ctx.index.vector_id_docid.get(ctx.txn, &index)?.unwrap().get();
dbg!(distance, f32::from_bits(*distance));
if universe.contains(docid) {
docids.push(docid);
if docids.len() == length {
break;
}
}
}
docids
}
// return the search docids if the vector field is not specified
None => docids,