[ENH]: Load HNSW index without disk intermediary

Author: tanujnay112

Cargo.lock

@@ -92,7 +92,7 @@ bytemuck = "1.21.0"
 rayon = "1.10.0"
 validator = { version = "0.19", features = ["derive"] }
 rust-embed = { version = "8.5.0", features = ["include-exclude", "debug-embed"] }
-hnswlib = { version = "0.8.1", git = "https://github.com/chroma-core/hnswlib.git" }
+hnswlib = { git = "https://github.com/chroma-core/hnswlib.git", rev = "736bfd15d05843e6cc6ac3b806edd21566821be" }
 reqwest = { version = "0.12.9", features = ["rustls-tls-native-roots", "http2"], default-features = false }
 random-port = "0.1.1"
 ndarray = { version = "0.16.1", features = ["approx"] }

Cargo.toml

@@ -230,6 +230,29 @@ impl PersistentIndex<HnswIndexConfig> for HnswIndex {
             dimensionality: index_config.dimensionality,
             persist_path: path.into(),
             ef_search,
+            hnsw_data: hnswlib::HnswData::default(),
+        })
+        .map_err(|e| WrappedHnswInitError::Other(e).boxed())?;
+
+        Ok(HnswIndex {
+            index,
+            id,
+            distance_function: index_config.distance_function.clone(),
+        })
+    }
+
+    fn load_from_hnsw_data(
+        hnsw_data: hnswlib::HnswData,
+        index_config: &IndexConfig,
+        ef_search: usize,
+        id: IndexUuid,
+    ) -> Result<Self, Box<dyn ChromaError>> {
+        let index = hnswlib::HnswIndex::load_from_hnsw_data(hnswlib::HnswIndexLoadConfig {
+            distance_function: map_distance_function(index_config.distance_function.clone()),
+            dimensionality: index_config.dimensionality,
+            persist_path: "".into(),
+            ef_search,
+            hnsw_data: hnsw_data,
         })
         .map_err(|e| WrappedHnswInitError::Other(e).boxed())?;
 

rust/index/src/hnsw.rs

@@ -207,20 +207,27 @@ impl HnswIndexProvider {
 
         let index_config = IndexConfig::new(dimensionality, distance_function);
 
-        let storage_path_str = match new_storage_path.to_str() {
-            Some(storage_path_str) => storage_path_str,
-            None => {
-                return Err(Box::new(HnswIndexProviderForkError::PathToStringError(
-                    new_storage_path,
-                )));
-            }
-        };
+        // let storage_path_str = match new_storage_path.to_str() {
+        //     Some(storage_path_str) => storage_path_str,
+        //     None => {
+        //         return Err(Box::new(HnswIndexProviderForkError::PathToStringError(
+        //             new_storage_path,
+        //         )));
+        //     }
+        // };
 
         // Check if the entry is in the cache, if it is, we assume
         // another thread has loaded the index and we return it.
         match self.get(&new_id, cache_key).await {
             Some(index) => Ok(index.clone()),
-            None => match HnswIndex::load(storage_path_str, &index_config, ef_search, new_id) {
+            None => match HnswIndex::load_from_hnsw_data(
+                self.fetch_hnsw_segment(&new_id, prefix_path)
+                    .await
+                    .map_err(|e| Box::new(HnswIndexProviderForkError::FileError(*e)))?,
+                &index_config,
+                ef_search,
+                new_id,
+            ) {
                 Ok(index) => {
                     let index = HnswIndexRef {
                         inner: Arc::new(RwLock::new(DistributedHnswInner {
@@ -277,10 +284,33 @@ impl HnswIndexProvider {
         prefix_path: &str,
     ) -> Result<(), Box<HnswIndexProviderFileError>> {
         // Fetch the files from storage and put them in the index storage path.
+        let hnsw_data = self.fetch_hnsw_segment(source_id, prefix_path).await?;
+        let getters = [
+            |hnsw_data: &hnswlib::HnswData| Arc::new(Vec::from(hnsw_data.header_buffer())),
+            |hnsw_data: &hnswlib::HnswData| Arc::new(Vec::from(hnsw_data.data_level0_buffer())),
+            |hnsw_data: &hnswlib::HnswData| Arc::new(Vec::from(hnsw_data.length_buffer())),
+            |hnsw_data: &hnswlib::HnswData| Arc::new(Vec::from(hnsw_data.link_list_buffer())),
+        ];
+
+        for (file, getter) in FILES.iter().zip(getters) {
+            let file_path = index_storage_path.join(file);
+            self.copy_bytes_to_local_file(&file_path, getter(&hnsw_data))
+                .await?;
+        }
+        Ok(())
+    }
+
+    async fn fetch_hnsw_segment(
+        &self,
+        source_id: &IndexUuid,
+        prefix_path: &str,
+    ) -> Result<hnswlib::HnswData, Box<HnswIndexProviderFileError>> {
+        let mut buffers = Vec::new();
+
         for file in FILES.iter() {
             let s3_fetch_span =
                 tracing::trace_span!(parent: Span::current(), "Read bytes from s3", file = file);
-            let buf = s3_fetch_span
+            let _ = s3_fetch_span
                 .in_scope(|| async {
                     let key = Self::format_key(prefix_path, source_id, file);
                     tracing::info!("Loading hnsw index file: {} into directory", key);
@@ -304,13 +334,24 @@ impl HnswIndexProvider {
                         bytes_read,
                         key,
                     );
-                    Ok(buf)
+                    buffers.push(buf);
+                    Ok(())
                 })
                 .await?;
-            let file_path = index_storage_path.join(file);
-            self.copy_bytes_to_local_file(&file_path, buf).await?;
         }
-        Ok(())
+        match hnswlib::HnswData::new_from_buffers(
+            buffers[0].clone(),
+            buffers[1].clone(),
+            buffers[2].clone(),
+            buffers[3].clone(),
+        ) {
+            Ok(hnsw_data) => Ok(hnsw_data),
+            Err(e) => Err(Box::new(HnswIndexProviderFileError::StorageError(
+                chroma_storage::StorageError::Message {
+                    message: e.to_string(),
+                },
+            ))),
+        }
     }
 
     pub async fn open(
@@ -356,20 +397,27 @@ impl HnswIndexProvider {
 
         let index_config = IndexConfig::new(dimensionality, distance_function);
 
-        let index_storage_path_str = match index_storage_path.to_str() {
-            Some(index_storage_path_str) => index_storage_path_str,
-            None => {
-                return Err(Box::new(HnswIndexProviderOpenError::PathToStringError(
-                    index_storage_path,
-                )));
-            }
-        };
+        // let index_storage_path_str = match index_storage_path.to_str() {
+        //     Some(index_storage_path_str) => index_storage_path_str,
+        //     None => {
+        //         return Err(Box::new(HnswIndexProviderOpenError::PathToStringError(
+        //             index_storage_path,
+        //         )));
+        //     }
+        // };
 
         // Check if the entry is in the cache, if it is, we assume
         // another thread has loaded the index and we return it.
         let index = match self.get(id, cache_key).await {
             Some(index) => Ok(index.clone()),
-            None => match HnswIndex::load(index_storage_path_str, &index_config, ef_search, *id) {
+            None => match HnswIndex::load_from_hnsw_data(
+                self.fetch_hnsw_segment(id, prefix_path)
+                    .await
+                    .map_err(|e| Box::new(HnswIndexProviderOpenError::FileError(*e)))?,
+                &index_config,
+                ef_search,
+                *id,
+            ) {
                 Ok(index) => {
                     let index = HnswIndexRef {
                         inner: Arc::new(RwLock::new(DistributedHnswInner {

rust/index/src/hnsw_provider.rs

@@ -67,6 +67,15 @@ pub trait PersistentIndex<C>: Index<C> {
     ) -> Result<Self, Box<dyn ChromaError>>
     where
         Self: Sized;
+
+    fn load_from_hnsw_data(
+        hnsw_data: hnswlib::HnswData,
+        index_config: &IndexConfig,
+        ef_search: usize,
+        id: IndexUuid,
+    ) -> Result<Self, Box<dyn ChromaError>>
+    where
+        Self: Sized;
 }
 
 /// IndexUuid is a wrapper around Uuid to provide a type for the index id.