[ENH] Add python & js client support to query on subset of IDs #4250

jairad26 · 2025-04-09T21:45:10Z

Description of changes

This PR adds python client and async python client support to query on a filtered set of IDs

example:

ids = ["1", "2", "3"]
documents = ["test", "test2", "apple"]
metadatas = [{"source": "test"}, {"source": "test2"}, {"source": "apple"}]

coll.add(
    ids=ids,
    documents=documents,
    metadatas=metadatas,
    # embeddings=numpy_embeddings
)

output = coll.query(
    ids=["1", "3"],
    query_texts=["test"],
    n_results=3,
    include=["documents", "metadatas", "distances"]
)

print(output)

This will output
% python test_filter_id.py
{'ids': [['1', '3']], 'embeddings': None, 'documents': [['test', 'apple']], 'uris': None, 'included': ['documents', 'metadatas', 'distances'], 'data': None, 'metadatas': [[{'source': 'test'}, {'source': 'apple'}]], 'distances': [[0.0, 0.7396076321601868]]}

Test plan

How are these changes tested?

Tests pass locally with pytest for python, yarn test for js, cargo test for rust
Added prop tests to test with other filtering and on its own

Documentation Changes

Are all docstrings for user-facing APIs updated if required? Do we need to make documentation changes in the docs repository?

github-actions · 2025-04-09T21:45:20Z

jairad26 · 2025-04-09T21:45:31Z

This stack of pull requests is managed by Graphite. Learn more about stacking.

hesreallyhim · 2025-04-14T14:51:24Z

chromadb/api/__init__.py

