[ENH]: Force statistics wrapper callers to specify an output collection

Author: tanujnay112

chromadb/test/distributed/test_statistics_wrapper.py

@@ -4,6 +4,8 @@
 
 import json
 import time
+from typing import Any
+
 from chromadb.api.client import Client as ClientCreator
 from chromadb.base_types import SparseVector
 from chromadb.config import System
@@ -31,7 +33,7 @@ def test_statistics_wrapper(basic_http_client: System) -> None:
     )
 
     # Enable statistics
-    attached_fn = attach_statistics_function(collection)
+    attached_fn = attach_statistics_function(collection, "test_collection_statistics")
     assert attached_fn is not None
     assert attached_fn.function_name == "statistics"
     assert attached_fn.output_collection == "test_collection_statistics"
@@ -54,7 +56,7 @@ def test_statistics_wrapper(basic_http_client: System) -> None:
     time.sleep(60)
 
     # Get statistics
-    stats = get_statistics(collection)
+    stats = get_statistics(collection, "test_collection_statistics")
     print("\nStatistics output:")
     print(json.dumps(stats, indent=2))
 
@@ -124,14 +126,14 @@ def test_backfill_statistics(basic_http_client: System) -> None:
     initial_version = get_collection_version(client, collection.name)
 
     # Enable statistics
-    attached_fn = attach_statistics_function(collection)
+    attached_fn = attach_statistics_function(collection, "my_collection_statistics")
     assert attached_fn.function_name == "statistics"
     assert attached_fn.output_collection == "my_collection_statistics"
 
     # Wait for statistics to be computed
     wait_for_version_increase(client, collection.name, initial_version)
 
-    stats = get_statistics(collection)
+    stats = get_statistics(collection, "my_collection_statistics")
     assert stats is not None
     assert "statistics" in stats
     assert "summary" in stats
@@ -190,7 +192,7 @@ def test_statistics_wrapper_custom_output_collection(basic_http_client: System)
     wait_for_version_increase(client, collection.name, initial_version)
 
     # Get statistics
-    stats = get_statistics(collection)
+    stats = get_statistics(collection, "my_custom_stats")
     assert "statistics" in stats
     assert "key" in stats["statistics"]
 
@@ -206,7 +208,7 @@ def test_statistics_wrapper_key_filter(basic_http_client: System) -> None:
     collection = client.create_collection(name="key_filter_test")
 
     # Enable statistics
-    attach_statistics_function(collection)
+    attach_statistics_function(collection, "key_filter_test_statistics")
 
     initial_version = get_collection_version(client, collection.name)
 
@@ -225,13 +227,15 @@ def test_statistics_wrapper_key_filter(basic_http_client: System) -> None:
     time.sleep(60)
 
     # Get all statistics (no key filter)
-    all_stats = get_statistics(collection)
+    all_stats = get_statistics(collection, "key_filter_test_statistics")
     assert "category" in all_stats["statistics"]
     assert "score" in all_stats["statistics"]
     assert "active" in all_stats["statistics"]
 
     # Get statistics filtered by "category" key only
-    category_stats = get_statistics(collection, key="category")
+    category_stats = get_statistics(
+        collection, "key_filter_test_statistics", key="category"
+    )
     assert "category" in category_stats["statistics"]
     assert "score" not in category_stats["statistics"]
     assert "active" not in category_stats["statistics"]
@@ -242,7 +246,7 @@ 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="score")
+    score_stats = get_statistics(collection, "key_filter_test_statistics", key="score")
     assert "score" in score_stats["statistics"]
     assert "category" not in score_stats["statistics"]
     assert "active" not in score_stats["statistics"]
@@ -263,7 +267,7 @@ def test_statistics_wrapper_incremental_updates(basic_http_client: System) -> No
     client.reset()
 
     collection = client.create_collection(name="incremental_test")
-    attach_statistics_function(collection)
+    attach_statistics_function(collection, "incremental_test_statistics")
 
     initial_version = get_collection_version(client, collection.name)
 
@@ -278,7 +282,7 @@ def test_statistics_wrapper_incremental_updates(basic_http_client: System) -> No
     next_version = get_collection_version(client, collection.name)
 
     # Check initial statistics
-    stats = get_statistics(collection)
+    stats = get_statistics(collection, "incremental_test_statistics")
     assert stats["statistics"]["category"]["A"]["count"] == 2
     assert stats["summary"]["total_count"] == 2
 
@@ -295,7 +299,7 @@ def test_statistics_wrapper_incremental_updates(basic_http_client: System) -> No
     time.sleep(70)
 
     # Check updated statistics
-    stats = get_statistics(collection)
+    stats = get_statistics(collection, "incremental_test_statistics")
     assert stats["statistics"]["category"]["A"]["count"] == 3
     assert stats["statistics"]["category"]["B"]["count"] == 1
     assert stats["summary"]["total_count"] == 4
@@ -333,14 +337,14 @@ def test_sparse_vector_statistics(basic_http_client: System) -> None:
             {"category": "A", "vec": sparse_vec3},
         ],
     )
-    attach_statistics_function(collection)
+    attach_statistics_function(collection, "sparse_vector_test1_statistics")
 
     initial_version = get_collection_version(client, collection.name)
 
     wait_for_version_increase(client, collection.name, initial_version)
 
     # Get statistics
-    stats = get_statistics(collection)
+    stats = get_statistics(collection, "sparse_vector_test1_statistics")
     print("\nSparse vector statistics output:")
     print(json.dumps(stats, indent=2))
 
