[ENH] Add local support for schema

Author: jairad26

chromadb/api/rust.py

@@ -193,7 +193,7 @@ def list_collections(
             CollectionModel(
                 id=collection.id,
                 name=collection.name,
-                serialized_schema=None,
+                serialized_schema=collection.schema,
                 configuration_json=collection.configuration,
                 metadata=collection.metadata,
                 dimension=collection.dimension,
@@ -229,14 +229,25 @@ def create_collection(
         else:
             configuration_json_str = None
 
+        if schema:
+            schema_str = json.dumps(schema.serialize_to_json())
+        else:
+            schema_str = None
+
         collection = self.bindings.create_collection(
-            name, configuration_json_str, metadata, get_or_create, tenant, database
+            name,
+            configuration_json_str,
+            schema_str,
+            metadata,
+            get_or_create,
+            tenant,
+            database,
         )
         collection_model = CollectionModel(
             id=collection.id,
             name=collection.name,
             configuration_json=collection.configuration,
-            serialized_schema=None,
+            serialized_schema=collection.schema,
             metadata=collection.metadata,
             dimension=collection.dimension,
             tenant=collection.tenant,
@@ -256,7 +267,7 @@ def get_collection(
             id=collection.id,
             name=collection.name,
             configuration_json=collection.configuration,
-            serialized_schema=None,
+            serialized_schema=collection.schema,
             metadata=collection.metadata,
             dimension=collection.dimension,
             tenant=collection.tenant,

chromadb/chromadb_rust_bindings.pyi

@@ -95,6 +95,7 @@ class Bindings:
         self,
         name: str,
         configuration_json_str: Optional[str] = None,
+        schema_str: Optional[str] = None,
         metadata: Optional[CollectionMetadata] = None,
         get_or_create: bool = False,
         tenant: str = DEFAULT_TENANT,

chromadb/test/api/test_schema_e2e.py

@@ -93,7 +93,6 @@ def build_from_config(config: Dict[str, Any]) -> "RecordingSearchEmbeddingFuncti
         return RecordingSearchEmbeddingFunction(config.get("label", "default"))
 
 
-@pytest.mark.skipif(is_spann_disabled_mode, reason=skip_reason_spann_disabled)
 def test_schema_spann_vector_config_persistence(
     client_factories: "ClientFactories",
 ) -> None:
@@ -135,29 +134,46 @@ def test_schema_spann_vector_config_persistence(
     assert vector_index is not None
     assert vector_index.enabled is True
     assert vector_index.config is not None
-    assert vector_index.config.spann is not None
-    spann_config = vector_index.config.spann
-    assert spann_config.search_nprobe == 16
-    assert spann_config.write_nprobe == 32
-    assert spann_config.ef_construction == 120
-    assert spann_config.max_neighbors == 24
+
+    if not is_spann_disabled_mode:
+        assert vector_index.config.spann is not None
+        spann_config = vector_index.config.spann
+        assert spann_config.search_nprobe == 16
+        assert spann_config.write_nprobe == 32
+        assert spann_config.ef_construction == 120
+        assert spann_config.max_neighbors == 24
+    else:
+        assert vector_index.config.spann is None
+        assert vector_index.config.hnsw is not None
+        hnsw_config = vector_index.config.hnsw
+        assert hnsw_config.ef_construction == 100
+        assert hnsw_config.ef_search == 100
+        assert hnsw_config.max_neighbors == 16
+        assert hnsw_config.resize_factor == 1.2
 
     ef = vector_index.config.embedding_function
     assert ef is not None
     assert ef.name() == "simple_ef"
     assert ef.get_config() == {"dim": 6}
 
     persisted_json = persisted_schema.serialize_to_json()
-    spann_json = persisted_json["keys"]["#embedding"]["float_list"]["vector_index"][
-        "config"
-    ]["spann"]
-    assert spann_json["search_nprobe"] == 16
-    assert spann_json["write_nprobe"] == 32
+    if not is_spann_disabled_mode:
+        spann_json = persisted_json["keys"]["#embedding"]["float_list"]["vector_index"][
+            "config"
+        ]["spann"]
+        assert spann_json["search_nprobe"] == 16
+        assert spann_json["write_nprobe"] == 32
+    else:
+        hnsw_json = persisted_json["keys"]["#embedding"]["float_list"]["vector_index"][
+            "config"
+        ]["hnsw"]
+        assert hnsw_json["ef_construction"] == 100
+        assert hnsw_json["ef_search"] == 100
+        assert hnsw_json["max_neighbors"] == 16
 
     client_reloaded = client_factories.create_client_from_system()
     reloaded_collection = client_reloaded.get_collection(
         name=collection_name,
-        embedding_function=SimpleEmbeddingFunction(dim=6),  # type: ignore[arg-type]
     )
 
     reloaded_schema = reloaded_collection.schema
@@ -167,9 +183,23 @@ def test_schema_spann_vector_config_persistence(
     reloaded_vector_index = reloaded_embedding_override.vector_index
     assert reloaded_vector_index is not None
     assert reloaded_vector_index.config is not None
-    assert reloaded_vector_index.config.spann is not None
-    assert reloaded_vector_index.config.spann.search_nprobe == 16
-    assert reloaded_vector_index.config.spann.write_nprobe == 32
+    if not is_spann_disabled_mode:
+        assert reloaded_vector_index.config.spann is not None
+        assert reloaded_vector_index.config.spann.search_nprobe == 16
+        assert reloaded_vector_index.config.spann.write_nprobe == 32
+    else:
+        assert reloaded_vector_index.config.hnsw is not None
+        assert reloaded_vector_index.config.hnsw.ef_construction == 100
+        assert reloaded_vector_index.config.hnsw.ef_search == 100
+        assert reloaded_vector_index.config.hnsw.max_neighbors == 16
+        assert reloaded_vector_index.config.hnsw.resize_factor == 1.2
+
+    config = reloaded_collection.configuration
+    assert config is not None
+    config_ef = config.get("embedding_function")
+    assert config_ef is not None
+    assert config_ef.name() == "simple_ef"
+    assert config_ef.get_config() == {"dim": 6}
 
 
 @register_sparse_embedding_function
@@ -650,6 +680,52 @@ def test_schema_delete_index_and_restore(
     assert set(search["ids"]) == {"key-enabled"}
 
 
+def test_disabled_metadata_index_filters_raise_invalid_argument_all_modes(
+    client_factories: "ClientFactories",
+) -> None:
+    """Disabled metadata inverted index should block filter-based operations in get, query, and delete for local, single node, and distributed."""
+    schema = Schema().delete_index(
+        key="restricted_tag", config=StringInvertedIndexConfig()
+    )
+    collection, _ = _create_isolated_collection(client_factories, schema=schema)
+
+    collection.add(
+        ids=["restricted-doc"],
+        embeddings=cast(Embeddings, [[0.1, 0.2, 0.3, 0.4]]),
+        metadatas=[{"restricted_tag": "blocked"}],
+        documents=["doc"],
+    )
+
+    assert collection.schema is not None
+    schema_entry = collection.schema.keys["restricted_tag"].string
+    assert schema_entry is not None
+    index_config = schema_entry.string_inverted_index
+    assert index_config is not None
+    assert index_config.enabled is False
+
+    filter_payload: Dict[str, Any] = {"restricted_tag": "blocked"}
+
+    def _expect_disabled_error(operation: Callable[[], Any]) -> None:
+        with pytest.raises(InvalidArgumentError) as exc_info:
+            operation()
+        assert "Cannot filter using metadata key 'restricted_tag'" in str(
+            exc_info.value
+        )
+
+    operations: List[Callable[[], Any]] = [
+        lambda: collection.get(where=filter_payload),
+        lambda: collection.query(
+            query_embeddings=cast(Embeddings, [[0.1, 0.2, 0.3, 0.4]]),
+            n_results=1,
+            where=filter_payload,
+        ),
+        lambda: collection.delete(where=filter_payload),
+    ]
+
+    for operation in operations:
+        _expect_disabled_error(operation)
+
+
 @pytest.mark.skipif(is_spann_disabled_mode, reason=skip_reason_spann_disabled)
 def test_disabled_metadata_index_filters_raise_invalid_argument(
     client_factories: "ClientFactories",

rust/cli/src/client/chroma_client.rs

@@ -4,7 +4,7 @@ use crate::client::prelude::CollectionModel;
 use crate::client::utils::send_request;
 use crate::utils::Profile;
 use axum::http::Method;
-use chroma_types::{CollectionConfiguration, CreateCollectionPayload, Metadata};
+use chroma_types::{CollectionConfiguration, CreateCollectionPayload, Metadata, Schema};
 use std::error::Error;
 use std::ops::Deref;
 use thiserror::Error;
@@ -104,6 +104,7 @@ impl ChromaClient {
         name: String,
         metadata: Option<Metadata>,
         configuration: Option<CollectionConfiguration>,
+        schema: Option<Schema>,
     ) -> Result<Collection, Box<dyn Error>> {
         let route = format!(
             "/api/v2/tenants/{}/databases/{}/collections",
@@ -115,7 +116,7 @@ impl ChromaClient {
             configuration,
             metadata,
             get_or_create: false,
-            schema: None,
+            schema,
         };
         let response = send_request::<CreateCollectionPayload, CollectionModel>(
             &self.host,

rust/cli/src/commands/copy.rs

@@ -239,6 +239,7 @@ async fn copy_collections(
                 collection.name.clone(),
                 collection.metadata.clone(),
                 Some(CollectionConfiguration::from(collection.config.clone())),
+                collection.schema.clone(),
             )
             .await
             .map_err(|_| ChromaClientError::CreateCollection(collection.name.clone()))?;

rust/cli/src/commands/vacuum.rs

@@ -11,7 +11,7 @@ use chroma_segment::local_segment_manager::LocalSegmentManager;
 use chroma_sqlite::db::SqliteDb;
 use chroma_sysdb::SysDb;
 use chroma_system::System;
-use chroma_types::{CollectionUuid, ListCollectionsRequest};
+use chroma_types::{CollectionUuid, KnnIndex, ListCollectionsRequest};
 use clap::Parser;
 use colored::Colorize;
 use dialoguer::Confirm;
@@ -101,11 +101,17 @@ async fn trigger_vector_segments_max_seq_id_migration(
     sqlite: &SqliteDb,
     sysdb: &mut SysDb,
     segment_manager: &LocalSegmentManager,
+    default_knn_index: KnnIndex,
 ) -> Result<(), Box<dyn Error>> {
     let collection_ids = get_collection_ids_to_migrate(sqlite).await?;
 
     for collection_id in collection_ids {
-        let collection = sysdb.get_collection_with_segments(collection_id).await?;
+        let mut collection = sysdb.get_collection_with_segments(collection_id).await?;
+
+        collection
+            .collection
+            .reconcile_schema_with_config(default_knn_index)
+            .map_err(|e| Box::new(e) as Box<dyn Error>)?;
 
         // If collection is uninitialized, that means nothing has been written yet.
         let dim = match collection.collection.dimension {
@@ -138,7 +144,7 @@ async fn configure_sql_embedding_queue(log: &SqliteLog) -> Result<(), Box<dyn Er
 pub async fn vacuum_chroma(config: FrontendConfig) -> Result<(), Box<dyn Error>> {
     let system = System::new();
     let registry = Registry::new();
-    let mut frontend = Frontend::try_from_config(&(config, system), &registry).await?;
+    let mut frontend = Frontend::try_from_config(&(config.clone(), system), &registry).await?;
 
     let sqlite = registry.get::<SqliteDb>()?;
     let segment_manager = registry.get::<LocalSegmentManager>()?;
@@ -147,7 +153,13 @@ pub async fn vacuum_chroma(config: FrontendConfig) -> Result<(), Box<dyn Error>>
 
     println!("Purging the log...\n");
 
-    trigger_vector_segments_max_seq_id_migration(&sqlite, &mut sysdb, &segment_manager).await?;
+    trigger_vector_segments_max_seq_id_migration(
+        &sqlite,
+        &mut sysdb,
+        &segment_manager,
+        config.default_knn_index,
+    )
+    .await?;
 
     let tenant = String::from("default_tenant");
     let database = String::from("default_database");

rust/frontend/src/config.rs

@@ -142,7 +142,7 @@ fn default_enable_span_indexing() -> bool {
 }
 
 fn default_enable_schema() -> bool {
-    false
+    true
 }
 
 pub fn default_min_records_for_task() -> u64 {

rust/frontend/src/executor/local.rs

@@ -204,17 +204,21 @@ impl LocalExecutor {
             allowed_offset_ids.push(offset_id);
         }
 
-        let distance_function = match collection_and_segments
+        let hnsw_config = collection_and_segments
             .collection
-            .config
-            .get_hnsw_config_with_legacy_fallback(&plan.scan.collection_and_segments.vector_segment)
-        {
-            Ok(Some(config)) => config.space,
-            Ok(None) => return Err(ExecutorError::CollectionMissingHnswConfiguration),
-            Err(err) => {
-                return Err(ExecutorError::Internal(Box::new(err)));
-            }
-        };
+            .schema
+            .as_ref()
+            .map(|schema| {
+                schema.get_internal_hnsw_config_with_legacy_fallback(
+                    &plan.scan.collection_and_segments.vector_segment,
+                )
+            })
+            .transpose()
+            .map_err(|err| ExecutorError::Internal(Box::new(err)))?
+            .flatten()
+            .ok_or(ExecutorError::CollectionMissingHnswConfiguration)?;
+
+        let distance_function = hnsw_config.space;
 
         let mut results = Vec::new();
         let mut returned_user_ids = Vec::new();

rust/frontend/src/get_collection_with_segments_provider.rs

@@ -4,7 +4,8 @@ use chroma_config::Configurable;
 use chroma_error::{ChromaError, ErrorCodes};
 use chroma_sysdb::SysDb;
 use chroma_types::{
-    CollectionAndSegments, CollectionUuid, GetCollectionWithSegmentsError, Schema, SchemaError,
+    CollectionAndSegments, CollectionUuid, GetCollectionWithSegmentsError, KnnIndex, Schema,
+    SchemaError,
 };
 use serde::{Deserialize, Serialize};
 use std::{
@@ -142,6 +143,7 @@ impl CollectionsWithSegmentsProvider {
     pub(crate) async fn get_collection_with_segments(
         &mut self,
         collection_id: CollectionUuid,
+        knn_index: KnnIndex,
     ) -> Result<CollectionAndSegments, CollectionsWithSegmentsProviderError> {
         if let Some(collection_and_segments_with_ttl) = self
             .collections_with_segments_cache
@@ -187,6 +189,7 @@ impl CollectionsWithSegmentsProvider {
         let reconciled_schema = Schema::reconcile_schema_and_config(
             collection_and_segments_sysdb.collection.schema.as_ref(),
             Some(&collection_and_segments_sysdb.collection.config),
+            knn_index,
         )
         .map_err(CollectionsWithSegmentsProviderError::InvalidSchema)?;
         collection_and_segments_sysdb.collection.schema = Some(reconciled_schema);