[ENH] Add schema support to collection configuration

Author: jairad26

chromadb/api/collection_configuration.py

@@ -1,4 +1,4 @@
-from typing import TypedDict, Dict, Any, Optional, cast, get_args
+from typing import TypedDict, Dict, Any, Optional, cast, get_args, Literal
 import json
 import copy
 from chromadb.api.types import (
@@ -15,6 +15,13 @@
 from multiprocessing import cpu_count
 import warnings
 
+ValueType = Literal["int", "float", "string", "boolean"]
+
+
+class CollectionSchema(TypedDict):
+    value_type: ValueType
+    metadata_index: bool
+
 
 class HNSWConfiguration(TypedDict, total=False):
     space: Space
@@ -44,6 +51,7 @@ class CollectionConfiguration(TypedDict, total=True):
     spann: Optional[SpannConfiguration]
     embedding_function: Optional[EmbeddingFunction]  # type: ignore
     query_embedding_function: Optional[EmbeddingFunction]  # type: ignore
+    schema: Optional[Dict[str, CollectionSchema]]
 
 
 def load_collection_configuration_from_json_str(
@@ -126,6 +134,7 @@ def load_collection_configuration_from_json(
         spann=spann_config,
         embedding_function=ef,  # type: ignore
         query_embedding_function=query_ef,  # type: ignore
+        schema=config_json_map.get("schema"),
     )
 
 
@@ -278,6 +287,7 @@ class CreateCollectionConfiguration(TypedDict, total=False):
     spann: Optional[CreateSpannConfiguration]
     embedding_function: Optional[EmbeddingFunction]  # type: ignore
     query_config: Optional[QueryConfig]
+    schema: Optional[Dict[str, CollectionSchema]]
 
 
 def create_collection_configuration_from_legacy_collection_metadata(
@@ -416,6 +426,7 @@ def create_collection_configuration_to_json(
         "spann": spann_config,
         "embedding_function": ef_config,
         "query_config": query_config,
+        "schema": config.get("schema"),
     }
 
 
@@ -488,6 +499,7 @@ class UpdateCollectionConfiguration(TypedDict, total=False):
     spann: Optional[UpdateSpannConfiguration]
     embedding_function: Optional[EmbeddingFunction]  # type: ignore
     query_config: Optional[QueryConfig]
+    schema: Optional[Dict[str, CollectionSchema]]
 
 
 def update_collection_configuration_from_legacy_collection_metadata(
@@ -587,6 +599,7 @@ def update_collection_configuration_to_json(
         "spann": spann_config,
         "embedding_function": ef_config,
         "query_config": query_config,
+        "schema": config.get("schema"),
     }
 
 
@@ -750,14 +763,34 @@ def overwrite_collection_configuration(
                     ef_config[k] = v
                 query_ef = updated_embedding_function.build_from_config(ef_config)
 
+    existing_schema = existing_config.get("schema")
+    new_diff_schema = update_config.get("schema")
+    updated_schema: Optional[Dict[str, CollectionSchema]] = None
+    if existing_schema is not None:
+        if new_diff_schema is not None:
+            updated_schema = overwrite_schema(existing_schema, new_diff_schema)
+        else:
+            updated_schema = existing_schema
+    else:
+        updated_schema = new_diff_schema
+
     return CollectionConfiguration(
         hnsw=updated_hnsw_config,
         spann=updated_spann_config,
         embedding_function=updated_embedding_function,
         query_embedding_function=query_ef,
+        schema=updated_schema,
     )
 
 
+def overwrite_schema(
+    existing_schema: Dict[str, CollectionSchema],
+    new_diff_schema: Dict[str, CollectionSchema],
+) -> Dict[str, CollectionSchema]:
+    """Overwrite a schema with a new configuration"""
+    return {**existing_schema, **new_diff_schema}
+
+
 def validate_embedding_function_conflict_on_create(
     embedding_function: Optional[EmbeddingFunction],  # type: ignore
     configuration_ef: Optional[EmbeddingFunction],  # type: ignore

go/pkg/sysdb/coordinator/model/collection_configuration.go

@@ -53,10 +53,16 @@ type SpannConfiguration struct {
 	MergeThreshold        int    `json:"merge_threshold"`
 }
 
+type CollectionSchema struct {
+	ValueType     string `json:"value_type"`
+	MetadataIndex bool   `json:"metadata_index"`
+}
+
 type InternalCollectionConfiguration struct {
 	VectorIndex       *VectorIndexConfiguration       `json:"vector_index"`
 	EmbeddingFunction *EmbeddingFunctionConfiguration `json:"embedding_function,omitempty"`
 	QueryConfig       interface{}                     `json:"query_config,omitempty"`
+	Schema            *map[string]CollectionSchema    `json:"schema,omitempty"`
 }
 
 // DefaultHnswCollectionConfiguration returns a default configuration using HNSW
@@ -129,4 +135,5 @@ type InternalUpdateCollectionConfiguration struct {
 	VectorIndex       *UpdateVectorIndexConfiguration `json:"vector_index,omitempty"`
 	EmbeddingFunction *EmbeddingFunctionConfiguration `json:"embedding_function,omitempty"`
 	QueryConfig       interface{}                     `json:"query_config,omitempty"`
+	Schema            *map[string]CollectionSchema    `json:"schema,omitempty"`
 }

go/pkg/sysdb/coordinator/table_catalog.go

@@ -858,6 +858,10 @@ func (tc *Catalog) updateCollectionConfiguration(
 		existingConfig.QueryConfig = updateConfig.QueryConfig
 	}
 
+	if updateConfig.Schema != nil {
+		existingConfig.Schema = updateConfig.Schema
+	}
+
 	// Serialize updated config back to JSON
 	updatedConfigBytes, err := json.Marshal(existingConfig)
 	if err != nil {

rust/python_bindings/src/bindings.rs

@@ -281,6 +281,7 @@ impl Bindings {
                     spann: None,
                     embedding_function: None,
                     query_config: None,
+                    schema: None,
                 },
                 self.frontend.get_default_knn_index(),
             )?),

rust/segment/src/distributed_hnsw.rs

@@ -435,6 +435,7 @@ pub mod test {
                 ),
                 embedding_function: None,
                 query_config: None,
+                schema: None,
             },
             ..Default::default()
         };

rust/segment/src/distributed_spann.rs

@@ -657,6 +657,7 @@ mod test {
                 vector_index: chroma_types::VectorIndexConfiguration::Spann(params),
                 embedding_function: None,
                 query_config: None,
+                schema: None,
             },
             metadata: None,
             dimension: None,
@@ -887,6 +888,7 @@ mod test {
                 vector_index: chroma_types::VectorIndexConfiguration::Spann(params),
                 embedding_function: None,
                 query_config: None,
+                schema: None,
             },
             ..Default::default()
         };

rust/sysdb/src/sqlite.rs

@@ -373,7 +373,9 @@ impl SqliteSysDb {
             let collections = collections.unwrap();
             let collection = collections.into_iter().next().unwrap();
             let mut existing_configuration = collection.config;
-            existing_configuration.update(&configuration);
+            existing_configuration
+                .update(&configuration)
+                .map_err(|e| UpdateCollectionError::Internal(e.boxed()))?;
             configuration_json_str = Some(
                 serde_json::to_string(&existing_configuration)
                     .map_err(UpdateCollectionError::Configuration)?,
@@ -1363,6 +1365,7 @@ mod tests {
                     spann: None,
                     embedding_function: None,
                     query_config: None,
+                    schema: None,
                 }),
             )
             .await