Leave provider change for upstream

Author: Sicheng-Pan

rust/index/src/spann/quantized_spann.rs

@@ -1073,11 +1073,12 @@ impl QuantizedSpannIndexWriter<USearchIndex> {
         usearch_provider: &USearchIndexProvider,
     ) -> Result<(), QuantizedSpannError> {
         // Scrub all clusters that received mutations.
-        let mutated_cluster_ids = self
+        let mut mutated_cluster_ids = self
             .cluster_deltas
             .iter()
             .filter_map(|entry| (!entry.value().ids.is_empty()).then_some(*entry.key()))
             .collect::<Vec<_>>();
+        mutated_cluster_ids.sort_unstable();
 
         for cluster_id in &mutated_cluster_ids {
             self.scrub(*cluster_id).await?;

rust/segment/src/quantized_spann.rs

@@ -543,8 +543,7 @@ mod test {
                 .await
                 .unwrap_or_else(|e| panic!("cycle {cycle}: finish failed: {e}"));
 
-            let flusher = writer
-                .commit()
+            let flusher = Box::pin(writer.commit())
                 .await
                 .unwrap_or_else(|e| panic!("cycle {cycle}: commit failed: {e}"));
 

rust/segment/src/spann_provider.rs

@@ -2,8 +2,6 @@ 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,
@@ -12,28 +10,17 @@ 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(Clone)]
+#[derive(Debug, Clone)]
 pub struct SpannProvider {
-    pub adaptive_search_nprobe: bool,
+    pub hnsw_provider: HnswIndexProvider,
     pub blockfile_provider: BlockfileProvider,
     pub garbage_collection_context: GarbageCollectionContext,
-    pub hnsw_provider: HnswIndexProvider,
     pub metrics: SpannMetrics,
     pub pl_block_size: usize,
-    #[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()
-    }
+    pub adaptive_search_nprobe: bool,
 }
 
-#[cfg(not(feature = "usearch"))]
 #[async_trait]
 impl Configurable<(HnswIndexProvider, BlockfileProvider, SpannProviderConfig)> for SpannProvider {
     async fn try_from_config(
@@ -49,51 +36,12 @@ impl Configurable<(HnswIndexProvider, BlockfileProvider, SpannProviderConfig)> f
         )
         .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,
-        })
-    }
-}
-
-#[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(),
+            adaptive_search_nprobe: config.2.adaptive_search_nprobe,
         })
     }
 }
@@ -119,22 +67,4 @@ 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,35 +57,17 @@ 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,
+            temp_dirs: vec![blockfile_dir, hnsw_dir],
             blockfile_provider: blockfile_provider.clone(),
             hnsw_provider: hnsw_provider.clone(),
             spann_provider: SpannProvider {
-                adaptive_search_nprobe: true,
+                hnsw_provider,
                 blockfile_provider,
                 garbage_collection_context,
-                hnsw_provider,
                 metrics: SpannMetrics::default(),
                 pl_block_size: 5 * 1024 * 1024,
-                #[cfg(feature = "usearch")]
-                usearch_provider,
+                adaptive_search_nprobe: true,
             },
             collection,
             metadata_segment: test_segment(collection_uuid, SegmentScope::METADATA),