[ENH] Wire up quantized writer in compaction

clients/new-js/packages/chromadb/src/api/types.gen.ts

@@ -436,6 +436,11 @@ export type Key = 'Document' | 'Embedding' | 'Metadata' | 'Score' | {
     MetadataField: string;
 };
 
+/**
+ * Quantization implementation for SPANN vector index.
+ */
+export type Quantization = 'none' | 'u_search4_bit_rabit_q';
+
 export type QueryRequestPayload = RawWhereFields & {
     ids?: Array<string> | null;
     include?: IncludeList;
@@ -559,9 +564,9 @@ export type SpannIndexConfig = {
     num_centers_to_merge_to?: number | null;
     num_samples_kmeans?: number | null;
     /**
-     * Enable quantization for vector search (cloud-only feature)
+     * Quantization implementation for vector search (cloud-only feature)
      */
-    quantize?: boolean;
+    quantize?: Quantization;
     reassign_neighbor_count?: number | null;
     search_nprobe?: number | null;
     search_rng_epsilon?: number | null;

rust/frontend/src/impls/service_based_frontend.rs

@@ -37,9 +37,9 @@ use chroma_types::{
     GetTenantRequest, GetTenantResponse, HealthCheckResponse, HeartbeatError, Include,
     IndexStatusError, IndexStatusResponse, KnnIndex, ListCollectionsRequest,
     ListCollectionsResponse, ListDatabasesError, ListDatabasesRequest, ListDatabasesResponse,
-    Operation, OperationRecord, QueryError, QueryRequest, QueryResponse, ResetError, ResetResponse,
-    Schema, SchemaError, SearchRequest, SearchResponse, Segment, SegmentScope, SegmentType,
-    SegmentUuid, UpdateCollectionError, UpdateCollectionRecordsError,
+    Operation, OperationRecord, Quantization, QueryError, QueryRequest, QueryResponse, ResetError,
+    ResetResponse, Schema, SchemaError, SearchRequest, SearchResponse, Segment, SegmentScope,
+    SegmentType, SegmentUuid, UpdateCollectionError, UpdateCollectionRecordsError,
     UpdateCollectionRecordsRequest, UpdateCollectionRecordsResponse, UpdateCollectionRequest,
     UpdateCollectionResponse, UpdateTenantError, UpdateTenantRequest, UpdateTenantResponse,
     UpsertCollectionRecordsError, UpsertCollectionRecordsRequest, UpsertCollectionRecordsResponse,
@@ -548,8 +548,8 @@ impl ServiceBasedFrontend {
 
         // Enable quantization for tenants in the config list (or all tenants if "*" is present)
         if let Some(ref mut schema) = reconciled_schema {
-            if let Some(spann_config) = schema.get_spann_config_mut() {
-                spann_config.quantize = self.should_enable_quantization_for_tenant(&tenant_id);
+            if self.should_enable_quantization_for_tenant(&tenant_id) {
+                schema.quantize(Quantization::USearch4BitRabitQ);
             }
         }
 

rust/index/src/spann/quantized_spann.rs

@@ -1456,7 +1456,7 @@ mod tests {
     use chroma_cache::{new_cache_for_test, new_non_persistent_cache_for_test};
     use chroma_distance::DistanceFunction;
     use chroma_storage::{local::LocalStorage, Storage};
-    use chroma_types::{CollectionUuid, DataRecord, SpannIndexConfig};
+    use chroma_types::{CollectionUuid, DataRecord, Quantization, SpannIndexConfig};
     use rand::{Rng, SeedableRng};
     use tempfile::TempDir;
 
@@ -1491,7 +1491,7 @@ mod tests {
             initial_lambda: None,
             num_centers_to_merge_to: None,
             max_neighbors: Some(8),
-            quantize: true,
+            quantize: Quantization::USearch4BitRabitQ,
         }
     }
 

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),