@@ -385,54 +389,54 @@ def test_statistics_high_cardinality(basic_http_client: System) -> None:
     num_fields = 10
     ids = [f"id{i}" for i in range(num_docs)]
     documents = [f"doc{i}" for i in range(num_docs)]
-    
-    metadatas = []
+
+    metadatas: list[dict[str, Any]] = []
     for i in range(num_docs):
-        meta = {}
+        meta: dict[str, Any] = {}
         for j in range(num_fields):
             meta[f"field_{j}"] = f"value_{j}_{i}"
         metadatas.append(meta)
 
     # Add in batches to avoid hitting request size limits
     batch_size = 100
     initial_version = get_collection_version(client, collection.name)
-    
+
     for i in range(0, num_docs, batch_size):
         collection.add(
             ids=ids[i : i + batch_size],
             documents=documents[i : i + batch_size],
-            metadatas=metadatas[i : i + batch_size],
+            metadatas=metadatas[i : i + batch_size],  # type: ignore[arg-type]
         )
 
     # Let all data be compacted
     wait_for_version_increase(client, collection.name, initial_version)
     initial_version = get_collection_version(client, collection.name)
 
     # Enable statistics
-    attach_statistics_function(collection)
+    attach_statistics_function(collection, "high_cardinality_test_statistics")
 
     # Wait for statistics to be computed
     wait_for_version_increase(client, collection.name, initial_version)
 
     # Get statistics
-    stats = get_statistics(collection)
-    
+    stats = get_statistics(collection, "high_cardinality_test_statistics")
+
     assert "statistics" in stats
-    
+
     # Verify we have stats for all fields
     for j in range(num_fields):
         field_key = f"field_{j}"
         assert field_key in stats["statistics"]
-        
+
         field_stats = stats["statistics"][field_key]
         assert len(field_stats) == num_docs
-        
+
         # Verify each value has count 1
         for i in range(num_docs):
             value = f"value_{j}_{i}"
             assert value in field_stats
             assert field_stats[value]["count"] == 1
-        
+
     # Verify total count
     assert stats["summary"]["total_count"] == num_docs
 

chromadb/utils/statistics.py

@@ -11,8 +11,8 @@
     >>> client = chromadb.Client()
     >>> collection = client.get_or_create_collection("my_collection")
     >>>
-    >>> # Attach statistics function
-    >>> attach_statistics_function(collection)
+    >>> # Attach statistics function with output collection name
+    >>> attach_statistics_function(collection, "my_collection_statistics")
     >>>
     >>> # Add some data
     >>> collection.add(
@@ -21,8 +21,8 @@
     ...     metadatas=[{"category": "A"}, {"category": "B"}]
     ... )
     >>>
-    >>> # Get statistics
-    >>> stats = get_statistics(collection)
+    >>> # Get statistics from the named output collection
+    >>> stats = get_statistics(collection, "my_collection_statistics")
     >>> print(stats)
 """
 
@@ -49,7 +49,7 @@ def get_statistics_fn_name(collection: "Collection") -> str:
 
 
 def attach_statistics_function(
-    collection: "Collection", stats_collection_name: Optional[str] = None
+    collection: "Collection", stats_collection_name: str
 ) -> "AttachedFunction":
     """Attach statistics collection function to a collection.
 
@@ -60,20 +60,16 @@ def attach_statistics_function(
     Args:
         collection: The collection to enable statistics for
         stats_collection_name: Name of the collection where statistics will be stored.
-                               If None, defaults to "{collection_name}_statistics".
 
     Returns:
         AttachedFunction: The attached statistics function
 
     Example:
-        >>> attach_statistics_function(collection)
+        >>> attach_statistics_function(collection, "my_collection_statistics")
         >>> collection.add(ids=["id1"], documents=["doc1"], metadatas=[{"key": "value"}])
         >>> # Statistics are automatically computed
-        >>> stats = get_statistics(collection)
+        >>> stats = get_statistics(collection, "my_collection_statistics")
     """
-    if stats_collection_name is None:
-        stats_collection_name = f"{collection.name}_statistics"
-
     return collection.attach_function(
         name=get_statistics_fn_name(collection),
         function_id="statistics",
@@ -125,7 +121,7 @@ def detach_statistics_function(
 
 
 def get_statistics(
-    collection: "Collection", key: Optional[str] = None
+    collection: "Collection", stats_collection_name: str, key: Optional[str] = None
 ) -> Dict[str, Any]:
     """Get the current statistics for a collection.
 
@@ -134,6 +130,7 @@ 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.
 
@@ -154,14 +151,14 @@ def get_statistics(
             }
 
     Example:
-        >>> attach_statistics_function(collection)
+        >>> attach_statistics_function(collection, "my_collection_statistics")
         >>> collection.add(
         ...     ids=["id1", "id2"],
         ...     documents=["doc1", "doc2"],
         ...     metadatas=[{"category": "A", "score": 10}, {"category": "B", "score": 10}]
         ... )
         >>> # Wait for statistics to be computed
-        >>> stats = get_statistics(collection)
+        >>> stats = get_statistics(collection, "my_collection_statistics")
         >>> print(stats)
         {
             "statistics": {
@@ -181,11 +178,9 @@ def get_statistics(
     # Import here to avoid circular dependency
     from chromadb.api.models.Collection import Collection
 
-    af = get_statistics_fn(collection)
-
     # Get the statistics output collection model from the server
     stats_collection_model = collection._client.get_collection(
-        name=af.output_collection,
+        name=stats_collection_name,
         tenant=collection.tenant,
         database=collection.database,
     )