Add degraded search (#107)

* add search result enum, factorize level walking

* feat(reader): add cancellation

* add more tests and fix nns_by_item

* add docs

* feat!(reader): make search typed

* bump tests

* update readme

* fix fmt and tests

Author: nnethercott

README.md

@@ -31,38 +31,32 @@ use heed::EnvOpenOptions;
 use rand::{rngs::StdRng, SeedableRng};
 
 fn main() -> Result<()> {
-    const DIM: usize = 3;
-    let vecs: Vec<[f32; DIM]> = vec![[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]];
-
     let env = unsafe {
         EnvOpenOptions::new()
-            .map_size(1024 * 1024 * 1024 * 1) // 1GiB
+            .map_size(1024 * 1024 * 1024) // 1GiB
             .open("./")
     }
     .unwrap();
 
-    let mut wtxn = env.write_txn().unwrap();
+    let mut wtxn = env.write_txn()?;
     let db: Database<Cosine> = env.create_database(&mut wtxn, None)?;
-    let writer: Writer<Cosine> = Writer::new(db, 0, DIM);
+    let writer: Writer<Cosine> = Writer::new(db, 0, 3);
 
-    // insert into lmdb
-    writer.add_item(&mut wtxn, 0, &vecs[0])?;
-    writer.add_item(&mut wtxn, 1, &vecs[1])?;
-    writer.add_item(&mut wtxn, 2, &vecs[2])?;
+    // build
+    writer.add_item(&mut wtxn, 0, &[1.0, 0.0, 0.0])?;
+    writer.add_item(&mut wtxn, 0, &[0.0, 1.0, 0.0])?;
 
-    // ...and build hnsw
     let mut rng = StdRng::seed_from_u64(42);
-
     let mut builder = writer.builder(&mut rng);
     builder.ef_construction(100).build::<16,32>(&mut wtxn)?;
     wtxn.commit()?;
 
-    // search hnsw using a new lmdb read transaction
+    // search
     let rtxn = env.read_txn()?;
     let reader = Reader::<Cosine>::open(&rtxn, 0, db)?;
 
     let query = vec![0.0, 1.0, 0.0];
-    let nns = reader.nns(1).ef_search(10).by_vector(&rtxn, &query)?;
+    let nns = reader.nns(1).ef_search(10).by_vector(&rtxn, &query)?.into_nns();
 
     dbg!("{:?}", &nns);
     Ok(())
@@ -81,7 +75,6 @@ db = hannoy.Database(tmp_dir, Metric.COSINE)
 with db.writer(3, m=4, ef=10) as writer:
     writer.add_item(0, [1.0, 0.0, 0.0])
     writer.add_item(1, [0.0, 1.0, 0.0])
-    writer.add_item(2, [0.0, 0.0, 1.0])
 
 reader = db.reader()
 nns = reader.by_vec([0.0, 1.0, 0.0], n=2)

src/lib.rs

@@ -101,7 +101,7 @@ use key::{Key, Prefix, PrefixCodec};
 use metadata::{Metadata, MetadataCodec};
 use node::{Node, NodeCodec};
 use node_id::{NodeId, NodeMode};
-pub use reader::{QueryBuilder, Reader};
+pub use reader::{QueryBuilder, Reader, Searched};
 pub use roaring::RoaringBitmapCodec;
 pub use writer::{HannoyBuilder, Writer};
 

src/python.rs

@@ -385,7 +385,7 @@ impl PyReader {
             };
         }
 
-        let neighbours = match &self.dyn_reader {
+        let found = match &self.dyn_reader {
             DynReader::Cosine(reader) => hnsw_search!(reader, &query)?,
             DynReader::Euclidean(reader) => hnsw_search!(reader, &query)?,
             DynReader::Manhattan(reader) => hnsw_search!(reader, &query)?,
@@ -394,7 +394,7 @@ impl PyReader {
             DynReader::BqManhattan(reader) => hnsw_search!(reader, &query)?,
             DynReader::Hamming(reader) => hnsw_search!(reader, &query)?,
         };
-        Ok(neighbours)
+        Ok(found.into_nns())
     }
 }