From 514bda02ccb8e1588a2e1fd20cddb832de28d9b1 Mon Sep 17 00:00:00 2001
From: Evan Culver <eculver@users.noreply.github.com>
Date: Tue, 7 Jan 2025 14:52:07 -0800
Subject: [PATCH] [BUG] Fix HNSW resize conditions to exclude deleted items
 (#3424)

## Description of changes

We were seeing errors in compaction that looked like:

```
WriteSegments(ApplyMaterializatedLogsError(HnswIndex(FFIException(\"The number of elements exceeds the specified limit\"))))
```
This is because the resize logic uses `len` which removes deleted
elements which then causes the conditional for resizing to be skipped.
This adds a new method `len_with_deleted` to return the length of the
index including the elements that have been deleted so that the resize
is not skipped. See:
https://github.com/chroma-core/chroma/blob/27564afa0e82d184a0fa853a7384d00861d8e9dc/rust/worker/src/segment/distributed_hnsw_segment.rs#L212-L222

## Test plan
A test was added. It will fail when using the existing `len` method, but
passes when using the newly added `len_with_deleted` method.

- [x] Tests pass locally with `pytest` for python, `yarn test` for js,
`cargo test` for rust

## Documentation Changes
N/A

Thanks to @sanketkedia for helping root cause this.

Fixes https://github.com/chroma-core/hosted-chroma/issues/855
---
 rust/index/bindings.cpp | 13 ++++++++++
 rust/index/src/hnsw.rs  | 56 ++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 68 insertions(+), 1 deletion(-)

diff --git a/rust/index/bindings.cpp b/rust/index/bindings.cpp
index a495c0f9a64..f8e4b00aeea 100644
--- a/rust/index/bindings.cpp
+++ b/rust/index/bindings.cpp
@@ -430,6 +430,7 @@ extern "C"
     }
 
     // Can not throw std::exception
+    // Note: includes deleted items
     int len(Index<float> *index)
     {
         if (!index->index_inited)
@@ -440,6 +441,18 @@ extern "C"
         return index->appr_alg->getCurrentElementCount() - index->appr_alg->getDeletedCount();
     }
 
+    // Can not throw std::exception
+    // Note: does not include deleted items
+    int len_with_deleted(Index<float> *index)
+    {
+        if (!index->index_inited)
+        {
+            return 0;
+        }
+
+        return index->appr_alg->getCurrentElementCount();
+    }
+
     // Can not throw std::exception
     size_t capacity(Index<float> *index)
     {
diff --git a/rust/index/src/hnsw.rs b/rust/index/src/hnsw.rs
index a687e80cecc..90105def396 100644
--- a/rust/index/src/hnsw.rs
+++ b/rust/index/src/hnsw.rs
@@ -308,6 +308,11 @@ impl HnswIndex {
         // Does not return an error
     }
 
+    pub fn len_with_deleted(&self) -> usize {
+        unsafe { len_with_deleted(self.ffi_ptr) as usize }
+        // Does not return an error
+    }
+
     pub fn is_empty(&self) -> bool {
         self.len() == 0
     }
@@ -397,9 +402,9 @@ extern "C" {
 
     #[cfg(test)]
     fn get_ef(index: *const IndexPtrFFI) -> c_int;
-
     fn set_ef(index: *const IndexPtrFFI, ef: c_int);
     fn len(index: *const IndexPtrFFI) -> c_int;
+    fn len_with_deleted(index: *const IndexPtrFFI) -> c_int;
     fn capacity(index: *const IndexPtrFFI) -> c_int;
     fn resize_index(index: *const IndexPtrFFI, new_size: usize);
     fn get_last_error(index: *const IndexPtrFFI) -> *const c_char;
@@ -925,4 +930,53 @@ pub mod test {
             .to_string()
             .contains("HNSW Integrity failure"))
     }
+
+    #[test]
+    fn it_can_resize_correctly() {
+        let n: usize = 10;
+        let d: usize = 960;
+        let distance_function = DistanceFunction::Euclidean;
+        let tmp_dir = tempdir().unwrap();
+        let persist_path = tmp_dir.path().to_str().unwrap().to_string();
+        let id = Uuid::new_v4();
+        let index = HnswIndex::init(
+            &IndexConfig {
+                dimensionality: d as i32,
+                distance_function: distance_function.clone(),
+            },
+            Some(&HnswIndexConfig {
+                max_elements: n,
+                m: 32,
+                ef_construction: 100,
+                ef_search: 100,
+                random_seed: 0,
+                persist_path: persist_path.clone(),
+            }),
+            IndexUuid(id),
+        );
+
+        let mut index = match index {
+            Err(e) => panic!("Error initializing index: {}", e),
+            Ok(index) => index,
+        };
+
+        let data: Vec<f32> = utils::generate_random_data(n, d);
+        let ids: Vec<usize> = (0..n).collect();
+
+        (0..n).for_each(|i| {
+            let data = &data[i * d..(i + 1) * d];
+            index.add(ids[i], data).expect("Should not error");
+        });
+
+        index.delete(0).unwrap();
+        let data = &data[d..2 * d];
+
+        let index_len = index.len_with_deleted();
+        let index_capacity = index.capacity();
+        if index_len + 1 > index_capacity {
+            index.resize(index_capacity * 2).unwrap();
+        }
+        // this will fail if the index is not resized correctly
+        index.add(100, data).unwrap();
+    }
 }