Handle group by.

Author: WaVEV

django_mongodb/compiler.py

@@ -3,6 +3,7 @@
 from django.db.models import Count, Expression
 from django.db.models.aggregates import Aggregate
 from django.db.models.constants import LOOKUP_SEP
+from django.db.models.expressions import Value
 from django.db.models.sql import compiler
 from django.db.models.sql.constants import GET_ITERATOR_CHUNK_SIZE, MULTI
 from django.utils.functional import cached_property
@@ -139,7 +140,7 @@ def build_query(self, columns=None):
         query = self.query_class(self)
         query.project_fields = self.get_project_fields(columns)
         query.lookup_pipeline = self.get_lookup_pipeline()
-        query.annotation_stage = self.get_group_pipeline()
+        query.aggregation_stage = self.get_aggregation_pipeline()
         try:
             query.mongo_query = {"$expr": self.query.where.as_mql(self, self.connection)}
         except FullResultSet:
@@ -238,7 +239,7 @@ def get_lookup_pipeline(self):
             result += self.query.alias_map[alias].as_mql(self, self.connection)
         return result
 
-    def get_group_pipeline(self):
+    def get_aggregation_pipeline(self):
         pipeline = None
         if any(isinstance(a, Aggregate) for a in self.query.annotations.values()):
             result = {}

django_mongodb/functions.py

@@ -1,6 +1,6 @@
 from django.db import NotSupportedError
 from django.db.models.aggregates import Aggregate, Avg, Count, Max, Min, StdDev, Sum, Variance
-from django.db.models.expressions import Func
+from django.db.models.expressions import Case, Func, Value, When
 from django.db.models.functions.comparison import Cast, Coalesce, Greatest, Least, NullIf
 from django.db.models.functions.datetime import (
     Extract,
@@ -41,7 +41,7 @@
     Min: "min",
     StdDev: "stddev",
     Sum: "sum",
-    Variance: "variance",
+    Variance: "stdDevPop",
 }
 
 
@@ -69,8 +69,19 @@
 }
 
 
-def aggregate(self, compiler, connection):  # noqa: ARG001
-    pass
+def aggregate(self, compiler, connection, **extra_context):  # noqa: ARG001
+    if self.filter:
+        copy = self.copy()
+        copy.filter = None
+        source_expressions = copy.get_source_expressions()
+        condition = When(self.filter, then=source_expressions[0])
+        copy.set_source_expressions([Case(condition)] + source_expressions[1:])
+        node = copy
+    else:
+        node = self
+    lhs_mql = process_lhs(node, compiler, connection)
+    operator = MONGO_AGGREGATION.get(self.__class__)
+    return {f"${operator}": lhs_mql}
 
 
 def cast(self, compiler, connection):
@@ -103,6 +114,19 @@ def cot(self, compiler, connection):
     return {"$divide": [1, {"$tan": lhs_mql}]}
 
 
+def count(self, compiler, connection, **extra_context):  # noqa: ARG001
+    if self.filter:
+        copy = self.copy()
+        copy.filter = None
+        source_expressions = copy.get_source_expressions()
+        condition = When(self.filter, then=Value(1))
+        copy.set_source_expressions([Case(condition)] + source_expressions[1:])
+        lhs_mql = process_lhs(copy, compiler, connection)
+    else:
+        lhs_mql = Value(1).as_mql(compiler, connection)
+    return {"$sum": lhs_mql}
+
+
 def extract(self, compiler, connection):
     lhs_mql = process_lhs(self, compiler, connection)
     operator = EXTRACT_OPERATORS.get(self.lookup_name)
@@ -208,6 +232,7 @@ def register_functions():
     Concat.as_mql = concat
     ConcatPair.as_mql = concat_pair
     Cot.as_mql = cot
+    Count.as_mql = count
     Extract.as_mql = extract
     Func.as_mql = func
     Left.as_mql = left

django_mongodb/query.py

@@ -50,6 +50,7 @@ def __init__(self, compiler):
         self.lookup_pipeline = None
         self.annotation_stage = None
         self.project_fields = None
+        self.aggregation_stage = None
 
     def __repr__(self):
         return f"<MongoQuery: {self.mongo_query!r} ORDER {self.ordering!r}>"
@@ -105,8 +106,8 @@ def get_cursor(self):
             pipeline.extend(self.lookup_pipeline)
         if self.mongo_query:
             pipeline.append({"$match": self.mongo_query})
-        if self.annotation_stage:
-            pipeline.extend(self.annotation_stage)
+        if self.aggregation_stage:
+            pipeline.extend(self.aggregation_stage)
         if self.project_fields:
             pipeline.append({"$project": self.project_fields})
         if self.ordering: