[ENH] Add python client support to query on subset of IDs

Author: jairad26

chromadb/api/__init__.py

@@ -270,6 +270,7 @@ def _query(
         self,
         collection_id: UUID,
         query_embeddings: Embeddings,
+        ids: Optional[IDs] = None,
         n_results: int = 10,
         where: Optional[Where] = None,
         where_document: Optional[WhereDocument] = None,
@@ -280,6 +281,7 @@ def _query(
         Args:
             collection_id: The UUID of the collection to query.
             query_embeddings: The embeddings to use as the query.
+            ids: The IDs to filter by during the query. Defaults to None.
             n_results: The number of results to return. Defaults to 10.
             where: Conditional filtering on metadata. Defaults to {}.
             where_document: Conditional filtering on documents. Defaults to {}.
@@ -724,6 +726,7 @@ def _query(
         self,
         collection_id: UUID,
         query_embeddings: Embeddings,
+        ids: Optional[IDs] = None,
         n_results: int = 10,
         where: Optional[Where] = None,
         where_document: Optional[WhereDocument] = None,

chromadb/api/async_api.py

@@ -264,6 +264,7 @@ async def _query(
         self,
         collection_id: UUID,
         query_embeddings: Embeddings,
+        ids: Optional[IDs] = None,
         n_results: int = 10,
         where: Optional[Where] = None,
         where_document: Optional[WhereDocument] = None,
@@ -718,6 +719,7 @@ async def _query(
         self,
         collection_id: UUID,
         query_embeddings: Embeddings,
+        ids: Optional[IDs] = None,
         n_results: int = 10,
         where: Optional[Where] = None,
         where_document: Optional[WhereDocument] = None,

chromadb/api/async_client.py

@@ -403,6 +403,7 @@ async def _query(
         self,
         collection_id: UUID,
         query_embeddings: Embeddings,
+        ids: Optional[IDs] = None,
         n_results: int = 10,
         where: Optional[Where] = None,
         where_document: Optional[WhereDocument] = None,
@@ -411,6 +412,7 @@ async def _query(
         return await self._server._query(
             collection_id=collection_id,
             query_embeddings=query_embeddings,
+            ids=ids,
             n_results=n_results,
             where=where,
             where_document=where_document,

chromadb/api/async_fastapi.py

@@ -623,6 +623,7 @@ async def _query(
         self,
         collection_id: UUID,
         query_embeddings: Embeddings,
+        ids: Optional[IDs] = None,
         n_results: int = 10,
         where: Optional[Where] = None,
         where_document: Optional[WhereDocument] = None,

chromadb/api/client.py

@@ -374,13 +374,15 @@ def _query(
         self,
         collection_id: UUID,
         query_embeddings: Embeddings,
+        ids: Optional[IDs] = None,
         n_results: int = 10,
         where: Optional[Where] = None,
         where_document: Optional[WhereDocument] = None,
         include: Include = IncludeMetadataDocumentsDistances,
     ) -> QueryResult:
         return self._server._query(
             collection_id=collection_id,
+            ids=ids,
             tenant=self.tenant,
             database=self.database,
             query_embeddings=query_embeddings,

chromadb/api/fastapi.py

@@ -591,6 +591,7 @@ def _query(
         self,
         collection_id: UUID,
         query_embeddings: Embeddings,
+        ids: Optional[IDs] = None,
         n_results: int = 10,
         where: Optional[Where] = None,
         where_document: Optional[WhereDocument] = None,
@@ -606,6 +607,7 @@ def _query(
             "post",
             f"/tenants/{tenant}/databases/{database}/collections/{collection_id}/query",
             json={
+                "ids": ids,
                 "query_embeddings": convert_np_embeddings_to_list(query_embeddings)
                 if query_embeddings is not None
                 else None,

chromadb/api/models/AsyncCollection.py

@@ -169,6 +169,7 @@ async def query(
         query_texts: Optional[OneOrMany[Document]] = None,
         query_images: Optional[OneOrMany[Image]] = None,
         query_uris: Optional[OneOrMany[URI]] = None,
+        ids: Optional[IDs] = None,
         n_results: int = 10,
         where: Optional[Where] = None,
         where_document: Optional[WhereDocument] = None,
@@ -184,6 +185,7 @@ async def query(
             query_embeddings: The embeddings to get the closes neighbors of. Optional.
             query_texts: The document texts to get the closes neighbors of. Optional.
             query_images: The images to get the closes neighbors of. Optional.
+            ids: A subset of ids to search within. Optional.
             n_results: The number of neighbors to return for each query_embedding or query_texts. Optional.
             where: A Where type dict used to filter results by. E.g. `{"$and": [{"color" : "red"}, {"price": {"$gte": 4.20}}]}`. Optional.
             where_document: A WhereDocument type dict used to filter by the documents. E.g. `{"$contains": "hello"}`. Optional.
@@ -213,6 +215,7 @@ async def query(
 
         query_results = await self._client._query(
             collection_id=self.id,
+            ids=ids,
             query_embeddings=query_request["embeddings"],
             n_results=query_request["n_results"],
             where=query_request["where"],

chromadb/api/models/Collection.py

@@ -172,6 +172,7 @@ def query(
         query_texts: Optional[OneOrMany[Document]] = None,
         query_images: Optional[OneOrMany[Image]] = None,
         query_uris: Optional[OneOrMany[URI]] = None,
+        ids: Optional[IDs] = None,
         n_results: int = 10,
         where: Optional[Where] = None,
         where_document: Optional[WhereDocument] = None,
@@ -188,6 +189,7 @@ def query(
             query_texts: The document texts to get the closes neighbors of. Optional.
             query_images: The images to get the closes neighbors of. Optional.
             query_uris: The URIs to be used with data loader. Optional.
+            ids: A subset of ids to search within. Optional.
             n_results: The number of neighbors to return for each query_embedding or query_texts. Optional.
             where: A Where type dict used to filter results by. E.g. `{"$and": [{"color" : "red"}, {"price": {"$gte": 4.20}}]}`. Optional.
             where_document: A WhereDocument type dict used to filter by the documents. E.g. `{"$contains": "hello"}`. Optional.
@@ -217,6 +219,7 @@ def query(
 
         query_results = self._client._query(
             collection_id=self.id,
+            ids=ids,
             query_embeddings=query_request["embeddings"],
             n_results=query_request["n_results"],
             where=query_request["where"],

chromadb/api/rust.py

@@ -472,6 +472,7 @@ def _query(
         self,
         collection_id: UUID,
         query_embeddings: Embeddings,
+        ids: Optional[IDs] = None,
         n_results: int = 10,
         where: Optional[Where] = None,
         where_document: Optional[WhereDocument] = None,
@@ -480,10 +481,12 @@ def _query(
         database: str = DEFAULT_DATABASE,
     ) -> QueryResult:
         query_amount = len(query_embeddings)
+        filtered_ids_amount = len(ids) if ids else 0
         self.product_telemetry_client.capture(
             CollectionQueryEvent(
                 collection_uuid=str(collection_id),
                 query_amount=query_amount,
+                filtered_ids_amount=filtered_ids_amount,
                 n_results=n_results,
                 with_metadata_filter=query_amount if where is not None else 0,
                 with_document_filter=query_amount if where_document is not None else 0,
@@ -496,6 +499,7 @@ def _query(
 
         rust_response = self.bindings.query(
             str(collection_id),
+            ids,
             query_embeddings,
             n_results,
             json.dumps(where) if where else None,

chromadb/api/segment.py

@@ -775,6 +775,7 @@ def _query(
         self,
         collection_id: UUID,
         query_embeddings: Embeddings,
+        ids: Optional[IDs] = None,
         n_results: int = 10,
         where: Optional[Where] = None,
         where_document: Optional[WhereDocument] = None,

chromadb/telemetry/product/events.py

@@ -137,6 +137,7 @@ class CollectionQueryEvent(ProductTelemetryEvent):
     batch_size: int
     collection_uuid: str
     query_amount: int
+    filtered_ids_amount: int
     with_metadata_filter: int
     with_document_filter: int
     n_results: int
@@ -149,6 +150,7 @@ def __init__(
         self,
         collection_uuid: str,
         query_amount: int,
+        filtered_ids_amount: int,
         with_metadata_filter: int,
         with_document_filter: int,
         n_results: int,
@@ -161,6 +163,7 @@ def __init__(
         super().__init__()
         self.collection_uuid = collection_uuid
         self.query_amount = query_amount
+        self.filtered_ids_amount = filtered_ids_amount
         self.with_metadata_filter = with_metadata_filter
         self.with_document_filter = with_document_filter
         self.n_results = n_results
@@ -182,6 +185,7 @@ def batch(self, other: "ProductTelemetryEvent") -> "CollectionQueryEvent":
         return CollectionQueryEvent(
             collection_uuid=self.collection_uuid,
             query_amount=total_amount,
+            filtered_ids_amount=self.filtered_ids_amount + other.filtered_ids_amount,
             with_metadata_filter=self.with_metadata_filter + other.with_metadata_filter,
             with_document_filter=self.with_document_filter + other.with_document_filter,
             n_results=self.n_results + other.n_results,