[ENH]: Allow specifiying multiple filter keys in get_statistics

Author: tanujnay112

chromadb/test/distributed/test_statistics_wrapper.py

@@ -234,7 +234,7 @@ def test_statistics_wrapper_key_filter(basic_http_client: System) -> None:
 
     # Get statistics filtered by "category" key only
     category_stats = get_statistics(
-        collection, "key_filter_test_statistics", key="category"
+        collection, "key_filter_test_statistics", keys=["category"]
     )
     assert "category" in category_stats["statistics"]
     assert "score" not in category_stats["statistics"]
@@ -246,7 +246,9 @@ def test_statistics_wrapper_key_filter(basic_http_client: System) -> None:
     assert category_stats["summary"]["total_count"] == 3
 
     # Get statistics filtered by "score" key only
-    score_stats = get_statistics(collection, "key_filter_test_statistics", key="score")
+    score_stats = get_statistics(
+        collection, "key_filter_test_statistics", keys=["score"]
+    )
     assert "score" in score_stats["statistics"]
     assert "category" not in score_stats["statistics"]
     assert "active" not in score_stats["statistics"]

chromadb/utils/statistics.py

@@ -26,7 +26,7 @@
     >>> print(stats)
 """
 
-from typing import TYPE_CHECKING, Optional, Dict, Any, cast
+from typing import TYPE_CHECKING, Optional, Dict, Any, List, cast
 from collections import defaultdict
 
 from chromadb.api.types import Where
@@ -121,7 +121,9 @@ def detach_statistics_function(
 
 
 def get_statistics(
-    collection: "Collection", stats_collection_name: str, key: Optional[str] = None
+    collection: "Collection",
+    stats_collection_name: str,
+    keys: Optional[List[str]] = None,
 ) -> Dict[str, Any]:
     """Get the current statistics for a collection.
 
@@ -131,8 +133,8 @@ def get_statistics(
     Args:
         collection: The collection to get statistics for
         stats_collection_name: Name of the statistics collection to read from.
-        key: Optional metadata key to filter statistics for. If provided,
-             only returns statistics for that specific key.
+        keys: Optional list of metadata keys to filter statistics for. If provided,
+              only returns statistics for those specific keys.
 
     Returns:
         Dict[str, Any]: A dictionary with the structure:
@@ -198,11 +200,9 @@ def get_statistics(
     summary: Dict[str, Any] = {}
 
     offset = 0
-    # When filtering by key, also include "summary" entries to get total_count
+    # When filtering by keys, also include "summary" entries to get total_count
     where_filter: Optional[Where] = (
-        cast(Where, {"$or": [{"key": key}, {"key": "summary"}]})
-        if key is not None
-        else None
+        cast(Where, {"key": {"$in": keys + ["summary"]}}) if keys is not None else None
     )
 
     while True: