[ENH] recognize and flush new metadata keys to schema on local compaction

Author: jairad26

chromadb/test/api/test_schema_e2e.py

@@ -363,7 +363,8 @@ def test_schema_defaults_enable_indexed_operations(
     # Ensure underlying schema persisted across fetches
     reloaded = client.get_collection(collection.name)
     assert reloaded.schema is not None
-    assert reloaded.schema.serialize_to_json() == schema.serialize_to_json()
+    if not is_spann_disabled_mode:
+        assert reloaded.schema.serialize_to_json() == schema.serialize_to_json()
 
 
 def test_get_or_create_and_get_collection_preserve_schema(
@@ -541,7 +542,8 @@ def test_schema_persistence_with_custom_overrides(
     reloaded_client = client_factories.create_client_from_system()
     reloaded_collection = reloaded_client.get_collection(name=collection.name)
     assert reloaded_collection.schema is not None
-    assert reloaded_collection.schema.serialize_to_json() == expected_schema_json
+    if not is_spann_disabled_mode:
+        assert reloaded_collection.schema.serialize_to_json() == expected_schema_json
 
     fetched = reloaded_collection.get(where={"title": "Schema Persistence"})
     assert set(fetched["ids"]) == {"persist-1"}
@@ -784,7 +786,6 @@ def _expect_disabled_error(operation: Callable[[], Any]) -> None:
         _expect_disabled_error(operation)
 
 
-@pytest.mark.skipif(is_spann_disabled_mode, reason=skip_reason_spann_disabled)
 def test_schema_discovers_new_keys_after_compaction(
     client_factories: "ClientFactories",
 ) -> None:
@@ -802,7 +803,8 @@ def test_schema_discovers_new_keys_after_compaction(
 
     collection.add(ids=ids, documents=documents, metadatas=metadatas)
 
-    wait_for_version_increase(client, collection.name, initial_version)
+    if not is_spann_disabled_mode:
+        wait_for_version_increase(client, collection.name, initial_version)
 
     reloaded = client.get_collection(collection.name)
     assert reloaded.schema is not None
@@ -828,7 +830,8 @@ def test_schema_discovers_new_keys_after_compaction(
         metadatas=upsert_metadatas,
     )
 
-    wait_for_version_increase(client, collection.name, next_version)
+    if not is_spann_disabled_mode:
+        wait_for_version_increase(client, collection.name, next_version)
 
     post_upsert = client.get_collection(collection.name)
     assert post_upsert.schema is not None
@@ -852,7 +855,6 @@ def test_schema_discovers_new_keys_after_compaction(
     assert "discover_upsert" in persisted.schema.keys
 
 
-@pytest.mark.skipif(is_spann_disabled_mode, reason=skip_reason_spann_disabled)
 def test_schema_rejects_conflicting_discoverable_key_types(
     client_factories: "ClientFactories",
 ) -> None:
@@ -868,7 +870,8 @@ def test_schema_rejects_conflicting_discoverable_key_types(
     documents = [f"doc {i}" for i in range(251)]
     collection.add(ids=ids, documents=documents, metadatas=metadatas)
 
-    wait_for_version_increase(client, collection.name, initial_version)
+    if not is_spann_disabled_mode:
+        wait_for_version_increase(client, collection.name, initial_version)
 
     collection.upsert(
         ids=["conflict-bad"],
@@ -1029,7 +1032,6 @@ def test_schema_embedding_configuration_enforced(
     assert "sparse_auto" not in numeric_metadata
 
 
-@pytest.mark.skipif(is_spann_disabled_mode, reason=skip_reason_spann_disabled)
 def test_schema_precedence_for_overrides_discoverables_and_defaults(
     client_factories: "ClientFactories",
 ) -> None:
@@ -1054,7 +1056,9 @@ def test_schema_precedence_for_overrides_discoverables_and_defaults(
 
     initial_version = get_collection_version(client, collection.name)
     collection.add(ids=ids, documents=documents, metadatas=metadatas)
-    wait_for_version_increase(client, collection.name, initial_version)
+
+    if not is_spann_disabled_mode:
+        wait_for_version_increase(client, collection.name, initial_version)
 
     schema_state = client.get_collection(collection.name).schema
     assert schema_state is not None

rust/log/src/local_compaction_manager.rs

@@ -140,6 +140,7 @@ impl Handler<BackfillMessage> for LocalCompactionManager {
             .sysdb
             .get_collection_with_segments(message.collection_id)
             .await?;
+        let schema_previously_persisted = collection_and_segments.collection.schema.is_some();
         collection_and_segments
             .collection
             .reconcile_schema_with_config(KnnIndex::Hnsw)?;
@@ -206,14 +207,31 @@ impl Handler<BackfillMessage> for LocalCompactionManager {
             .begin()
             .await
             .map_err(|_| CompactionManagerError::MetadataApplyLogsFailed)?;
-        metadata_writer
+        let apply_outcome = metadata_writer
             .apply_logs(
                 mt_data_chunk,
                 collection_and_segments.metadata_segment.id,
+                if schema_previously_persisted {
+                    collection_and_segments.collection.schema.clone()
+                } else {
+                    None
+                },
                 &mut *tx,
             )
             .await
             .map_err(|_| CompactionManagerError::MetadataApplyLogsFailed)?;
+        if schema_previously_persisted {
+            if let Some(updated_schema) = apply_outcome.schema_update {
+                metadata_writer
+                    .update_collection_schema(
+                        collection_and_segments.collection.collection_id,
+                        &updated_schema,
+                        &mut *tx,
+                    )
+                    .await
+                    .map_err(|_| CompactionManagerError::MetadataApplyLogsFailed)?;
+            }
+        }
         tx.commit()
             .await
             .map_err(|_| CompactionManagerError::MetadataApplyLogsFailed)?;

rust/segment/Cargo.toml

@@ -11,6 +11,7 @@ roaring = { workspace = true }
 sea-query = { workspace = true }
 sea-query-binder = { workspace = true, features = ["sqlx-sqlite"] }
 serde = { workspace = true }
+serde_json = { workspace = true }
 sqlx = { workspace = true }
 serde-pickle = "1.2.0"
 tantivy = { workspace = true }

rust/segment/src/sqlite_metadata.rs

@@ -7,17 +7,17 @@ use chroma_error::{ChromaError, ErrorCodes};
 use chroma_sqlite::{
     db::SqliteDb,
     helpers::{delete_metadata, update_metadata},
-    table::{EmbeddingFulltextSearch, EmbeddingMetadata, Embeddings, MaxSeqId},
+    table::{Collections, EmbeddingFulltextSearch, EmbeddingMetadata, Embeddings, MaxSeqId},
 };
 use chroma_types::{
     operator::{
         CountResult, Filter, GetResult, Limit, Projection, ProjectionOutput, ProjectionRecord, Scan,
     },
     plan::{Count, Get},
-    BooleanOperator, Chunk, CompositeExpression, DocumentExpression, DocumentOperator, LogRecord,
-    MetadataComparison, MetadataExpression, MetadataSetValue, MetadataValue,
-    MetadataValueConversionError, Operation, OperationRecord, PrimitiveOperator, SegmentUuid,
-    SetOperator, UpdateMetadataValue, Where, CHROMA_DOCUMENT_KEY,
+    BooleanOperator, Chunk, CollectionUuid, CompositeExpression, DocumentExpression,
+    DocumentOperator, LogRecord, MetadataComparison, MetadataExpression, MetadataSetValue,
+    MetadataValue, MetadataValueConversionError, Operation, OperationRecord, PrimitiveOperator,
+    Schema, SegmentUuid, SetOperator, UpdateMetadataValue, Where, CHROMA_DOCUMENT_KEY,
 };
 use sea_query::{
     Alias, DeleteStatement, Expr, ExprTrait, Func, InsertStatement, LikeExpr, OnConflict, Query,
@@ -41,6 +41,8 @@ pub enum SqliteMetadataError {
     SeaQuery(#[from] sea_query::error::Error),
     #[error(transparent)]
     Sqlx(#[from] sqlx::Error),
+    #[error("Could not serialize schema: {0}")]
+    SerializeSchema(#[from] serde_json::Error),
 }
 
 impl ChromaError for SqliteMetadataError {
@@ -53,6 +55,10 @@ pub struct SqliteMetadataWriter {
     pub db: SqliteDb,
 }
 
+pub struct ApplyLogsOutcome {
+    pub schema_update: Option<Schema>,
+}
+
 impl SqliteMetadataWriter {
     pub fn new(db: SqliteDb) -> Self {
         Self { db }
@@ -278,18 +284,66 @@ impl SqliteMetadataWriter {
         Ok(self.db.get_conn().begin().await?)
     }
 
+    pub async fn update_collection_schema<C>(
+        &self,
+        collection_id: CollectionUuid,
+        schema: &Schema,
+        tx: &mut C,
+    ) -> Result<(), SqliteMetadataError>
+    where
+        for<'connection> &'connection mut C: sqlx::Executor<'connection, Database = sqlx::Sqlite>,
+    {
+        let schema_str = serde_json::to_string(schema)?;
+        let (sql, values) = Query::update()
+            .table(Collections::Table)
+            .value(Collections::SchemaStr, schema_str)
+            .and_where(
+                Expr::col((Collections::Table, Collections::Id)).eq(collection_id.to_string()),
+            )
+            .build_sqlx(SqliteQueryBuilder);
+        sqlx::query_with(&sql, values).execute(&mut *tx).await?;
+        Ok(())
+    }
+
+    fn ensure_schema_for_update_value(
+        schema: &mut Option<Schema>,
+        key: &str,
+        value: &UpdateMetadataValue,
+    ) -> bool {
+        if key == CHROMA_DOCUMENT_KEY {
+            return false;
+        }
+        match value {
+            UpdateMetadataValue::None => false,
+            _ => {
+                if let Some(schema_mut) = schema.as_mut() {
+                    if let Ok(metadata_value) = MetadataValue::try_from(value) {
+                        return schema_mut
+                            .ensure_key_from_metadata(key, metadata_value.value_type());
+                    }
+                }
+                false
+            }
+        }
+    }
+
     pub async fn apply_logs<C>(
         &self,
         logs: Chunk<LogRecord>,
         segment_id: SegmentUuid,
+        schema: Option<Schema>,
         tx: &mut C,
-    ) -> Result<(), SqliteMetadataError>
+    ) -> Result<ApplyLogsOutcome, SqliteMetadataError>
     where
         for<'connection> &'connection mut C: sqlx::Executor<'connection, Database = sqlx::Sqlite>,
     {
         if logs.is_empty() {
-            return Ok(());
+            return Ok(ApplyLogsOutcome {
+                schema_update: None,
+            });
         }
+        let mut schema = schema;
+        let mut schema_modified = false;
         let mut max_seq_id = u64::MIN;
         for (
             LogRecord {
@@ -323,6 +377,11 @@ impl SqliteMetadataWriter {
                         Self::add_record(tx, segment_id, log_offset_unsigned, id.clone()).await?
                     {
                         if let Some(meta) = metadata_owned {
+                            for (key, value) in meta.iter() {
+                                if Self::ensure_schema_for_update_value(&mut schema, key, value) {
+                                    schema_modified = true;
+                                }
+                            }
                             update_metadata::<EmbeddingMetadata, _, _>(tx, offset_id, meta).await?;
                         }
 
@@ -336,6 +395,11 @@ impl SqliteMetadataWriter {
                         Self::update_record(tx, segment_id, log_offset_unsigned, id.clone()).await?
                     {
                         if let Some(meta) = metadata_owned {
+                            for (key, value) in meta.iter() {
+                                if Self::ensure_schema_for_update_value(&mut schema, key, value) {
+                                    schema_modified = true;
+                                }
+                            }
                             update_metadata::<EmbeddingMetadata, _, _>(tx, offset_id, meta).await?;
                         }
 
@@ -351,6 +415,11 @@ impl SqliteMetadataWriter {
                             .await?;
 
                     if let Some(meta) = metadata_owned {
+                        for (key, value) in meta.iter() {
+                            if Self::ensure_schema_for_update_value(&mut schema, key, value) {
+                                schema_modified = true;
+                            }
+                        }
                         update_metadata::<EmbeddingMetadata, _, _>(tx, offset_id, meta).await?;
                     }
 
@@ -371,7 +440,9 @@ impl SqliteMetadataWriter {
 
         Self::upsert_max_seq_id(tx, segment_id, max_seq_id).await?;
 
-        Ok(())
+        Ok(ApplyLogsOutcome {
+            schema_update: if schema_modified { schema } else { None },
+        })
     }
 }
 
@@ -910,7 +981,17 @@ mod tests {
             ref_seg.apply_logs(test_data.logs.clone(), metadata_seg_id);
             let mut tx = runtime.block_on(sqlite_seg_writer.begin()).expect("Should be able to start transaction");
             let data: Chunk<LogRecord> = Chunk::new(test_data.logs.clone().into());
-            runtime.block_on(sqlite_seg_writer.apply_logs(data, metadata_seg_id, &mut *tx)).expect("Should be able to apply logs");
+            runtime.block_on(sqlite_seg_writer.apply_logs(
+                data,
+                metadata_seg_id,
+                test_data
+                    .collection_and_segments
+                    .collection
+                    .schema
+                    .clone(),
+                &mut *tx,
+            ))
+            .expect("Should be able to apply logs");
             runtime.block_on(tx.commit()).expect("Should be able to commit log");
 
             let sqlite_seg_reader = SqliteMetadataReader {
@@ -938,7 +1019,17 @@ mod tests {
             ref_seg.apply_logs(test_data.logs.clone(), metadata_seg_id);
             let mut tx = runtime.block_on(sqlite_seg_writer.begin()).expect("Should be able to start transaction");
             let data: Chunk<LogRecord> = Chunk::new(test_data.logs.clone().into());
-            runtime.block_on(sqlite_seg_writer.apply_logs(data, metadata_seg_id, &mut *tx)).expect("Should be able to apply logs");
+            runtime.block_on(sqlite_seg_writer.apply_logs(
+                data,
+                metadata_seg_id,
+                test_data
+                    .collection_and_segments
+                    .collection
+                    .schema
+                    .clone(),
+                &mut *tx,
+            ))
+            .expect("Should be able to apply logs");
             runtime.block_on(tx.commit()).expect("Should be able to commit log");
 
             let sqlite_seg_reader = SqliteMetadataReader {
@@ -1020,7 +1111,12 @@ mod tests {
             .expect("Should be able to start transaction");
         let data: Chunk<LogRecord> = Chunk::new(logs.into());
         sqlite_seg_writer
-            .apply_logs(data, metadata_seg_id, &mut *tx)
+            .apply_logs(
+                data,
+                metadata_seg_id,
+                collection_and_segments.collection.schema.clone(),
+                &mut *tx,
+            )
             .await
             .expect("Should be able to apply logs");
         tx.commit().await.expect("Should be able to commit log");
@@ -1140,7 +1236,12 @@ mod tests {
             .expect("Should be able to start transaction");
         let data: Chunk<LogRecord> = Chunk::new(logs.into());
         sqlite_seg_writer
-            .apply_logs(data, metadata_seg_id, &mut *tx)
+            .apply_logs(
+                data,
+                metadata_seg_id,
+                collection_and_segments.collection.schema.clone(),
+                &mut *tx,
+            )
             .await
             .expect("Should be able to apply logs");
         tx.commit().await.expect("Should be able to commit log");