Repository: huggingface/dataset-viewer

Issue:
Optimize mongo query
The metrics about the cache entries is a list of tuples `(kind, http_status, error_code, count)`.

Currently we compute with:

https://github.com/huggingface/datasets-server/blob/deb708ae737a2f8da51b74c1ca4a489c4ff39b51/libs/libcommon/src/libcommon/simple_cache.py#L490-L506

These queries are very slow (eg each `entries(kind=kind, http_status=http_status).distinct("error_code")` command can take about 10s)

An alternative is to run an aggregation:

```python
def format_group(group: Dict[str, Any]) -> CountEntry:
    kind = group["kind"]
    if not isinstance(kind, str):
        raise TypeError("kind must be a str")
    http_status = group["http_status"]
    if not isinstance(http_status, int):
        raise TypeError("http_status must be an int")
    error_code = group["error_code"]
    if not isinstance(error_code, str) and error_code is not None:
        raise TypeError("error_code must be a str or None")
    count = group["count"]
    if not isinstance(count, int):
        raise TypeError("count must be an int")
    return {"kind": kind, "http_status": http_status, "error_code": error_code, "count": count}


def get_responses_count_by_kind_status_and_error_code() -> List[CountEntry]:
    groups = CachedResponseDocument.objects().aggregate(
        [
            {
                "$group": {
                    "_id": {"kind": "$kind", "http_status": "$http_status", "error_code": "$error_code"},
                    "count": {"$sum": 1},
                }
            },
            {
                "$project": {
                    "kind": "$_id.kind",
                    "http_status": "$_id.http_status",
                    "error_code": "$_id.error_code",
                    "count": "$count",
                }
            },
        ]
    )
    return [format_group(group) for group in groups]
```

but from my tests on the prod database, it does not compute quickly either (3 minutes).

Note that we have the `(kind, http_status, error_code)` index in the database:

<img width="1487" alt="Capture d’écran 2023-08-02 à 11 07 00" src="https://github.com/huggingface/datasets-server/assets/1676121/87d587e4-1051-455f-b4ec-ee9f55a47b40">



Discussion:
A list of long operations: about 10s

<img width="1592" alt="Capture d’écran 2023-08-02 à 11 13 03" src="https://github.com/huggingface/datasets-server/assets/1676121/3c91a1c6-e568-400b-a2d1-5f8c90f03604">


Resolving patch:
diff --git a/libs/libcommon/src/libcommon/simple_cache.py b/libs/libcommon/src/libcommon/simple_cache.py
--- a/libs/libcommon/src/libcommon/simple_cache.py
+++ b/libs/libcommon/src/libcommon/simple_cache.py
@@ -487,23 +487,43 @@ class CountEntry(TypedDict):
     count: int
 
 
+def format_group(group: Dict[str, Any]) -> CountEntry:
+    kind = group["kind"]
+    if not isinstance(kind, str):
+        raise TypeError("kind must be a str")
+    http_status = group["http_status"]
+    if not isinstance(http_status, int):
+        raise TypeError("http_status must be an int")
+    error_code = group["error_code"]
+    if not isinstance(error_code, str) and error_code is not None:
+        raise TypeError("error_code must be a str or None")
+    count = group["count"]
+    if not isinstance(count, int):
+        raise TypeError("count must be an int")
+    return {"kind": kind, "http_status": http_status, "error_code": error_code, "count": count}
+
+
 def get_responses_count_by_kind_status_and_error_code() -> List[CountEntry]:
-    # TODO: rework with aggregate
-    # see
-    # https://stackoverflow.com/questions/47301829/mongodb-distinct-count-for-combination-of-two-fields?noredirect=1&lq=1#comment81555081_47301829
-    # and https://docs.mongoengine.org/guide/querying.html#mongodb-aggregation-api
-    entries = CachedResponseDocument.objects().only("kind", "http_status", "error_code")
-    return [
-        {
-            "kind": str(kind),
-            "http_status": int(http_status),
-            "error_code": error_code,
-            "count": entries(kind=kind, http_status=http_status, error_code=error_code).count(),
-        }
-        for kind in sorted(entries.distinct("kind"))
-        for http_status in sorted(entries(kind=kind).distinct("http_status"))
-        for error_code in entries(kind=kind, http_status=http_status).distinct("error_code")
-    ]
+    groups = CachedResponseDocument.objects().aggregate(
+        [
+            {"$sort": {"kind": 1, "http_status": 1, "error_code": 1}},
+            {
+                "$group": {
+                    "_id": {"kind": "$kind", "http_status": "$http_status", "error_code": "$error_code"},
+                    "count": {"$sum": 1},
+                }
+            },
+            {
+                "$project": {
+                    "kind": "$_id.kind",
+                    "http_status": "$_id.http_status",
+                    "error_code": "$_id.error_code",
+                    "count": "$count",
+                }
+            },
+        ]
+    )
+    return [format_group(group) for group in groups]
 
 
 # /cache-reports/... endpoints


Test patch:
diff --git a/libs/libcommon/tests/test_simple_cache.py b/libs/libcommon/tests/test_simple_cache.py
--- a/libs/libcommon/tests/test_simple_cache.py
+++ b/libs/libcommon/tests/test_simple_cache.py
@@ -345,10 +345,10 @@ def test_count_by_status_and_error_code() -> None:
         error_code="error_code",
     )
 
-    assert get_responses_count_by_kind_status_and_error_code() == [
-        {"kind": "test_kind", "http_status": 200, "error_code": None, "count": 1},
-        {"kind": "test_kind2", "http_status": 500, "error_code": "error_code", "count": 1},
-    ]
+    metrics = get_responses_count_by_kind_status_and_error_code()
+    assert len(metrics) == 2
+    assert {"kind": "test_kind", "http_status": 200, "error_code": None, "count": 1} in metrics
+    assert {"kind": "test_kind2", "http_status": 500, "error_code": "error_code", "count": 1} in metrics
 
 
 def test_get_cache_reports() -> None: