Update provider

Author: Sicheng-Pan

rust/segment/src/spann_provider.rs

@@ -2,6 +2,8 @@ use async_trait::async_trait;
 use chroma_blockstore::provider::BlockfileProvider;
 use chroma_config::{registry::Registry, Configurable};
 use chroma_error::ChromaError;
+#[cfg(feature = "usearch")]
+use chroma_index::usearch::USearchIndexProvider;
 use chroma_index::{
     config::SpannProviderConfig,
     hnsw_provider::HnswIndexProvider,
@@ -10,17 +12,28 @@ use chroma_index::{
 use chroma_types::{Cmek, Collection, Segment};
 
 use crate::distributed_spann::{SpannSegmentWriter, SpannSegmentWriterError};
+#[cfg(feature = "usearch")]
+use crate::quantized_spann::{QuantizedSpannSegmentError, QuantizedSpannSegmentWriter};
 
-#[derive(Debug, Clone)]
+#[derive(Clone)]
 pub struct SpannProvider {
-    pub hnsw_provider: HnswIndexProvider,
+    pub adaptive_search_nprobe: bool,
     pub blockfile_provider: BlockfileProvider,
     pub garbage_collection_context: GarbageCollectionContext,
+    pub hnsw_provider: HnswIndexProvider,
     pub metrics: SpannMetrics,
     pub pl_block_size: usize,
-    pub adaptive_search_nprobe: bool,
+    #[cfg(feature = "usearch")]
+    pub usearch_provider: USearchIndexProvider,
 }
 
+impl std::fmt::Debug for SpannProvider {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("SpannProvider").finish()
+    }
+}
+
+#[cfg(not(feature = "usearch"))]
 #[async_trait]
 impl Configurable<(HnswIndexProvider, BlockfileProvider, SpannProviderConfig)> for SpannProvider {
     async fn try_from_config(
@@ -36,12 +49,51 @@ impl Configurable<(HnswIndexProvider, BlockfileProvider, SpannProviderConfig)> f
         )
         .await?;
         Ok(SpannProvider {
-            hnsw_provider: config.0.clone(),
+            adaptive_search_nprobe: config.2.adaptive_search_nprobe,
             blockfile_provider: config.1.clone(),
             garbage_collection_context,
+            hnsw_provider: config.0.clone(),
             metrics: SpannMetrics::default(),
             pl_block_size: config.2.pl_block_size,
+        })
+    }
+}
+
+#[cfg(feature = "usearch")]
+#[async_trait]
+impl
+    Configurable<(
+        HnswIndexProvider,
+        BlockfileProvider,
+        SpannProviderConfig,
+        USearchIndexProvider,
+    )> for SpannProvider
+{
+    async fn try_from_config(
+        config: &(
+            HnswIndexProvider,
+            BlockfileProvider,
+            SpannProviderConfig,
+            USearchIndexProvider,
+        ),
+        registry: &Registry,
+    ) -> Result<Self, Box<dyn ChromaError>> {
+        let garbage_collection_context = GarbageCollectionContext::try_from_config(
+            &(
+                config.2.pl_garbage_collection.clone(),
+                config.2.hnsw_garbage_collection.clone(),
+            ),
+            registry,
+        )
+        .await?;
+        Ok(SpannProvider {
             adaptive_search_nprobe: config.2.adaptive_search_nprobe,
+            blockfile_provider: config.1.clone(),
+            garbage_collection_context,
+            hnsw_provider: config.0.clone(),
+            metrics: SpannMetrics::default(),
+            pl_block_size: config.2.pl_block_size,
+            usearch_provider: config.3.clone(),
         })
     }
 }
@@ -67,4 +119,22 @@ impl SpannProvider {
         )
         .await
     }
+
+    #[cfg(feature = "usearch")]
+    pub async fn write_quantized_usearch(
+        &self,
+        collection: &Collection,
+        vector_segment: &Segment,
+        record_segment: &Segment,
+    ) -> Result<QuantizedSpannSegmentWriter, QuantizedSpannSegmentError> {
+        QuantizedSpannSegmentWriter::from_segment(
+            self.pl_block_size,
+            collection,
+            vector_segment,
+            record_segment,
+            &self.blockfile_provider,
+            &self.usearch_provider,
+        )
+        .await
+    }
 }

rust/segment/src/test.rs

@@ -57,17 +57,35 @@ impl TestDistributedSegment {
             .await
             .expect("Expected to construct gc context for spann");
 
+        #[allow(unused_mut)]
+        let mut temp_dirs = vec![blockfile_dir, hnsw_dir];
+
+        #[cfg(feature = "usearch")]
+        let (usearch_dir, usearch_provider) = {
+            let (dir, storage) = chroma_storage::test_storage();
+            let provider = chroma_index::usearch::USearchIndexProvider::new(
+                storage,
+                chroma_cache::new_non_persistent_cache_for_test(),
+            );
+            (dir, provider)
+        };
+
+        #[cfg(feature = "usearch")]
+        temp_dirs.push(usearch_dir);
+
         Self {
-            temp_dirs: vec![blockfile_dir, hnsw_dir],
+            temp_dirs,
             blockfile_provider: blockfile_provider.clone(),
             hnsw_provider: hnsw_provider.clone(),
             spann_provider: SpannProvider {
-                hnsw_provider,
+                adaptive_search_nprobe: true,
                 blockfile_provider,
                 garbage_collection_context,
+                hnsw_provider,
                 metrics: SpannMetrics::default(),
                 pl_block_size: 5 * 1024 * 1024,
-                adaptive_search_nprobe: true,
+                #[cfg(feature = "usearch")]
+                usearch_provider,
             },
             collection,
             metadata_segment: test_segment(collection_uuid, SegmentScope::METADATA),

rust/worker/src/execution/orchestration/attached_function_orchestrator.rs

@@ -728,10 +728,7 @@ impl Handler<TaskResult<CollectionAndSegments, GetCollectionAndSegmentsError>>
                 )
                 .await
             {
-                Some(writer) => (
-                    Some(writer.index_uuid()),
-                    VectorSegmentWriter::Hnsw(writer),
-                ),
+                Some(writer) => (Some(writer.index_uuid()), VectorSegmentWriter::Hnsw(writer)),
                 None => return,
             },
         };

rust/worker/src/execution/orchestration/log_fetch_orchestrator.rs

@@ -578,21 +578,13 @@ impl Handler<TaskResult<GetCollectionAndSegmentsOutput, GetCollectionAndSegments
                 .ok_or_terminate(
                     self.context
                         .spann_provider
-                        .write_quantized_usearch(
-                            &collection,
-                            &vector_segment,
-                            &record_segment,
-                        )
+                        .write_quantized_usearch(&collection, &vector_segment, &record_segment)
                         .await,
                     ctx,
                 )
                 .await
             {
-                Some(writer) => (
-                    None,
-                    VectorSegmentWriter::QuantizedSpann(writer),
-                    true,
-                ),
+                Some(writer) => (None, VectorSegmentWriter::QuantizedSpann(writer), true),
                 None => return,
             },
             SegmentType::Spann => match self