handle the empty group by.

Author: WaVEV

django_mongodb/compiler.py

@@ -3,9 +3,9 @@
 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.expressions import Col, Value
 from django.db.models.sql import compiler
-from django.db.models.sql.constants import GET_ITERATOR_CHUNK_SIZE, MULTI
+from django.db.models.sql.constants import GET_ITERATOR_CHUNK_SIZE, MULTI, SINGLE
 from django.utils.functional import cached_property
 
 from .base import Cursor
@@ -17,6 +17,62 @@ class SQLCompiler(compiler.SQLCompiler):
 
     query_class = MongoQuery
 
+    def pre_sql_setup(self):
+        super().pre_sql_setup()
+        self.annotations = {}
+        group = {}
+        group_expressions = set()
+        aggregation_idx = 1
+        for target, expr in self.query.annotation_select.items():
+            if not expr.contains_aggregate:
+                result_expr = expr
+            else:
+                replacements = {}
+                for sub_expr in self._get_aggregate_expressions(expr):
+                    alias = f"__aggregation{aggregation_idx}"
+                    group[alias] = sub_expr.as_mql(self, self.connection)
+                    aggregation_idx += 1
+                    column_target = expr.output_field.__class__()
+                    column_target.set_attributes_from_name(alias)
+                    replacements[sub_expr] = Col(self.collection_name, column_target)
+                result_expr = expr.replace_expressions(replacements)
+
+            self.annotations[target] = result_expr
+        if group:
+            """
+            order_by = self.get_order_by()
+            for expr, (_, _, is_ref) in order_by:
+                # Skip references to the SELECT clause, as all expressions in
+                # the SELECT clause are already part of the GROUP BY.
+                if not is_ref:
+                    group_expressions |= set(expr.get_group_by_cols())
+            having_group_by = self.having.get_group_by_cols() if self.having else ()
+            for expr in having_group_by:
+                group_expressions.add(expr)
+            """
+
+            ids = (
+                None
+                if not group_expressions
+                else {
+                    col.target.column: col.as_mql(self, self.connection)
+                    for col in group_expressions
+                }
+            )
+            group["_id"] = ids
+
+            pipeline = [{"$group": group}]
+            if ids:
+                pipeline.append(
+                    {"$addFields": {key: f"$_id.{value[1:]}" for key, value in ids.items()}}
+                )
+                if "_id" not in ids:
+                    pipeline.append({"$unSet": "$_id"})
+
+            self._group_pipeline = pipeline
+        else:
+            self._group_pipeline = None
+
     def execute_sql(
         self, result_type=MULTI, chunked_fetch=False, chunk_size=GET_ITERATOR_CHUNK_SIZE
     ):
@@ -34,11 +90,13 @@ def execute_sql(
         except EmptyResultSet:
             return iter([]) if result_type == MULTI else None
 
-        return (
-            (self._make_result(row, columns) for row in query.fetch())
-            if result_type == MULTI
-            else self._make_result(next(query.fetch()), columns)
-        )
+        if result_type == MULTI:
+            return (self._make_result(row, columns) for row in query.fetch())
+
+        try:
+            return self._make_result(next(query.fetch()), columns)
+        except StopIteration:
+            return None
 
     def results_iter(
         self,
@@ -65,7 +123,7 @@ def results_iter(
         return rows
 
     def has_results(self):
-        return bool(self.get_count(check_exists=True))
+        return bool(self.execute_sql(SINGLE))
 
     def _make_result(self, entity, columns):
         """
@@ -138,9 +196,9 @@ def build_query(self, columns=None):
         """Check if the query is supported and prepare a MongoQuery."""
         self.check_query()
         query = self.query_class(self)
-        query.project_fields = self.get_project_fields(columns)
-        query.lookup_pipeline = self.get_lookup_pipeline()
         query.aggregation_stage = self.get_aggregation_pipeline()
+        query.lookup_pipeline = self.get_lookup_pipeline()
+        query.project_fields = self.get_project_fields(columns)
         try:
             query.mongo_query = {"$expr": self.query.where.as_mql(self, self.connection)}
         except FullResultSet:
@@ -180,7 +238,7 @@ def project_field(column):
 
         return (
             tuple(map(project_field, columns))
-            + tuple(self.query.annotation_select.items())
+            + tuple(self.annotations.items())
             + tuple(map(project_field, related_columns))
         )
 
@@ -239,52 +297,34 @@ def get_lookup_pipeline(self):
             result += self.query.alias_map[alias].as_mql(self, self.connection)
         return result
 
-    def get_aggregation_pipeline(self):
-        pipeline = None
-        if any(isinstance(a, Aggregate) for a in self.query.annotations.values()):
-            result = {}
-            # self.get_group_by(self.select, [])
-            for alias, annotation in self.query.annotation_select.items():
-                value = annotation.as_mql(self, self.connection)
-                if isinstance(value, list):
-                    value = value[0]
-                result[alias] = value
-
-            expressions = set()
-            for expr, *_ in self.select:
-                expressions |= set(expr.get_group_by_cols())
-            order_by = self.get_order_by()
-            for expr, (_, _, is_ref) in order_by:
-                # Skip references to the SELECT clause, as all expressions in
-                # the SELECT clause are already part of the GROUP BY.
-                if not is_ref:
-                    expressions |= set(expr.get_group_by_cols())
-            having_group_by = self.having.get_group_by_cols() if self.having else ()
-            for expr in having_group_by:
-                expressions.add(expr)
-
-            ids = (
-                None
-                if not expressions
-                else {col.target.column: col.as_mql(self, self.connection) for col in expressions}
-            )
-            result["_id"] = ids
-
-            pipeline = [{"$group": result}]
-            if ids:
-                pipeline.append(
-                    {"$addFields": {key: f"$_id.{value[1:]}" for key, value in ids.items()}}
+    def _get_aggregate_expressions2(self, expr):
+        stack = [(None, expr)]
+        while stack:
+            parent, expr = stack.pop()
+            if isinstance(expr, Aggregate):
+                yield parent
+            elif hasattr(expr, "get_source_expressions"):
+                stack.extend(
+                    [((expr, idx), se) for idx, se in enumerate(expr.get_source_expressions())]
                 )
-                if "_id" not in ids:
-                    pipeline.append({"$unSet": "$_id"})
 
-        return pipeline
+    def _get_aggregate_expressions(self, expr):
+        stack = [expr]
+        while stack:
+            expr = stack.pop()
+            if isinstance(expr, Aggregate):
+                yield expr
+            elif hasattr(expr, "get_source_expressions"):
+                stack.extend(expr.get_source_expressions())
+
+    def get_aggregation_pipeline(self):
+        return self._group_pipeline
 
     def get_project_fields(self, columns=None):
         fields = {}
         for name, expr in columns or []:
             try:
-                column = name if isinstance(expr, Aggregate) else expr.target.column
+                column = expr.target.column
             except AttributeError:
                 # Generate the MQL for an annotation.
                 try:

django_mongodb/functions.py

@@ -39,7 +39,7 @@
     Count: "sum",
     Max: "max",
     Min: "min",
-    StdDev: "stddev",
+    StdDev: "stdDevPop",
     Sum: "sum",
     Variance: "stdDevPop",
 }