Force bftree insert errors to be handled.

Author: Mark Hildebrand

diskann-bftree/.clippy.toml

@@ -0,0 +1,7 @@
+allow-unwrap-in-tests = true
+allow-expect-in-tests = true
+allow-panic-in-tests = true
+
+disallowed-methods = [
+    { path = "bf_tree::BfTree::insert", reason = "This method is fallible but returns a `bool`. Use `crate::bftree_insert` instead" }
+]

diskann-bftree/Cargo.toml

@@ -37,6 +37,11 @@ tokio = { workspace = true, features = ["full"] }
 default = []
 experimental_diversity_search = ["diskann/experimental_diversity_search"]
 
-[lints]
-workspace = true
+[lints.clippy]
+undocumented_unsafe_blocks = "warn"
+unwrap_used = "warn"
+expect_used = "warn"
+panic = "warn"
+disallowed_methods = "warn"
+
 

diskann-bftree/src/lib.rs

@@ -160,3 +160,38 @@ impl Default for TestCallCount {
         Self::new()
     }
 }
+
+/// `bf_tree::BfTree::insert` can fail, but uses a `LeafInsertResult` that is not `#[must_use]`.
+///
+/// This makes it too easy to drop errors.
+///
+/// We use a `clippy` lint to explicitly disallow `bf_tree::BfTree::insert` and instead
+/// funnel calls through this method instead to get a proper error.
+#[expect(
+    clippy::disallowed_methods,
+    reason = "this is the allowed way to call this method"
+)]
+fn bftree_insert(tree: &bf_tree::BfTree, key: &[u8], value: &[u8]) -> Result<(), InsertError> {
+    match tree.insert(key, value) {
+        bf_tree::LeafInsertResult::Success => Ok(()),
+        bf_tree::LeafInsertResult::InvalidKV(s) => Err(InsertError(s)),
+    }
+}
+
+#[derive(Debug)]
+pub struct InsertError(String);
+
+impl std::fmt::Display for InsertError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "insert into a `bftree` failed: {}", self.0)
+    }
+}
+
+impl std::error::Error for InsertError {}
+
+impl From<InsertError> for ANNError {
+    #[track_caller]
+    fn from(error: InsertError) -> Self {
+        ANNError::new(diskann::ANNErrorKind::IndexError, error)
+    }
+}

diskann-bftree/src/neighbors.rs

@@ -17,7 +17,7 @@ use diskann::{
 };
 
 use super::ConfigError;
