[PERF][IMPROVEMENT] argilla server: improve computation for dataset progress and metrics (#5618)

# Description
<!-- Please include a summary of the changes and the related issue.
Please also include relevant motivation and context. List any
dependencies that are required for this change. -->

Computing dataset progress and metrics using the search engine
drastically reduces the time required to run the Argilla server in HF
spaces when persistent storage is enabled.

## TODO
- [x] Add search engine tests
- [x] Adapt existing test to mock engine results

**Type of change**
<!-- Please delete options that are not relevant. Remember to title the
PR according to the type of change -->

- Refactor (change restructuring the codebase without changing
functionality)
- Improvement (change adding some improvement to an existing
functionality)

**How Has This Been Tested**
<!-- Please add some reference about how your feature has been tested.
-->

**Checklist**
<!-- Please go over the list and make sure you've taken everything into
account -->

- I added relevant documentation
- I followed the style guidelines of this project
- I did a self-review of my code
- I made corresponding changes to the documentation
- I confirm My changes generate no new warnings
- I have added tests that prove my fix is effective or that my feature
works
- I have added relevant notes to the CHANGELOG.md file (See
https://keepachangelog.com/)

Author: frascuchon

argilla-server/src/argilla_server/api/handlers/v1/datasets/datasets.py

@@ -143,37 +143,44 @@ async def get_dataset(
 @router.get("/me/datasets/{dataset_id}/metrics", response_model=DatasetMetrics)
 async def get_current_user_dataset_metrics(
     *,
-    db: AsyncSession = Depends(get_async_db),
     dataset_id: UUID,
+    db: AsyncSession = Depends(get_async_db),
+    search_engine: SearchEngine = Depends(get_search_engine),
     current_user: User = Security(auth.get_current_user),
 ):
     dataset = await Dataset.get_or_raise(db, dataset_id)
 
     await authorize(current_user, DatasetPolicy.get(dataset))
 
-    return await datasets.get_user_dataset_metrics(db, current_user.id, dataset.id)
+    result = await datasets.get_user_dataset_metrics(db, search_engine, current_user, dataset)
+
+    return DatasetMetrics(responses=result)
 
 
 @router.get("/datasets/{dataset_id}/progress", response_model=DatasetProgress)
 async def get_dataset_progress(
     *,
-    db: AsyncSession = Depends(get_async_db),
     dataset_id: UUID,
+    db: AsyncSession = Depends(get_async_db),
+    search_engine: SearchEngine = Depends(get_search_engine),
     current_user: User = Security(auth.get_current_user),
 ):
     dataset = await Dataset.get_or_raise(db, dataset_id)
 
     await authorize(current_user, DatasetPolicy.get(dataset))
 
-    return await datasets.get_dataset_progress(db, dataset.id)
+    result = await datasets.get_dataset_progress(search_engine, dataset)
+
+    return DatasetProgress(**result)
 
 
 @router.get("/datasets/{dataset_id}/users/progress", response_model=UsersProgress)
 async def get_dataset_users_progress(
     *,
-    current_user: User = Security(auth.get_current_user),
     dataset_id: UUID,
     db: AsyncSession = Depends(get_async_db),
+    search_engine: SearchEngine = Depends(get_search_engine),
+    current_user: User = Security(auth.get_current_user),
 ):
     dataset = await Dataset.get_or_raise(db, dataset_id)
 

argilla-server/src/argilla_server/contexts/datasets.py

@@ -12,12 +12,8 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-import asyncio
 import copy
 from collections import defaultdict
-
-import sqlalchemy
-
 from datetime import datetime
 from typing import (
     TYPE_CHECKING,
@@ -35,6 +31,8 @@
     Union,
 )
 from uuid import UUID
+
+import sqlalchemy
 from fastapi.encoders import jsonable_encoder
 from sqlalchemy import Select, and_, func, select
 from sqlalchemy.ext.asyncio import AsyncSession
@@ -62,7 +60,7 @@
 from argilla_server.api.schemas.v1.vectors import Vector as VectorSchema
 from argilla_server.contexts import accounts, distribution
 from argilla_server.database import get_async_db
-from argilla_server.enums import DatasetStatus, UserRole, RecordStatus
+from argilla_server.enums import DatasetStatus, UserRole
 from argilla_server.errors.future import NotUniqueError, UnprocessableEntityError
 from argilla_server.jobs import dataset_jobs
 from argilla_server.models import (
@@ -72,7 +70,6 @@
     Question,
     Record,
     Response,
-    ResponseStatus,
     Suggestion,
     User,
     Vector,
@@ -377,88 +374,38 @@ async def _configure_query_relationships(
     return query
 
 
-async def get_user_dataset_metrics(db: AsyncSession, user_id: UUID, dataset_id: UUID) -> dict:
-    responses_submitted, responses_discarded, responses_draft, responses_pending = await asyncio.gather(
-        db.execute(
-            select(func.count(Response.id))
-            .join(Record, and_(Record.id == Response.record_id, Record.dataset_id == dataset_id))
-            .filter(
-                Response.user_id == user_id,
-                Response.status == ResponseStatus.submitted,
-            ),
-        ),
-        db.execute(
-            select(func.count(Response.id))
-            .join(Record, and_(Record.id == Response.record_id, Record.dataset_id == dataset_id))
-            .filter(
-                Response.user_id == user_id,
-                Response.status == ResponseStatus.discarded,
-            ),
-        ),
-        db.execute(
-            select(func.count(Response.id))
-            .join(Record, and_(Record.id == Response.record_id, Record.dataset_id == dataset_id))
-            .filter(
-                Response.user_id == user_id,
-                Response.status == ResponseStatus.draft,
-            ),
-        ),
-        db.execute(
-            select(func.count(Record.id))
-            .outerjoin(Response, and_(Response.record_id == Record.id, Response.user_id == user_id))
-            .filter(
-                Record.dataset_id == dataset_id,
-                Record.status == RecordStatus.pending,
-                Response.id == None,  # noqa
-            ),
-        ),
-    )
+async def get_user_dataset_metrics(
+    db: AsyncSession,
+    search_engine: SearchEngine,
+    user: User,
+    dataset: Dataset,
+) -> dict:
+    total_records = await Record.count_by(db, dataset_id=dataset.id)
+    result = await search_engine.get_dataset_user_progress(dataset, user)
 
-    responses_submitted = responses_submitted.scalar_one()
-    responses_discarded = responses_discarded.scalar_one()
-    responses_draft = responses_draft.scalar_one()
-    responses_pending = responses_pending.scalar_one()
-    responses_total = responses_submitted + responses_discarded + responses_draft + responses_pending
+    submitted_responses = result.get("submitted", 0)
+    discarded_responses = result.get("discarded", 0)
+    draft_responses = result.get("draft", 0)
+    pending_responses = total_records - submitted_responses - discarded_responses - draft_responses
 
     return {
-        "responses": {
-            "total": responses_total,
-            "submitted": responses_submitted,
-            "discarded": responses_discarded,
-            "draft": responses_draft,
-            "pending": responses_pending,
-        },
+        "total": total_records,
+        "submitted": submitted_responses,
+        "discarded": discarded_responses,
+        "draft": draft_responses,
+        "pending": pending_responses,
     }
 
 
-async def get_dataset_progress(db: AsyncSession, dataset_id: UUID) -> dict:
-    records_completed, records_pending = await asyncio.gather(
-        db.execute(
-            select(func.count(Record.id)).where(
-                and_(
-                    Record.dataset_id == dataset_id,
-                    Record.status == RecordStatus.completed,
-                )
-            ),
-        ),
-        db.execute(
-            select(func.count(Record.id)).where(
-                and_(
-                    Record.dataset_id == dataset_id,
-                    Record.status == RecordStatus.pending,
-                )
-            ),
-        ),
-    )
-
-    records_completed = records_completed.scalar_one()
-    records_pending = records_pending.scalar_one()
-    records_total = records_completed + records_pending
-
+async def get_dataset_progress(
+    search_engine: SearchEngine,
+    dataset: Dataset,
+) -> dict:
+    result = await search_engine.get_dataset_progress(dataset)
     return {
-        "total": records_total,
-        "completed": records_completed,
-        "pending": records_pending,
+        "total": result.get("total", 0),
+        "completed": result.get("completed", 0),
+        "pending": result.get("pending", 0),
     }
 
 

argilla-server/src/argilla_server/search_engine/base.py

@@ -22,6 +22,7 @@
     Optional,
     Union,
     TypeVar,
+    Literal,
 )
 from uuid import UUID
 
@@ -45,7 +46,7 @@
     "SearchResponses",
     "SortBy",
     "MetadataMetrics",
-    "TermsMetadataMetrics",
+    "TermsMetrics",
     "IntegerMetadataMetrics",
     "FloatMetadataMetrics",
     "SuggestionFilterScope",
@@ -149,12 +150,12 @@ class Config:
         arbitrary_types_allowed = True
 
 
-class TermsMetadataMetrics(BaseModel):
+class TermsMetrics(BaseModel):
     class TermCount(BaseModel):
         term: str
         count: int
 
-    type: MetadataPropertyType = Field(MetadataPropertyType.terms)
+    type: Literal["terms"] = "terms"
     total: int
     values: List[TermCount] = Field(default_factory=list)
 
@@ -175,7 +176,7 @@ class FloatMetadataMetrics(NumericMetadataMetrics[float]):
     type: MetadataPropertyType = Field(MetadataPropertyType.float)
 
 
-MetadataMetrics = Union[TermsMetadataMetrics, IntegerMetadataMetrics, FloatMetadataMetrics]
+MetadataMetrics = Union[TermsMetrics, IntegerMetadataMetrics, FloatMetadataMetrics]
 
 
 class SearchEngine(metaclass=ABCMeta):
@@ -267,6 +268,14 @@ async def update_record_suggestion(self, suggestion: Suggestion):
     async def delete_record_suggestion(self, suggestion: Suggestion):
         pass
 
+    @abstractmethod
+    async def get_dataset_progress(self, dataset: Dataset) -> dict:
+        pass
+
+    @abstractmethod
+    async def get_dataset_user_progress(self, dataset: Dataset, user: User) -> dict:
+        pass
+
     @abstractmethod
     async def search(
         self,

argilla-server/src/argilla_server/search_engine/commons.py

@@ -32,6 +32,7 @@
     Suggestion,
     Vector,
     VectorSettings,
+    User,
 )
 from argilla_server.search_engine.base import (
     AndFilter,
@@ -50,7 +51,7 @@
     SearchResponses,
     SuggestionFilterScope,
     TermsFilter,
-    TermsMetadataMetrics,
+    TermsMetrics,
     TextQuery,
 )
 
@@ -390,6 +391,26 @@ async def delete_record_suggestion(self, suggestion: Suggestion):
             body={"script": f'ctx._source["suggestions"].remove("{suggestion.question.name}")'},
         )
 
+    async def get_dataset_progress(self, dataset: Dataset) -> dict:
+        index_name = es_index_name_for_dataset(dataset)
+
+        metrics = await self._compute_terms_metrics_for(index_name, "status")
+
+        return {"total": metrics.total, **{metric.term: metric.count for metric in metrics.values}}
+
+    async def get_dataset_user_progress(self, dataset: Dataset, user: User) -> dict:
+        index_name = es_index_name_for_dataset(dataset)
+
+        result = await self._compute_terms_metrics_for(
+            index_name,
+            field_name=es_field_for_response_property("status"),
+            query=es_nested_query(
+                path="responses",
+                query=es_term_query(es_field_for_response_property("user_id"), str(user.id)),
+            ),
+        )
+        return {"total": result.total, **{metric.term: metric.count for metric in result.values}}
+
     async def search(
         self,
         dataset: Dataset,
@@ -598,21 +619,27 @@ async def _metrics_for_numeric_property(
 
         return metrics_class(min=stats["min"], max=stats["max"])
 
-    async def _metrics_for_terms_property(
-        self, index_name: str, metadata_property: MetadataProperty, query: Optional[dict] = None
-    ) -> TermsMetadataMetrics:
-        field_name = es_field_for_metadata_property(metadata_property)
+    async def _compute_terms_metrics_for(
+        self, index_name: str, field_name: str, query: Optional[dict] = None
+    ) -> TermsMetrics:
         query = query or {"match_all": {}}
 
         total_terms = await self.__value_count_aggregation(index_name, field_name=field_name, query=query)
         if total_terms == 0:
-            return TermsMetadataMetrics(total=total_terms)
+            return TermsMetrics(total=total_terms)
 
-        terms_buckets = await self.__terms_aggregation(index_name, field_name=field_name, query=query, size=total_terms)
+        terms_buckets = await self._terms_aggregation(index_name, field_name=field_name, query=query, size=total_terms)
         terms_values = [
-            TermsMetadataMetrics.TermCount(term=bucket["key"], count=bucket["doc_count"]) for bucket in terms_buckets
+            TermsMetrics.TermCount(term=bucket["key"], count=bucket["doc_count"]) for bucket in terms_buckets
         ]
-        return TermsMetadataMetrics(total=total_terms, values=terms_values)
+
+        return TermsMetrics(total=total_terms, values=terms_values)
+
+    async def _metrics_for_terms_property(
+        self, index_name: str, metadata_property: MetadataProperty, query: Optional[dict] = None
+    ) -> TermsMetrics:
+        field_name = es_field_for_metadata_property(metadata_property)
+        return await self._compute_terms_metrics_for(index_name, field_name, query)
 
     def _configure_index_mappings(self, dataset: Dataset) -> dict:
         return {
@@ -845,7 +872,7 @@ def _map_record_fields_to_es(cls, fields: dict, dataset_fields: List[Field]) ->
 
         return fields
 
-    async def __terms_aggregation(self, index_name: str, field_name: str, query: dict, size: int) -> List[dict]:
+    async def _terms_aggregation(self, index_name: str, field_name: str, query: dict, size: int) -> List[dict]:
         aggregation_name = "terms_agg"
 
         terms_agg = {aggregation_name: {"terms": {"field": field_name, "size": min(size, self.max_terms_size)}}}