[ENH] Add where filtering support on list collections for local chroma

jairad26 · jairad26 · commit 1542c7f50955 · 2025-07-04T16:33:59.000-07:00
diff --git a/chromadb/api/__init__.py b/chromadb/api/__init__.py
@@ -343,11 +343,13 @@ class ClientAPI(BaseAPI, ABC):
     @abstractmethod
     def list_collections(
         self,
+        where: Optional[Where] = None,
         limit: Optional[int] = None,
         offset: Optional[int] = None,
     ) -> Sequence[Collection]:
         """List all collections.
         Args:
+            where: Conditional filtering on collection metadata. Defaults to None.
             limit: The maximum number of entries to return. Defaults to None.
             offset: The number of entries to skip before returning. Defaults to None.
 
@@ -576,6 +578,7 @@ def count_collections(
     @abstractmethod
     def list_collections(
         self,
+        where: Optional[Where] = None,
         limit: Optional[int] = None,
         offset: Optional[int] = None,
         tenant: str = DEFAULT_TENANT,
diff --git a/chromadb/api/async_api.py b/chromadb/api/async_api.py
@@ -336,11 +336,13 @@ class AsyncClientAPI(AsyncBaseAPI, ABC):
     @abstractmethod
     async def list_collections(
         self,
+        where: Optional[Where] = None,
         limit: Optional[int] = None,
         offset: Optional[int] = None,
     ) -> Sequence[AsyncCollection]:
         """List all collections.
         Args:
+            where: Conditional filtering on collection metadata. Defaults to None.
             limit: The maximum number of entries to return. Defaults to None.
             offset: The number of entries to skip before returning. Defaults to None.
 