-use crate::TestCallCount;
+use crate::{bftree_insert, TestCallCount};
 
 pub struct NeighborProvider<I: VectorId> {
     adjacency_list_index: BfTree,
@@ -142,10 +142,10 @@ impl<I: VectorId> NeighborProvider<I> {
     ///
     /// Each key is a `VectorId`, written in bytes. The value is a
     /// neighbor list of length exactly `self.dim`. Specifically,
-    /// an array of exactly `n` neighbors is written as :  
+    /// an array of exactly `n` neighbors is written as :
     /// ```text
     /// |  I0  | ... |  In  | padding |  ...  | [n; 0] |
-    ///     
+    ///
     /// -----------------------------------------------
     ///                     |
     ///                 self.dim
@@ -177,7 +177,7 @@ impl<I: VectorId> NeighborProvider<I> {
         let key = bytemuck::bytes_of(&vector_id);
         let value = cast_slice::<I, u8>(&buf[..self.dim]);
 
-        self.adjacency_list_index.insert(key, value);
+        bftree_insert(&self.adjacency_list_index, key, value)?;
 
         Ok(())
     }
@@ -189,7 +189,7 @@ impl<I: VectorId> NeighborProvider<I> {
     /// bytes of it into the bf-tree.
     ///
     /// # Errors
-    ///  
+    ///
     ///  - Neighbor length is larger than `self.max_degree()`
     ///  - Buffer length (in `I` cells) is smaller than `max_degree() + 1`
     pub fn set_neighbors(&self, vector_id: I, neighbors: &[I], buf: &mut [I]) -> ANNResult<()> {

diskann-bftree/src/provider.rs

@@ -862,10 +862,17 @@ where
     }
 
     fn get_distance(&mut self, id: u32) -> Result<f32, AccessError> {
-        self.provider
+        match self
+            .provider
             .quant_vectors
             .get_vector_into(id.into_usize(), &mut self.element)
-            .map(|_: ()| self.computer.evaluate_similarity(&self.element))
+        {
+            Ok(()) => self
+                .computer
+                .evaluate(&self.element)
+                .map_err(RankedError::Error),
+            Err(err) => Err(err),
+        }
     }
 }
 

diskann-bftree/src/quant.rs

@@ -17,15 +17,16 @@ use diskann_quantization::{
 use diskann_vector::PreprocessedDistanceFunction;
 
 use super::ConfigError;
-use crate::TestCallCount;
+use crate::{bftree_insert, TestCallCount};
 
 pub struct QuantQueryComputer(QueryComputer<GlobalAllocator>);
 
-impl PreprocessedDistanceFunction<&[u8], f32> for QuantQueryComputer {
-    fn evaluate_similarity(&self, x: &[u8]) -> f32 {
-        self.0
-            .evaluate_similarity(Opaque::new(x))
-            .expect("spherical query distance failed")
+impl QuantQueryComputer {
+    pub(crate) fn evaluate(&self, x: &[u8]) -> ANNResult<f32> {
+        match self.0.evaluate_similarity(Opaque::new(x)) {
+            Ok(distance) => Ok(distance),
+            Err(err) => Err(ANNError::new(diskann::ANNErrorKind::IndexError, err)),
+        }
     }
 }
 
@@ -188,7 +189,7 @@ impl QuantVectorProvider {
             )
             .map_err(|e| ANNError::log_sq_error(e))?;
 
-        self.quant_vector_index.insert(key, quant_vector);
+        bftree_insert(&self.quant_vector_index, key, quant_vector)?;
 
         Ok(())
     }
@@ -209,7 +210,7 @@ impl QuantVectorProvider {
         // Update pq vector with id = i to v
         let key = bytemuck::bytes_of(&i);
 
-        self.quant_vector_index.insert(key, v);
+        bftree_insert(&self.quant_vector_index, key, v)?;
 
         Ok(())
     }
@@ -264,7 +265,7 @@ mod tests {
 
     use diskann::ANNErrorKind;
     use diskann_quantization::spherical::iface::Opaque;
-    use diskann_vector::{DistanceFunction, PreprocessedDistanceFunction};
+    use diskann_vector::DistanceFunction;
     use tokio::task::JoinSet;
 
     use super::*;
@@ -331,7 +332,7 @@ mod tests {
 
         // Query Computer — verify it returns finite distances.
         let c = provider.query_computer(&[-0.5f32, -0.5]).unwrap();
-        let dist = c.evaluate_similarity(&provider.get_vector_sync(3).unwrap());
+        let dist = c.evaluate(&provider.get_vector_sync(3).unwrap()).unwrap();
         assert!(dist.is_finite(), "query distance should be finite");
 
         // Distance Computer — verify distances between compressed vectors are finite

diskann-bftree/src/vectors.rs

@@ -14,7 +14,7 @@ use diskann::{error::RankedError, utils::VectorRepr, ANNError, ANNErrorKind, ANN
 use thiserror::Error;
 
 use super::ConfigError;
-use crate::TestCallCount;
+use crate::{bftree_insert, TestCallCount};
 
 pub struct VectorProvider<T: VectorRepr> {
     dim: usize,
@@ -131,7 +131,7 @@ impl<T: VectorRepr> VectorProvider<T> {
         let key = bytemuck::bytes_of(&i);
         let value = cast_slice::<T, u8>(v);
 
-        self.vector_index.insert(key, value);
+        bftree_insert(&self.vector_index, key, value)?;
 
         Ok(())
     }