@@ -270,6 +270,7 @@ def _query(
        self,
        collection_id: UUID,
        query_embeddings: Embeddings,
+        ids: Optional[IDs] = None,


add to docstring(?)

added, thanks

hesreallyhim · 2025-04-14T15:12:41Z

chromadb/api/rust.py

@@ -472,6 +472,7 @@ def _query(
        self,
        collection_id: UUID,
        query_embeddings: Embeddings,
+        ids: Optional[IDs] = None,


add to telemetry capture call below? just observing the pattern

added, thanks

hesreallyhim · 2025-04-14T17:05:49Z

chromadb/api/async_api.py

@@ -264,6 +264,7 @@ async def _query(
        self,
        collection_id: UUID,
        query_embeddings: Embeddings,
+        ids: Optional[IDs] = None,


docstring(?)

hesreallyhim · 2025-04-14T17:12:55Z

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,


what do you think about the OneOrMany pattern here?

not a fan, because why would someone want to filter on 1 id? it makes it a little more confusing on what it does. also the rust takes in a list of ids, so i'd prefer maintaining that.

yeah that's very reasonable

I think we should keep consistency with other APIs

hesreallyhim · 2025-04-15T13:47:15Z

i think i'm missing something, so for my edification, which part of the code actually processes the ids filter? Is it in here? https://github.com/hesreallyhim/chroma/blob/cac7f29be72907e71640981e7cf0dc94c68d06ce/rust/frontend/src/impls/in_memory_frontend.rs#L609

EDIT: Sorry that might be a really broad question, I don't know how the rust code works.

hesreallyhim · 2025-04-15T14:13:06Z

it would be cool to add filter syntax to ids, like {"$and": [{"$gt": 1000}, {"$lt": 2000}]} but i think it doesn't really make sense since we don't allow comparison like that for strings anyway, and IDs are strings...

hesreallyhim · 2025-04-15T14:16:35Z

@jairad26 would it be helpful for me to put together a PR to add this feature to the docs?

jairad26 · 2025-04-15T16:18:25Z

i think i'm missing something, so for my edification, which part of the code actually processes the ids filter? Is it in here? https://github.com/hesreallyhim/chroma/blob/cac7f29be72907e71640981e7cf0dc94c68d06ce/rust/frontend/src/impls/in_memory_frontend.rs#L609

EDIT: Sorry that might be a really broad question, I don't know how the rust code works.

yep that's exactly the spot! all that does is filter by ids pre-vector search.

it would be cool to add filter syntax to ids, like {"$and": [{"$gt": 1000}, {"$lt": 2000}]} but i think it doesn't really make sense since we don't allow comparison like that for strings anyway, and IDs are strings...

yes that would be nice, but would only really work if the ids are ints. not 100% sure but i dont think uuids for example guarantee order like that. also random string generator definitely wouldn't work with that

@jairad26 would it be helpful for me to put together a PR to add this feature to the docs?

I'm moving this PR to draft for now, since i want to add support for the js client and that requires a couple other PRs to get merged. after that I can prob handle the docs. Thank you though!

HammadB · 2025-04-24T16:24:08Z

chromadb/test/property/test_filtering.py

@@ -19,7 +19,6 @@
 import chromadb.test.property.strategies as strategies


please add explicit unit tests as well with small, medium size data

HammadB

Generally is good. Please add more explicit unit testing, especially around edge cases like deleted ids, upserted ids.

Also I think OneOrMany is preferable for API consistency, even though unlikely to be used.

HammadB · 2025-04-29T16:24:15Z

chromadb/test/property/test_filtering.py

-        random_query = normalized_record_set["embeddings"][
-            random.randint(0, total_count - 1)
-        ]
+        # Use data.draw to select index


thanks for this change

HammadB · 2025-04-29T16:24:39Z

chromadb/test/property/test_filtering.py

+        client.get_settings().chroma_api_impl
+        == "chromadb.api.async_fastapi.AsyncFastAPI"
+    ):
+        pytest.skip(


aside - when can we remove this?

i haven’t been able to find a root cause, will set aside a couple hours for it this week

HammadB · 2025-04-29T16:27:55Z

chromadb/test/property/test_filtering.py

+    ids_to_query = data.draw(
+        st.lists(
+            st.sampled_from(normalized_record_set["ids"]),
+            min_size=ids_subset_size,


why choose id_subset_size when min_size and max_size do the same thing? This introduces more entropy requirements.

ah yea ok, will make it create more diverse sets of data

No I think its the same thing is my point. St.lists() choosing between min and max size is equivalent to data.draw(st.integers(min_value=0, max_value=total_count))

oh got it okay. removing id_subset_size.

it's this now

ids_to_query = data.draw( st.lists( st.sampled_from(normalized_record_set["ids"]), min_size=0, max_size=total_count, unique=True, ) )

…a-core#4250)

jairad26 mentioned this pull request Apr 9, 2025

[CHORE] remove unused page, page_size, and sort args on get #4204

Merged

1 task

jairad26 mentioned this pull request Apr 9, 2025

[CHORE] Propogate error messages correctly to user #4235

Merged

1 task

jairad26 force-pushed the jai/query-on-doc-ids branch from 89174d3 to 3323962 Compare April 9, 2025 22:24

jairad26 marked this pull request as ready for review April 9, 2025 22:24

jairad26 force-pushed the jai/query-on-doc-ids branch 6 times, most recently from 699620e to 4a0795f Compare April 10, 2025 02:06

jairad26 changed the base branch from jai/propogate-api-ef-error to main April 10, 2025 13:54

jairad26 force-pushed the jai/query-on-doc-ids branch 5 times, most recently from 49df6bb to 194e27b Compare April 11, 2025 21:06

hesreallyhim reviewed Apr 14, 2025

View reviewed changes

jairad26 force-pushed the jai/query-on-doc-ids branch 2 times, most recently from ddb0949 to cac7f29 Compare April 14, 2025 16:56

hesreallyhim reviewed Apr 14, 2025

View reviewed changes

jairad26 marked this pull request as draft April 15, 2025 16:19

jairad26 force-pushed the jai/query-on-doc-ids branch from cac7f29 to 5a95f79 Compare April 15, 2025 18:57

jairad26 force-pushed the jai/query-on-doc-ids branch from 74c9b85 to 99a1335 Compare April 21, 2025 22:55

jairad26 mentioned this pull request Apr 21, 2025

[CHORE] Fix type safety #4339

Merged

1 task

jairad26 changed the title ~~[ENH] Add python &js client support to query on subset of IDs~~ [ENH] Add python & js client support to query on subset of IDs Apr 21, 2025

jairad26 force-pushed the jai/query-on-doc-ids branch from 99a1335 to 2239c58 Compare April 23, 2025 23:43

jairad26 mentioned this pull request Apr 23, 2025

[ENH] Add close on python clients #4367

Draft

1 task

HammadB reviewed Apr 24, 2025

View reviewed changes

HammadB requested changes Apr 24, 2025

View reviewed changes

jairad26 force-pushed the jai/query-on-doc-ids branch 8 times, most recently from 621b98a to fc2f354 Compare April 25, 2025 23:41

jairad26 mentioned this pull request Apr 25, 2025

[ENH] Add Cloudflare Worker AI Embedding Function #4389

Merged

1 task

jairad26 force-pushed the jai/query-on-doc-ids branch from fc2f354 to bb6c085 Compare April 26, 2025 00:02

hesreallyhim mentioned this pull request Apr 28, 2025

Example on how to filter a collection over id's amikos-tech/chroma-cookbook#18

Open

jairad26 force-pushed the jai/query-on-doc-ids branch 2 times, most recently from 907ec35 to b4d8d7b Compare April 28, 2025 21:00

HammadB reviewed Apr 29, 2025

View reviewed changes

HammadB approved these changes Apr 29, 2025

View reviewed changes

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

03ccf12

jairad26 force-pushed the jai/query-on-doc-ids branch from b4d8d7b to 03ccf12 Compare April 29, 2025 17:34

jairad26 enabled auto-merge (squash) April 29, 2025 17:36

jairad26 merged commit 827f0f9 into main Apr 29, 2025
70 checks passed

warpbuild-benchmark-bot bot added a commit to WarpBuilds/chroma that referenced this pull request Apr 29, 2025

[ENH] Add python & js client support to query on subset of IDs (chrom…

dc975f7

…a-core#4250)

jairad26 mentioned this pull request Apr 29, 2025

[Feature Request]: Vector query on specific set of documents ids #3941

Closed

		@@ -19,7 +19,6 @@
		import chromadb.test.property.strategies as strategies

[ENH] Add python & js client support to query on subset of IDs #4250

[ENH] Add python & js client support to query on subset of IDs #4250

Uh oh!

Conversation

jairad26 commented Apr 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description of changes

Test plan

Documentation Changes

Uh oh!

github-actions bot commented Apr 9, 2025

Reviewer Checklist

Testing, Bugs, Errors, Logs, Documentation

System Compatibility

Quality

Uh oh!

jairad26 commented Apr 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

hesreallyhim commented Apr 15, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

hesreallyhim commented Apr 15, 2025

Uh oh!

hesreallyhim commented Apr 15, 2025

Uh oh!

jairad26 commented Apr 15, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

HammadB left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

jairad26 commented Apr 9, 2025 •

edited

Loading

jairad26 commented Apr 9, 2025 •

edited

Loading

hesreallyhim commented Apr 15, 2025 •

edited

Loading