Merge pull request #77 from meilisearch/panic-at-search

Panic at search when the trees need a rebuild

Author: irevoire

Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "arroy"
 description = "Annoy-inspired Approximate Nearest Neighbors in Rust, based on LMDB and optimized for memory usage"
-version = "0.3.1"
+version = "0.4.0"
 documentation = "https://docs.rs/arroy"
 repository = "https://github.com/meilisearch/arroy"
 keywords = ["ANN-search", "Graph-algorithms", "Vector-Search", "Store"]

src/error.rs

@@ -1,5 +1,7 @@
 use std::io;
 
+use crate::{key::Key, node_id::NodeMode, ItemId};
+
 /// The different set of errors that arroy can encounter.
 #[derive(Debug, thiserror::Error)]
 pub enum Error {
@@ -12,7 +14,7 @@ pub enum Error {
     Io(#[from] io::Error),
 
     /// The user is trying to insert or search for a vector that is not of the right dimensions.
-    #[error("Invalid vector dimensions. Got {received} but expected {expected}.")]
+    #[error("Invalid vector dimensions. Got {received} but expected {expected}")]
     InvalidVecDimension {
         /// The expected number of dimensions.
         expected: usize,
@@ -21,7 +23,7 @@ pub enum Error {
     },
 
     /// An internal error returned when arroy cannot generate internal IDs.
-    #[error("Database full. Arroy cannot generate enough internal IDs for your items.")]
+    #[error("Database full. Arroy cannot generate enough internal IDs for your items")]
     DatabaseFull,
 
     /// The user tried to append an item in the database but the last inserted item
@@ -30,7 +32,7 @@ pub enum Error {
     InvalidItemAppend,
 
     /// The user is trying to query a database with a distance that is not of the right type.
-    #[error("Invalid distance provided. Got {received} but expected {expected}.")]
+    #[error("Invalid distance provided. Got {received} but expected {expected}")]
     UnmatchingDistance {
         /// The expected distance type.
         expected: String,
@@ -40,10 +42,37 @@ pub enum Error {
 
     /// Arroy is not able to find the metadata for a given index.
     /// It is probably because the user forget to build the database.
-    #[error("Metadata are missing, did you build your database before trying to read it.")]
-    MissingMetadata,
+    #[error(
+        "Metadata are missing on index {0}, You must build your database before attempting to read it"
+    )]
+    MissingMetadata(u16),
+
+    /// The last time items in the database were updated, the [`Writer::build`] method wasn't called.
+    #[error("The trees have not been built after an update on index {0}")]
+    NeedBuild(u16),
 
     /// Internal error
-    #[error("Internal error: Node is missing")]
-    MissingNode,
+    #[error("Internal error: {mode}({item}) is missing in index `{index}`")]
+    MissingKey {
+        /// The index that caused the error
+        index: u16,
+        /// The kind of item that was being queried
+        mode: &'static str,
+        /// The item ID queried
+        item: ItemId,
+    },
+}
+
+impl Error {
+    pub(crate) fn missing_key(key: Key) -> Self {
+        Self::MissingKey {
+            index: key.index,
+            mode: match key.node.mode {
+                NodeMode::Item => "Item",
+                NodeMode::Tree => "Tree",
+                NodeMode::Metadata => "Metadata",
+            },
+            item: key.node.item,
+        }
+    }
 }

src/reader.rs

@@ -5,7 +5,7 @@ use std::iter::repeat;
 use std::marker;
 use std::num::NonZeroUsize;
 
-use heed::types::DecodeIgnore;
+use heed::types::{Bytes, DecodeIgnore};
 use heed::RoTxn;
 use ordered_float::OrderedFloat;
 use roaring::RoaringBitmap;
@@ -36,7 +36,7 @@ impl<'t, D: Distance> Reader<'t, D> {
         let metadata_key = Key::metadata(index);
         let metadata = match database.remap_data_type::<MetadataCodec>().get(rtxn, &metadata_key)? {
             Some(metadata) => metadata,
-            None => return Err(Error::MissingMetadata),
+            None => return Err(Error::MissingMetadata(index)),
         };
 
         if D::name() != metadata.distance {
@@ -45,6 +45,9 @@ impl<'t, D: Distance> Reader<'t, D> {
                 received: D::name(),
             });
         }
+        if database.remap_data_type::<Bytes>().get(rtxn, &Key::updated(index))?.is_some() {
+            return Err(Error::NeedBuild(index));
+        }
 
         Ok(Reader {
             database: database.remap_data_type(),
@@ -222,7 +225,8 @@ impl<'t, D: Distance> Reader<'t, D> {
                 None => break,
             };
 
-            match self.database.get(rtxn, &Key::new(self.index, item))?.unwrap() {
+            let key = Key::new(self.index, item);
+            match self.database.get(rtxn, &key)?.ok_or(Error::missing_key(key))? {
                 Node::Leaf(_) => {
                     if candidates.map_or(true, |c| c.contains(item.item)) {
                         nns.push(item.unwrap_item());
@@ -250,7 +254,8 @@ impl<'t, D: Distance> Reader<'t, D> {
 
         let mut nns_distances = Vec::with_capacity(nns.len());
         for nn in nns {
-            let leaf = match self.database.get(rtxn, &Key::item(self.index, nn))?.unwrap() {
+            let key = Key::item(self.index, nn);
+            let leaf = match self.database.get(rtxn, &key)?.ok_or(Error::missing_key(key))? {
                 Node::Leaf(leaf) => leaf,
                 Node::Descendants(_) | Node::SplitPlaneNormal(_) => unreachable!(),
             };

src/tests/reader.rs

@@ -34,7 +34,7 @@ fn open_unfinished_db() {
 
     let rtxn = handle.env.read_txn().unwrap();
     let ret = Reader::<Euclidean>::open(&rtxn, 0, handle.database).map(|_| ()).unwrap_err();
-    insta::assert_display_snapshot!(ret, @"Metadata are missing, did you build your database before trying to read it.");
+    insta::assert_display_snapshot!(ret, @"Metadata are missing on index 0, You must build your database before attempting to read it");
 }
 
 #[test]
@@ -50,7 +50,7 @@ fn open_db_with_wrong_dimension() {
     let rtxn = handle.env.read_txn().unwrap();
     let reader = Reader::<Euclidean>::open(&rtxn, 0, handle.database).unwrap();
     let ret = reader.nns_by_vector(&rtxn, &[1.0, 2.0, 3.0], 5, None, None).unwrap_err();
-    insta::assert_display_snapshot!(ret, @"Invalid vector dimensions. Got 3 but expected 2.");
+    insta::assert_display_snapshot!(ret, @"Invalid vector dimensions. Got 3 but expected 2");
 }
 
 #[test]
@@ -206,7 +206,7 @@ fn filtering() {
 
 #[test]
 fn search_in_empty_database() {
-    // See https://github.com/meilisearch/arroy/issues/74
+    // See https://github.com/meilisearch/arroy/issues/75
     let handle = create_database::<Euclidean>();
 
     let mut wtxn = handle.env.write_txn().unwrap();
@@ -219,3 +219,41 @@ fn search_in_empty_database() {
     let ret = reader.nns_by_vector(&rtxn, &[0., 0.], 10, None, None).unwrap();
     insta::assert_debug_snapshot!(ret, @"[]");
 }
+
+#[test]
+fn try_reading_in_a_non_built_database() {
+    // See https://github.com/meilisearch/arroy/issues/74
+    let handle = create_database::<Euclidean>();
+
+    let mut wtxn = handle.env.write_txn().unwrap();
+    let writer = Writer::new(handle.database, 0, 2);
+    writer.add_item(&mut wtxn, 0, &[0.0, 0.0]).unwrap();
+    // We don't build the database
+    wtxn.commit().unwrap();
+
+    let rtxn = handle.env.read_txn().unwrap();
+    let error = Reader::open(&rtxn, 0, handle.database).unwrap_err();
+    insta::assert_debug_snapshot!(error, @r###"
+    MissingMetadata(
+        0,
+    )
+    "###);
+    drop(rtxn);
+
+    // we build the database once to get valid metadata
+    let mut wtxn = handle.env.write_txn().unwrap();
+    let writer = Writer::new(handle.database, 0, 2);
+    writer.build(&mut wtxn, &mut rng(), None).unwrap();
+    let writer = Writer::new(handle.database, 0, 2);
+    writer.del_item(&mut wtxn, 0).unwrap();
+    // We don't build the database; this leaves the database in a corrupted state
+    wtxn.commit().unwrap();
+
+    let rtxn = handle.env.read_txn().unwrap();
+    let error = Reader::open(&rtxn, 0, handle.database).unwrap_err();
+    insta::assert_debug_snapshot!(error, @r###"
+    NeedBuild(
+        0,
+    )
+    "###);
+}

src/tests/snapshots/arroy__tests__writer__write_and_update_lot_of_random_points-2.snap

@@ -213,5 +213,4 @@ Tree 124: Descendants(Descendants { descendants: [14, 16, 22, 38, 41, 42, 43, 44
 Tree 125: Descendants(Descendants { descendants: [2, 10, 21, 23, 26, 28, 29, 30, 34, 37, 39, 47, 52, 54, 55, 67, 68, 69, 71, 72, 76, 78, 86, 88, 90, 94, 97, 98] })
 Tree 126: SplitPlaneNormal(SplitPlaneNormal { left: Tree(124), right: Tree(125), normal: [-0.2156, 0.0925, 0.0242, 0.1246, -0.0758, -0.0672, -0.1816, 0.4853, 0.1643, 0.1418, 0.2752, -0.1825, -0.2671, 0.0674, -0.1950, -0.0743, -0.1493, -0.0368, -0.1929, 0.2540, -0.0441, 0.3135, -0.2219, 0.1059, 0.2115, 0.0749, 0.1629, -0.1050, -0.0211, 0.0051] })
 Root: Metadata { dimensions: 30, items: RoaringBitmap<100 values between 0 and 99>, roots: [8, 17, 24, 35, 44, 55, 64, 75, 86, 97], distance: "euclidean" }
-updated_item_ids: RoaringBitmap<[]>
 

src/tests/snapshots/arroy__tests__writer__write_and_update_lot_of_random_points.snap

@@ -203,5 +203,4 @@ Tree 95: SplitPlaneNormal(SplitPlaneNormal { left: Tree(89), right: Tree(94), no
 Tree 96: SplitPlaneNormal(SplitPlaneNormal { left: Tree(88), right: Tree(95), normal: [0.0684, -0.2434, -0.1543, 0.0337, 0.1784, 0.1295, 0.3061, -0.0175, -0.0679, -0.1419, -0.0180, 0.2335, 0.3782, -0.1127, 0.2746, -0.2657, -0.0579, 0.1808, 0.2623, -0.3329, 0.1801, -0.2106, 0.1692, -0.1191, 0.0152, 0.1214, 0.0252, -0.0236, -0.1220, -0.1448] })
 Tree 97: SplitPlaneNormal(SplitPlaneNormal { left: Tree(87), right: Tree(96), normal: [-0.0316, -0.1036, 0.3252, -0.0985, -0.1286, -0.4867, -0.0286, 0.1159, -0.0177, 0.2205, 0.0177, 0.3192, 0.1521, -0.1484, 0.1918, -0.2794, -0.1141, 0.2598, 0.1453, 0.1133, -0.1149, -0.0455, 0.0697, -0.2537, 0.1797, -0.0423, -0.0470, 0.0886, 0.0868, 0.2083] })
 Root: Metadata { dimensions: 30, items: RoaringBitmap<100 values between 0 and 99>, roots: [8, 17, 24, 35, 44, 55, 64, 75, 86, 97], distance: "euclidean" }
-updated_item_ids: RoaringBitmap<[]>