[CLN] Use InternalUpdateConfiguration in Rust, cleanup go code

jairad26 · jairad26 · commit c561c113fbb1 · 2025-07-10T12:54:09.000-07:00
diff --git a/go/pkg/sysdb/coordinator/table_catalog.go b/go/pkg/sysdb/coordinator/table_catalog.go
@@ -815,8 +815,8 @@ func (tc *Catalog) updateCollectionConfiguration(
 
 	// Update existing configuration with new values
 	if updateConfig.VectorIndex != nil {
-		if updateConfig.VectorIndex.Type == "hnsw" && updateConfig.VectorIndex.Hnsw != nil {
-			if existingConfig.VectorIndex == nil || existingConfig.VectorIndex.Type != "hnsw" {
+		if updateConfig.VectorIndex.Hnsw != nil {
+			if existingConfig.VectorIndex == nil || existingConfig.VectorIndex.Hnsw == nil {
 				return existingConfigJsonStr, nil
 			}
 			if updateConfig.VectorIndex.Hnsw.EfSearch != nil {
@@ -837,8 +837,8 @@ func (tc *Catalog) updateCollectionConfiguration(
 			if updateConfig.VectorIndex.Hnsw.BatchSize != nil {
 				existingConfig.VectorIndex.Hnsw.BatchSize = *updateConfig.VectorIndex.Hnsw.BatchSize
 			}
-		} else if updateConfig.VectorIndex.Type == "spann" && updateConfig.VectorIndex.Spann != nil {
-			if existingConfig.VectorIndex == nil || existingConfig.VectorIndex.Type != "spann" {
+		} else if updateConfig.VectorIndex.Spann != nil {
+			if existingConfig.VectorIndex == nil || existingConfig.VectorIndex.Spann == nil {
 				return existingConfigJsonStr, nil
 			}
 			if updateConfig.VectorIndex.Spann.EfSearch != nil {
diff --git a/rust/frontend/src/server.rs b/rust/frontend/src/server.rs
@@ -18,10 +18,10 @@ use chroma_types::{
     DeleteCollectionRecordsResponse, DeleteDatabaseRequest, DeleteDatabaseResponse,
     GetCollectionRequest, GetDatabaseRequest, GetDatabaseResponse, GetRequest, GetResponse,
     GetTenantRequest, GetTenantResponse, GetUserIdentityResponse, HeartbeatResponse, IncludeList,
-    InternalCollectionConfiguration, ListCollectionsRequest, ListCollectionsResponse,
-    ListDatabasesRequest, ListDatabasesResponse, Metadata, QueryRequest, QueryResponse,
-    UpdateCollectionConfiguration, UpdateCollectionRecordsResponse, UpdateCollectionResponse,
-    UpdateMetadata, UpsertCollectionRecordsResponse,
+    InternalCollectionConfiguration, InternalUpdateCollectionConfiguration, ListCollectionsRequest,
+    ListCollectionsResponse, ListDatabasesRequest, ListDatabasesResponse, Metadata, QueryRequest,
+    QueryResponse, UpdateCollectionConfiguration, UpdateCollectionRecordsResponse,
+    UpdateCollectionResponse, UpdateMetadata, UpsertCollectionRecordsResponse,
 };
 use chroma_types::{ForkCollectionResponse, RawWhereFields};
 use mdac::{Rule, Scorecard, ScorecardTicket};
@@ -1091,13 +1091,18 @@ async fn update_collection(
     let collection_id =
         CollectionUuid::from_str(&collection_id).map_err(|_| ValidationError::CollectionId)?;
 
+    let configuration = match payload.new_configuration {
+        Some(c) => Some(InternalUpdateCollectionConfiguration::try_from(c)?),
+        None => None,
+    };
+
     let request = chroma_types::UpdateCollectionRequest::try_new(
         collection_id,
         payload.new_name,
         payload
             .new_metadata
             .map(CollectionMetadataUpdate::UpdateMetadata),
-        payload.new_configuration,
+        configuration,
     )?;
 
     server.frontend.update_collection(request).await?;
diff --git a/rust/python_bindings/src/bindings.rs b/rust/python_bindings/src/bindings.rs
@@ -19,9 +19,9 @@ use chroma_types::{
     CountResponse, CreateCollectionRequest, CreateDatabaseRequest, CreateTenantRequest, Database,
     DeleteCollectionRequest, DeleteDatabaseRequest, GetCollectionRequest, GetDatabaseRequest,
     GetResponse, GetTenantRequest, GetTenantResponse, HeartbeatError, IncludeList,
-    InternalCollectionConfiguration, KnnIndex, ListCollectionsRequest, ListDatabasesRequest,
-    Metadata, QueryResponse, UpdateCollectionConfiguration, UpdateCollectionRequest,
-    UpdateMetadata, WrappedSerdeJsonError,
+    InternalCollectionConfiguration, InternalUpdateCollectionConfiguration, KnnIndex,
+    ListCollectionsRequest, ListDatabasesRequest, Metadata, QueryResponse,
+    UpdateCollectionConfiguration, UpdateCollectionRequest, UpdateMetadata, WrappedSerdeJsonError,
 };
 use pyo3::{exceptions::PyValueError, pyclass, pyfunction, pymethods, types::PyAnyMethods, Python};
 use std::time::SystemTime;
@@ -344,11 +344,16 @@ impl Bindings {
             None => None,
         };
 
+        let configuration = match configuration_json {
+            Some(c) => Some(InternalUpdateCollectionConfiguration::try_from(c)?),
+            None => None,
+        };
+
         let request = UpdateCollectionRequest::try_new(
             collection_id,
             new_name,
             new_metadata.map(CollectionMetadataUpdate::UpdateMetadata),
-            configuration_json,
+            configuration,
         )?;
 
         let mut frontend = self.frontend.clone();
diff --git a/rust/sysdb/src/sqlite.rs b/rust/sysdb/src/sqlite.rs
@@ -12,9 +12,9 @@ use chroma_types::{
     CreateTenantError, CreateTenantResponse, Database, DatabaseUuid, DeleteCollectionError,
     DeleteDatabaseError, DeleteDatabaseResponse, GetCollectionWithSegmentsError,
     GetCollectionsError, GetDatabaseError, GetSegmentsError, GetTenantError, GetTenantResponse,
-    InternalCollectionConfiguration, ListDatabasesError, Metadata, MetadataValue, ResetError,
-    ResetResponse, Segment, SegmentScope, SegmentType, SegmentUuid, UpdateCollectionConfiguration,
-    UpdateCollectionError,
+    InternalCollectionConfiguration, InternalUpdateCollectionConfiguration, ListDatabasesError,
+    Metadata, MetadataValue, ResetError, ResetResponse, Segment, SegmentScope, SegmentType,
+    SegmentUuid, UpdateCollectionError,
 };
 use futures::TryStreamExt;
 use sea_query_binder::SqlxBinder;
@@ -356,7 +356,7 @@ impl SqliteSysDb {
         name: Option<String>,
         metadata: Option<CollectionMetadataUpdate>,
         dimension: Option<u32>,
-        configuration: Option<UpdateCollectionConfiguration>,
+        configuration: Option<InternalUpdateCollectionConfiguration>,
     ) -> Result<(), UpdateCollectionError> {
         let mut tx = self
             .db
@@ -1048,8 +1048,9 @@ mod tests {
     use super::*;
     use chroma_sqlite::db::test_utils::get_new_sqlite_db;
     use chroma_types::{
-        SegmentScope, SegmentType, SegmentUuid, UpdateHnswConfiguration, UpdateMetadata,
-        UpdateMetadataValue, VectorIndexConfiguration,
+        InternalUpdateCollectionConfiguration, SegmentScope, SegmentType, SegmentUuid,
+        UpdateHnswConfiguration, UpdateMetadata, UpdateMetadataValue,
+        UpdateVectorIndexConfiguration, VectorIndexConfiguration,
     };
 
     #[tokio::test]
@@ -1354,13 +1355,14 @@ mod tests {
                 Some("new_name".to_string()),
                 Some(CollectionMetadataUpdate::UpdateMetadata(metadata)),
                 Some(1024),
-                Some(UpdateCollectionConfiguration {
-                    hnsw: Some(UpdateHnswConfiguration {
-                        ef_search: Some(20),
-                        num_threads: Some(4),
-                        ..Default::default()
-                    }),
-                    spann: None,
+                Some(InternalUpdateCollectionConfiguration {
+                    vector_index: Some(UpdateVectorIndexConfiguration::Hnsw(Some(
+                        UpdateHnswConfiguration {
+                            ef_search: Some(10),
+                            num_threads: Some(2),
+                            ..Default::default()
+                        },
+                    ))),
                     embedding_function: None,
                 }),
             )
diff --git a/rust/sysdb/src/sysdb.rs b/rust/sysdb/src/sysdb.rs
@@ -14,9 +14,9 @@ use chroma_types::{
     DeleteCollectionError, DeleteDatabaseError, DeleteDatabaseResponse, GetCollectionSizeError,
     GetCollectionWithSegmentsError, GetCollectionsError, GetDatabaseError, GetDatabaseResponse,
     GetSegmentsError, GetTenantError, GetTenantResponse, InternalCollectionConfiguration,
-    ListCollectionVersionsError, ListDatabasesError, ListDatabasesResponse, Metadata, ResetError,
-    ResetResponse, SegmentFlushInfo, SegmentFlushInfoConversionError, SegmentUuid,
-    UpdateCollectionConfiguration, UpdateCollectionError, VectorIndexConfiguration,
+    InternalUpdateCollectionConfiguration, ListCollectionVersionsError, ListDatabasesError,
+    ListDatabasesResponse, Metadata, ResetError, ResetResponse, SegmentFlushInfo,
+    SegmentFlushInfoConversionError, SegmentUuid, UpdateCollectionError, VectorIndexConfiguration,
 };
 use chroma_types::{
     BatchGetCollectionSoftDeleteStatusError, BatchGetCollectionVersionFilePathsError, Collection,
@@ -299,7 +299,7 @@ impl SysDb {
         name: Option<String>,
         metadata: Option<CollectionMetadataUpdate>,
         dimension: Option<u32>,
-        configuration: Option<UpdateCollectionConfiguration>,
+        configuration: Option<InternalUpdateCollectionConfiguration>,
     ) -> Result<(), UpdateCollectionError> {
         match self {
             SysDb::Grpc(grpc) => {
@@ -984,7 +984,7 @@ impl GrpcSysDb {
         name: Option<String>,
         metadata: Option<CollectionMetadataUpdate>,
         dimension: Option<u32>,
-        configuration: Option<UpdateCollectionConfiguration>,
+        configuration: Option<InternalUpdateCollectionConfiguration>,
     ) -> Result<(), UpdateCollectionError> {
         let mut configuration_json_str = None;
         if let Some(configuration) = configuration {
diff --git a/rust/types/src/api_types.rs b/rust/types/src/api_types.rs
@@ -1,5 +1,5 @@
 use crate::collection_configuration::InternalCollectionConfiguration;
-use crate::collection_configuration::UpdateCollectionConfiguration;
+use crate::collection_configuration::InternalUpdateCollectionConfiguration;
 use crate::error::QueryConversionError;
 use crate::operator::GetResult;
 use crate::operator::KnnBatchResult;
@@ -724,15 +724,15 @@ pub struct UpdateCollectionRequest {
     pub new_name: Option<String>,
     #[validate(custom(function = "validate_non_empty_collection_update_metadata"))]
     pub new_metadata: Option<CollectionMetadataUpdate>,
-    pub new_configuration: Option<UpdateCollectionConfiguration>,
+    pub new_configuration: Option<InternalUpdateCollectionConfiguration>,
 }
 
 impl UpdateCollectionRequest {
     pub fn try_new(
         collection_id: CollectionUuid,
         new_name: Option<String>,
         new_metadata: Option<CollectionMetadataUpdate>,
-        new_configuration: Option<UpdateCollectionConfiguration>,
+        new_configuration: Option<InternalUpdateCollectionConfiguration>,
     ) -> Result<Self, ChromaValidationError> {
         let request = Self {
             collection_id,
diff --git a/rust/types/src/collection_configuration.rs b/rust/types/src/collection_configuration.rs
@@ -152,39 +152,47 @@ impl InternalCollectionConfiguration {
         }
     }
 
-    pub fn update(&mut self, configuration: &UpdateCollectionConfiguration) {
+    pub fn update(&mut self, configuration: &InternalUpdateCollectionConfiguration) {
         // Update vector_index if it exists in the update configuration
 
-        if let Some(hnsw_config) = &configuration.hnsw {
-            if let VectorIndexConfiguration::Hnsw(current_config) = &mut self.vector_index {
-                // Update only the non-None fields from the update configuration
-                if let Some(ef_search) = hnsw_config.ef_search {
-                    current_config.ef_search = ef_search;
-                }
-                if let Some(max_neighbors) = hnsw_config.max_neighbors {
-                    current_config.max_neighbors = max_neighbors;
-                }
-                if let Some(num_threads) = hnsw_config.num_threads {
-                    current_config.num_threads = num_threads;
-                }
-                if let Some(resize_factor) = hnsw_config.resize_factor {
-                    current_config.resize_factor = resize_factor;
-                }
-                if let Some(sync_threshold) = hnsw_config.sync_threshold {
-                    current_config.sync_threshold = sync_threshold;
-                }
-                if let Some(batch_size) = hnsw_config.batch_size {
-                    current_config.batch_size = batch_size;
-                }
-            }
-        }
-        if let Some(spann_config) = &configuration.spann {
-            if let VectorIndexConfiguration::Spann(current_config) = &mut self.vector_index {
-                if let Some(search_nprobe) = spann_config.search_nprobe {
-                    current_config.search_nprobe = search_nprobe;
+        if let Some(vector_index) = &configuration.vector_index {
+            match vector_index {
+                UpdateVectorIndexConfiguration::Hnsw(hnsw_config) => {
+                    if let VectorIndexConfiguration::Hnsw(current_config) = &mut self.vector_index {
+                        if let Some(update_config) = hnsw_config {
+                            if let Some(ef_search) = update_config.ef_search {
+                                current_config.ef_search = ef_search;
+                            }
+                            if let Some(max_neighbors) = update_config.max_neighbors {
+                                current_config.max_neighbors = max_neighbors;
+                            }
+                            if let Some(num_threads) = update_config.num_threads {
+                                current_config.num_threads = num_threads;
+                            }
+                            if let Some(resize_factor) = update_config.resize_factor {
+                                current_config.resize_factor = resize_factor;
+                            }
+                            if let Some(sync_threshold) = update_config.sync_threshold {
+                                current_config.sync_threshold = sync_threshold;
+                            }
+                            if let Some(batch_size) = update_config.batch_size {
+                                current_config.batch_size = batch_size;
+                            }
+                        }
+                    }
                 }
-                if let Some(ef_search) = spann_config.ef_search {
-                    current_config.ef_search = ef_search;
+                UpdateVectorIndexConfiguration::Spann(spann_config) => {
+                    if let VectorIndexConfiguration::Spann(current_config) = &mut self.vector_index
+                    {
+                        if let Some(update_config) = spann_config {
+                            if let Some(search_nprobe) = update_config.search_nprobe {
+                                current_config.search_nprobe = search_nprobe;
+                            }
+                            if let Some(ef_search) = update_config.ef_search {
+                                current_config.ef_search = ef_search;
+                            }
+                        }
+                    }
                 }
             }
         }
@@ -398,6 +406,48 @@ pub struct UpdateCollectionConfiguration {
     pub embedding_function: Option<EmbeddingFunctionConfiguration>,
 }
 
+#[derive(Deserialize, Serialize, ToSchema, Debug, Clone)]
+pub struct InternalUpdateCollectionConfiguration {
+    pub vector_index: Option<UpdateVectorIndexConfiguration>,
+    pub embedding_function: Option<EmbeddingFunctionConfiguration>,
+}
+
+#[derive(Debug, Error)]
+pub enum UpdateCollectionConfigurationToInternalUpdateConfigurationError {
+    #[error("Multiple vector index configurations provided")]
+    MultipleVectorIndexConfigurations,
+}
+
+impl ChromaError for UpdateCollectionConfigurationToInternalUpdateConfigurationError {
+    fn code(&self) -> ErrorCodes {
+        match self {
+            Self::MultipleVectorIndexConfigurations => ErrorCodes::InvalidArgument,
+        }
+    }
+}
+
+impl TryFrom<UpdateCollectionConfiguration> for InternalUpdateCollectionConfiguration {
+    type Error = UpdateCollectionConfigurationToInternalUpdateConfigurationError;
+
+    fn try_from(value: UpdateCollectionConfiguration) -> Result<Self, Self::Error> {
+        match (value.hnsw, value.spann) {
+            (Some(_), Some(_)) => Err(Self::Error::MultipleVectorIndexConfigurations),
+            (Some(hnsw), None) => Ok(InternalUpdateCollectionConfiguration {
+                vector_index: Some(UpdateVectorIndexConfiguration::Hnsw(Some(hnsw))),
+                embedding_function: value.embedding_function,
+            }),
+            (None, Some(spann)) => Ok(InternalUpdateCollectionConfiguration {
+                vector_index: Some(UpdateVectorIndexConfiguration::Spann(Some(spann))),
+                embedding_function: value.embedding_function,
+            }),
+            (None, None) => Ok(InternalUpdateCollectionConfiguration {
+                vector_index: None,
+                embedding_function: value.embedding_function,
+            }),
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::hnsw_configuration::HnswConfiguration;
