[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 None.
             where_document: Conditional filtering on documents. Defaults to None.
@@ -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

@@ -600,6 +600,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

@@ -570,6 +570,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,
@@ -585,6 +586,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,
@@ -791,10 +792,12 @@ def _query(
         )
 
         query_amount = len(query_embeddings)
+        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=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,

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,

chromadb/test/property/test_filtering.py

@@ -19,7 +19,6 @@
 import chromadb.test.property.strategies as strategies
 import hypothesis.strategies as st
 import logging
-import random
 import re
 from chromadb.test.utils.wait_for_version_increase import wait_for_version_increase
 import numpy as np
@@ -325,6 +324,7 @@ def test_filterable_metadata_get_limit_offset(
         min_size=1,
     ),
     should_compact=st.booleans(),
+    data=st.data(),
 )
 def test_filterable_metadata_query(
     caplog: pytest.LogCaptureFixture,
@@ -333,6 +333,7 @@ def test_filterable_metadata_query(
     record_set: strategies.RecordSet,
     filters: List[strategies.Filter],
     should_compact: bool,
+    data: st.DataObject,
 ) -> None:
     caplog.set_level(logging.ERROR)
 
@@ -355,19 +356,21 @@ def test_filterable_metadata_query(
             wait_for_version_increase(client, collection.name, initial_version)  # type: ignore
 
     total_count = len(normalized_record_set["ids"])
-    # Pick a random vector
+    # Pick a random vector using Hypothesis data
     random_query: Embedding
+
+    query_index = data.draw(st.integers(min_value=0, max_value=total_count - 1))
     if collection.has_embeddings:
         assert normalized_record_set["embeddings"] is not None
         assert all(isinstance(e, list) for e in normalized_record_set["embeddings"])
-        random_query = normalized_record_set["embeddings"][
-            random.randint(0, total_count - 1)
-        ]
+        # Use data.draw to select index
+        random_query = normalized_record_set["embeddings"][query_index]
     else:
         assert isinstance(normalized_record_set["documents"], list)
         assert collection.embedding_function is not None
+        # Use data.draw to select index
         random_query = collection.embedding_function(
-            [normalized_record_set["documents"][random.randint(0, total_count - 1)]]
+            [normalized_record_set["documents"][query_index]]
         )[0]
     for filter in filters:
         result_ids = set(
@@ -402,7 +405,7 @@ def test_empty_filter(client: ClientAPI) -> None:
         query_embeddings=test_query_embedding,
         where={"q": {"$eq": 4}},  # type: ignore[dict-item]
         n_results=3,
-        include=["embeddings", "distances", "metadatas"],  # type: ignore[list-item]
+        include=["embeddings", "distances", "metadatas"],
     )
     assert res["ids"] == [[]]
     if res["embeddings"] is not None:
@@ -459,9 +462,108 @@ def check_empty_res(res: GetResult) -> None:
 
     coll.add(ids=test_ids, embeddings=test_embeddings, metadatas=test_metadatas)
 
-    res = coll.get(ids=["nope"], include=["embeddings", "metadatas", "documents"])  # type: ignore[list-item]
+    res = coll.get(ids=["nope"], include=["embeddings", "metadatas", "documents"])
     check_empty_res(res)
     res = coll.get(
-        include=["embeddings", "metadatas", "documents"], where={"test": 100}  # type: ignore[list-item]
+        include=["embeddings", "metadatas", "documents"], where={"test": 100}
     )
     check_empty_res(res)
+
+
+@settings(
+    deadline=90000,
+    suppress_health_check=[
+        HealthCheck.function_scoped_fixture,
+        HealthCheck.large_base_example,
+    ],
+)
+@given(
+    collection=collection_st,
+    record_set=recordset_st,
+    n_results_st=st.integers(min_value=1, max_value=100),
+    should_compact=st.booleans(),
+    data=st.data(),
+)
+def test_query_ids_filter_property(
+    caplog: pytest.LogCaptureFixture,
+    client: ClientAPI,
+    collection: strategies.Collection,
+    record_set: strategies.RecordSet,
+    n_results_st: int,
+    should_compact: bool,
+    data: st.DataObject,
+) -> None:
+    """Property test for querying with only the ids filter."""
+    if (
+        client.get_settings().chroma_api_impl
+        == "chromadb.api.async_fastapi.AsyncFastAPI"
+    ):
+        pytest.skip(
+            "Skipping test for async client due to potential resource/timeout issues"
+        )
+    caplog.set_level(logging.ERROR)
+    reset(client)
+    coll = client.create_collection(
+        name=collection.name,
+        metadata=collection.metadata,  # type: ignore
+        embedding_function=collection.embedding_function,
+    )
+    initial_version = coll.get_model()["version"]
+    normalized_record_set = invariants.wrap_all(record_set)
+
+    if len(normalized_record_set["ids"]) == 0:
+        # Cannot add empty record set
+        return
+
+    coll.add(**record_set)  # type: ignore[arg-type]
+
+    if not NOT_CLUSTER_ONLY:
+        if should_compact and len(normalized_record_set["ids"]) > 10:
+            wait_for_version_increase(client, collection.name, initial_version)  # type: ignore
+
+    total_count = len(normalized_record_set["ids"])
+    n_results = min(n_results_st, total_count)
+
+    # Generate a random subset of ids to filter on using Hypothesis data
+    ids_subset_size = data.draw(st.integers(min_value=0, max_value=total_count))
+    ids_to_query = data.draw(
+        st.lists(
+            st.sampled_from(normalized_record_set["ids"]),
+            min_size=ids_subset_size,
+            max_size=ids_subset_size,
+            unique=True,
+        )
+    )
+
+    # Pick a random query vector using Hypothesis data
+    random_query: Embedding
+    query_index = data.draw(st.integers(min_value=0, max_value=total_count - 1))
+    if collection.has_embeddings:
+        assert normalized_record_set["embeddings"] is not None
+        assert all(isinstance(e, list) for e in normalized_record_set["embeddings"])
+        # Use data.draw to select index
+        random_query = normalized_record_set["embeddings"][query_index]
+    else:
+        assert isinstance(normalized_record_set["documents"], list)
+        assert collection.embedding_function is not None
+        # Use data.draw to select index
+        random_query = collection.embedding_function(
+            [normalized_record_set["documents"][query_index]]
+        )[0]
+
+    # Perform the query with only the ids filter
+    result = coll.query(
+        query_embeddings=[random_query],
+        ids=ids_to_query,
+        n_results=n_results,
+    )
+
+    result_ids = set(result["ids"][0])
+    filter_ids_set = set(ids_to_query)
+
+    # The core assertion: all returned IDs must be within the filter set
+    assert result_ids.issubset(filter_ids_set)
+
+    # Also check that the number of results is reasonable
+    assert len(result_ids) <= n_results
+    assert len(result_ids) <= len(filter_ids_set)