GroupBy

Group By & Aggregation is now available on Chroma Cloud.

GroupBy is useful when the unit of results you want to return is not the chunk, but instead an entity like a page, document, or product. GroupBy is also powerful paired with agentic search to quickly surface pages/documents/products the agent should take a closer look at.

As a simple example - suppose you are searching for arXiv papers about reinforcement learning and you want to return a list of strong candidate papers back to the human (or LLM).

Without GroupBy, the results may be dominated by documents from a single paper.

With GroupBy, you can ensure much greater document diversity in the results.

How it works#

By default, vector search returns the nearest neighbors based solely on similarity. With GroupBy, you organize these ranked results into groups based on metadata keys before the final results are returned. This ensures that your results are diverse across your data rather than a cluster of similar results from a single source or category.

To use GroupBy, you configure a GroupBy expression in your search. You specify the metadata keys to partition by and an aggregation function—MinK or MaxK—to select the top records within each group. Chroma then transparently performs this intra-group selection, flattens the groups, and applies your final limit.

GroupBy complements our existing query engine features - extending it to enable:

• Diversification: Prevent "result crowding" by ensuring a variety of categories or authors are represented in the top results.
• Deduplication: Use metadata keys like document_id to ensure you only retrieve the most relevant chunk from any single source.
• Multi-Key Sorting: Apply tie-breaking logic, such as sorting within groups by priority or rating alongside the vector similarity #score.

Implementation#

You can implement GroupBy using the Search API in both Python, TypeScript and Rust. Continuing the example above - here is an example pattern for diversified arXiv paper search:

from chromadb import Search, K, Knn, GroupBy, MinK

# Get the top 2 results per category, ordered by relevance
search = (Search()
    .rank(Knn(query="reinforcement_learning", limit=100))
    .group_by(GroupBy(
        keys=K("paper_id"),
        aggregate=MinK(keys=K.SCORE, k=2)
    ))
    .limit(20))

results = collection.search(search)

Get Started#

Group By & Aggregation is available now in the latest versions of our Python, JavaScript, and Rust SDKs. To learn more, please visit our technical documentation.