@@ -562,6 +564,7 @@ class AsyncServerAPI(AsyncBaseAPI, AsyncAdminAPI, Component):
     @abstractmethod
     async def list_collections(
         self,
+        where: Optional[Where] = None,
         limit: Optional[int] = None,
         offset: Optional[int] = None,
         tenant: str = DEFAULT_TENANT,
diff --git a/chromadb/api/async_client.py b/chromadb/api/async_client.py
@@ -159,10 +159,13 @@ async def heartbeat(self) -> int:
 
     @override
     async def list_collections(
-        self, limit: Optional[int] = None, offset: Optional[int] = None
+        self,
+        where: Optional[Where] = None,
+        limit: Optional[int] = None,
+        offset: Optional[int] = None,
     ) -> Sequence[AsyncCollection]:
         models = await self._server.list_collections(
-            limit, offset, tenant=self.tenant, database=self.database
+            where, limit, offset, tenant=self.tenant, database=self.database
         )
         return [AsyncCollection(client=self._server, model=model) for model in models]
 
diff --git a/chromadb/api/async_fastapi.py b/chromadb/api/async_fastapi.py
@@ -2,6 +2,7 @@
 from uuid import UUID
 import urllib.parse
 import orjson
+import json
 from typing import Any, Optional, cast, Tuple, Sequence, Dict
 import logging
 import httpx
@@ -251,18 +252,21 @@ async def get_user_identity(self) -> UserIdentity:
     @override
     async def list_collections(
         self,
+        where: Optional[Where] = None,
         limit: Optional[int] = None,
         offset: Optional[int] = None,
         tenant: str = DEFAULT_TENANT,
         database: str = DEFAULT_DATABASE,
     ) -> Sequence[CollectionModel]:
+        where_str = json.dumps(where) if where else None
         resp_json = await self._make_request(
             "get",
             f"/tenants/{tenant}/databases/{database}/collections",
             params=BaseHTTPClient._clean_params(
                 {
                     "limit": limit,
                     "offset": offset,
+                    "where": where_str,
                 }
             ),
         )
diff --git a/chromadb/api/client.py b/chromadb/api/client.py
@@ -124,12 +124,15 @@ def heartbeat(self) -> int:
 
     @override
     def list_collections(
-        self, limit: Optional[int] = None, offset: Optional[int] = None
+        self,
+        where: Optional[Where] = None,
+        limit: Optional[int] = None,
+        offset: Optional[int] = None,
     ) -> Sequence[Collection]:
         return [
             Collection(client=self._server, model=model)
             for model in self._server.list_collections(
-                limit, offset, tenant=self.tenant, database=self.database
+                where, limit, offset, tenant=self.tenant, database=self.database
             )
         ]
 
diff --git a/chromadb/api/fastapi.py b/chromadb/api/fastapi.py
@@ -17,6 +17,7 @@
 from chromadb.api.base_http_client import BaseHTTPClient
 from chromadb.types import Database, Tenant, Collection as CollectionModel
 from chromadb.api import ServerAPI
+import json
 
 from chromadb.api.types import (
     Documents,
@@ -205,19 +206,22 @@ def get_user_identity(self) -> UserIdentity:
     @override
     def list_collections(
         self,
+        where: Optional[Where] = None,
         limit: Optional[int] = None,
         offset: Optional[int] = None,
         tenant: str = DEFAULT_TENANT,
         database: str = DEFAULT_DATABASE,
     ) -> Sequence[CollectionModel]:
         """Returns a list of all collections"""
+        where_str = json.dumps(where) if where else None
         json_collections = self._make_request(
             "get",
             f"/tenants/{tenant}/databases/{database}/collections",
             params=BaseHTTPClient._clean_params(
                 {
                     "limit": limit,
                     "offset": offset,
+                    "where": where_str,
                 }
             ),
         )
diff --git a/chromadb/api/rust.py b/chromadb/api/rust.py
@@ -180,12 +180,15 @@ def count_collections(
     @override
     def list_collections(
         self,
+        where: Optional[Where] = None,
         limit: Optional[int] = None,
         offset: Optional[int] = None,
         tenant: str = DEFAULT_TENANT,
         database: str = DEFAULT_DATABASE,
     ) -> Sequence[CollectionModel]:
-        collections = self.bindings.list_collections(limit, offset, tenant, database)
+        collections = self.bindings.list_collections(
+            limit, offset, tenant, database, json.dumps(where) if where else None
+        )
         return [
             CollectionModel(
                 id=collection.id,
diff --git a/chromadb/api/segment.py b/chromadb/api/segment.py
@@ -325,6 +325,7 @@ def get_collection(
     @rate_limit
     def list_collections(
         self,
+        where: Optional[Where] = None,
         limit: Optional[int] = None,
         offset: Optional[int] = None,
         tenant: str = DEFAULT_TENANT,
diff --git a/chromadb/chromadb_rust_bindings.pyi b/chromadb/chromadb_rust_bindings.pyi
@@ -86,6 +86,7 @@ class Bindings:
     ) -> int: ...
     def list_collections(
         self,
+        where: Optional[str] = None,
         limit: Optional[int] = None,
         offset: Optional[int] = None,
         tenant: str = DEFAULT_TENANT,
diff --git a/chromadb/db/__init__.py b/chromadb/db/__init__.py
@@ -29,7 +29,10 @@ def get_collection(self, name: str) -> Sequence:  # type: ignore
 
     @abstractmethod
     def list_collections(
-        self, limit: Optional[int] = None, offset: Optional[int] = None
+        self,
+        where: Optional[Where] = None,
+        limit: Optional[int] = None,
+        offset: Optional[int] = None,
     ) -> Sequence:  # type: ignore
         pass
 
diff --git a/chromadb/test/property/test_collections.py b/chromadb/test/property/test_collections.py
@@ -117,6 +117,35 @@ def list_collections_with_limit_offset(self, limit: int, offset: int) -> None:
         else:
             assert len(colls) == limit
 
+    # @rule()
+    # def list_collections_with_where_filter(self) -> None:
+    #     if not self.model:
+    #         return
+    #     key_values = []
+    #     for name, metadata in self.model.items():
+    #         if metadata is not None:
+    #             for k, v in metadata.items():
+    #                 key_values.append((k, v))
+    #     if not key_values:
+    #         return
+    #     key, value = key_values[0]
+    #     expected = [
+    #         name
+    #         for name, metadata in self.model.items()
+    #         if metadata is not None and key in metadata and metadata[key] == value
+    #     ]
+    #     filtered = self.client.list_collections(
+    #         where={key: {"$eq": value}}  # type:ignore
+    #     )
+    #     filtered_names = [c.name for c in filtered]
+    #     for c in filtered:
+    #         assert (
+    #             c.metadata is not None
+    #             and key in c.metadata
+    #             and c.metadata[key] == value
+    #         )
+    #     assert sorted(filtered_names) == sorted(expected)
+
     @rule(
         target=collections,
         new_metadata=st.one_of(st.none(), strategies.collection_metadata),
diff --git a/go/pkg/sysdb/grpc/collection_service.go b/go/pkg/sysdb/grpc/collection_service.go
@@ -164,6 +164,7 @@ func (s *Server) GetCollections(ctx context.Context, req *coordinatorpb.GetColle
 	databaseName := req.Database
 	limit := req.Limit
 	offset := req.Offset
+	// where := req.Where
 
 	res := &coordinatorpb.GetCollectionsResponse{}
 
diff --git a/idl/chromadb/proto/coordinator.proto b/idl/chromadb/proto/coordinator.proto
@@ -188,6 +188,7 @@ message GetCollectionsRequest {
   optional int32 offset = 7;
   optional CollectionIdsFilter ids_filter = 8;
   optional bool include_soft_deleted = 9;
+  optional Where where = 10;
 }
 
 message GetCollectionsResponse {
diff --git a/rust/cli/src/commands/vacuum.rs b/rust/cli/src/commands/vacuum.rs
@@ -152,7 +152,8 @@ pub async fn vacuum_chroma(config: FrontendConfig) -> Result<(), Box<dyn Error>>
     let tenant = String::from("default_tenant");
     let database = String::from("default_database");
 
-    let list_collections_request = ListCollectionsRequest::try_new(tenant, database, None, 0)?;
+    let list_collections_request =
+        ListCollectionsRequest::try_new(tenant, database, None, 0, None)?;
     let collections = frontend.list_collections(list_collections_request).await?;
 
     let pb = ProgressBar::new(collections.len() as u64);
diff --git a/rust/frontend/src/impls/service_based_frontend.rs b/rust/frontend/src/impls/service_based_frontend.rs
@@ -335,6 +335,7 @@ impl ServiceBasedFrontend {
             database_name,
             limit,
             offset,
+            r#where,
             ..
         }: ListCollectionsRequest,
     ) -> Result<ListCollectionsResponse, GetCollectionsError> {
@@ -344,6 +345,7 @@ impl ServiceBasedFrontend {
                 database: Some(database_name.clone()),
                 limit,
                 offset,
+                r#where,
                 ..Default::default()
             })
             .await
diff --git a/rust/frontend/src/server.rs b/rust/frontend/src/server.rs
@@ -11,19 +11,19 @@ use chroma_metering::{
 };
 use chroma_system::System;
 use chroma_types::{
-    AddCollectionRecordsResponse, ChecklistResponse, Collection, CollectionConfiguration,
-    CollectionMetadataUpdate, CollectionUuid, CountCollectionsRequest, CountCollectionsResponse,
-    CountRequest, CountResponse, CreateCollectionRequest, CreateDatabaseRequest,
-    CreateDatabaseResponse, CreateTenantRequest, CreateTenantResponse,
+    parse_where, AddCollectionRecordsResponse, ChecklistResponse, Collection,
+    CollectionConfiguration, CollectionMetadataUpdate, CollectionUuid, CountCollectionsRequest,
+    CountCollectionsResponse, CountRequest, CountResponse, CreateCollectionRequest,
+    CreateDatabaseRequest, CreateDatabaseResponse, CreateTenantRequest, CreateTenantResponse,
     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,
+    UpdateMetadata, UpsertCollectionRecordsResponse, WhereValidationError,
 };
-use chroma_types::{ForkCollectionResponse, RawWhereFields};
+use chroma_types::{ForkCollectionResponse, RawWhereFields, Where};
 use mdac::{Rule, Scorecard, ScorecardTicket};
 use opentelemetry::global;
 use opentelemetry::metrics::{Counter, Meter};
@@ -790,6 +790,7 @@ struct ListCollectionsParams {
     limit: Option<u32>,
     #[serde(default)]
     offset: u32,
+    r#where: Option<String>,
 }
 
 /// Lists all collections in the specified database.
@@ -811,7 +812,11 @@ struct ListCollectionsParams {
 async fn list_collections(
     headers: HeaderMap,
     Path((tenant, database)): Path<(String, String)>,
-    Query(ListCollectionsParams { limit, offset }): Query<ListCollectionsParams>,
+    Query(ListCollectionsParams {
+        limit,
+        offset,
+        r#where,
+    }): Query<ListCollectionsParams>,
     State(mut server): State<FrontendServer>,
 ) -> Result<Json<ListCollectionsResponse>, ServerError> {
     server
@@ -841,7 +846,16 @@ async fn list_collections(
     server.quota_enforcer.enforce(&quota_payload).await?;
     let _guard = server
         .scorecard_request(&["op:list_collections", format!("tenant:{}", tenant).as_str()])?;
-    let request = ListCollectionsRequest::try_new(tenant, database, limit, offset)?;
+
+    let parsed_where: Option<Where> = match r#where {
+        Some(r#where) => {
+            let value = serde_json::from_str(r#where.as_str())
+                .map_err(|_| WhereValidationError::WhereClause)?;
+            parse_where(&value).ok()
+        }
+        None => None,
+    };
+    let request = ListCollectionsRequest::try_new(tenant, database, limit, offset, parsed_where)?;
     Ok(Json(server.frontend.list_collections(request).await?))
 }
 
diff --git a/rust/python_bindings/src/bindings.rs b/rust/python_bindings/src/bindings.rs
@@ -15,13 +15,13 @@ use chroma_sqlite::config::SqliteDBConfig;
 use chroma_sysdb::{SqliteSysDbConfig, SysDbConfig};
 use chroma_system::System;
 use chroma_types::{
-    Collection, CollectionConfiguration, CollectionMetadataUpdate, CountCollectionsRequest,
-    CountResponse, CreateCollectionRequest, CreateDatabaseRequest, CreateTenantRequest, Database,
-    DeleteCollectionRequest, DeleteDatabaseRequest, GetCollectionRequest, GetDatabaseRequest,
-    GetResponse, GetTenantRequest, GetTenantResponse, HeartbeatError, IncludeList,
-    InternalCollectionConfiguration, KnnIndex, ListCollectionsRequest, ListDatabasesRequest,
-    Metadata, QueryResponse, UpdateCollectionConfiguration, UpdateCollectionRequest,
-    UpdateMetadata, WrappedSerdeJsonError,
+    parse_where, Collection, CollectionConfiguration, CollectionMetadataUpdate,
+    CountCollectionsRequest, CountResponse, CreateCollectionRequest, CreateDatabaseRequest,
+    CreateTenantRequest, Database, DeleteCollectionRequest, DeleteDatabaseRequest,
+    GetCollectionRequest, GetDatabaseRequest, GetResponse, GetTenantRequest, GetTenantResponse,
+    HeartbeatError, IncludeList, InternalCollectionConfiguration, KnnIndex, ListCollectionsRequest,
+    ListDatabasesRequest, Metadata, QueryResponse, UpdateCollectionConfiguration,
+    UpdateCollectionRequest, UpdateMetadata, Where, WhereValidationError, WrappedSerdeJsonError,
 };
 use pyo3::{exceptions::PyValueError, pyclass, pyfunction, pymethods, types::PyAnyMethods, Python};
 use std::time::SystemTime;
@@ -229,16 +229,30 @@ impl Bindings {
         Ok(count)
     }
 
-    #[pyo3(signature = (limit = None, offset = 0, tenant = DEFAULT_TENANT.to_string(), database = DEFAULT_DATABASE.to_string()))]
+    #[pyo3(signature = (limit = None, offset = 0, tenant = DEFAULT_TENANT.to_string(), database = DEFAULT_DATABASE.to_string(), r#where = None))]
     fn list_collections(
         &self,
         limit: Option<u32>,
         offset: Option<u32>,
         tenant: String,
         database: String,
+        r#where: Option<String>,
     ) -> ChromaPyResult<Vec<Collection>> {
-        let request =
-            ListCollectionsRequest::try_new(tenant, database, limit, offset.unwrap_or(0))?;
+        let parsed_where: Option<Where> = match r#where {
+            Some(r#where) => {
+                let value = serde_json::from_str(r#where.as_str())
+                    .map_err(|_| WhereValidationError::WhereClause)?;
+                parse_where(&value).ok()
+            }
+            None => None,
+        };
+        let request = ListCollectionsRequest::try_new(
+            tenant,
+            database,
+            limit,
+            offset.unwrap_or(0),
+            parsed_where,
+        )?;
         let mut frontend = self.frontend.clone();
         let collections = self
             .runtime
diff --git a/rust/sqlite/migrations/sysdb/00010-collection-metadata-indices.sqlite.sql b/rust/sqlite/migrations/sysdb/00010-collection-metadata-indices.sqlite.sql
@@ -0,0 +1,4 @@
+CREATE INDEX IF NOT EXISTS collection_metadata_int_value ON collection_metadata (key, int_value) WHERE int_value IS NOT NULL;
+CREATE INDEX IF NOT EXISTS collection_metadata_float_value ON collection_metadata (key, float_value) WHERE float_value IS NOT NULL;
+CREATE INDEX IF NOT EXISTS collection_metadata_string_value ON collection_metadata (key, str_value) WHERE str_value IS NOT NULL;
+CREATE INDEX IF NOT EXISTS collection_metadata_bool_value ON collection_metadata (key, bool_value) WHERE bool_value IS NOT NULL;
diff --git a/rust/sysdb/src/sqlite.rs b/rust/sysdb/src/sqlite.rs
diff --git a/rust/sysdb/src/sysdb.rs b/rust/sysdb/src/sysdb.rs
diff --git a/rust/sysdb/src/test_sysdb.rs b/rust/sysdb/src/test_sysdb.rs
diff --git a/rust/sysdb/src/types.rs b/rust/sysdb/src/types.rs
diff --git a/rust/types/src/api_types.rs b/rust/types/src/api_types.rs
diff --git a/rust/types/src/metadata.rs b/rust/types/src/metadata.rs
