Aggregation rework

Author: WaVEV

django_mongodb_backend/aggregates.py

@@ -1,13 +1,13 @@
+from django.core.exceptions import EmptyResultSet, FullResultSet
 from django.db import NotSupportedError
 from django.db.models.aggregates import (
     Aggregate,
-    AggregateFilter,
     Count,
     StdDev,
     StringAgg,
     Variance,
 )
-from django.db.models.expressions import Case, Col, Value, When
+from django.db.models.expressions import Case, Value, When
 from django.db.models.lookups import IsNull
 
 from .query_utils import process_lhs
@@ -20,26 +20,26 @@ def aggregate(self, compiler, connection, operator=None, resolve_inner_expressio
     # TODO: isinstance(self.filter, Col) works around failure of
     # aggregation.tests.AggregateTestCase.test_distinct_on_aggregate. Is this
     # correct?
-    if self.filter is not None and not isinstance(self.filter, Col):
+    if self.filter is not None:
         # Generate a CASE statement for this aggregate.
-        node = self.copy()
-        node.filter = None
-        source_expressions = node.get_source_expressions()
-        condition = When(self.filter, then=source_expressions[0])
-        node.set_source_expressions([Case(condition), *source_expressions[1:]])
+        try:
+            lhs_mql = self.filter.as_mql(compiler, connection, as_expr=True)
+        except NotSupportedError:
+            source_expressions = self.get_source_expressions()
+            condition = Case(When(self.filter.condition, then=source_expressions[0]))
+            lhs_mql = condition.as_mql(compiler, connection)
+        except FullResultSet:
+            lhs_mql = source_expressions[0].as_mql(compiler, connection, as_expr=True)
+        except EmptyResultSet:
+            lhs_mql = Value(None).as_mql(compiler, connection, as_expr=True)
     else:
-        node = self
-    lhs_mql = process_lhs(node, compiler, connection, as_expr=True)
+        lhs_mql = process_lhs(self, compiler, connection, as_expr=True)
     if resolve_inner_expression:
         return lhs_mql
     operator = operator or MONGO_AGGREGATIONS.get(self.__class__, self.function.lower())
     return {f"${operator}": lhs_mql}
 
 
-def aggregate_filter(self, compiler, connection):
-    return self.condition.as_mql(compiler, connection, as_expr=True)
-
-
 def count(self, compiler, connection, resolve_inner_expression=False):
     """
     When resolve_inner_expression=True, return the MQL that resolves as a
@@ -48,14 +48,19 @@ def count(self, compiler, connection, resolve_inner_expression=False):
     """
     if not self.distinct or resolve_inner_expression:
         if self.filter:
-            node = self.copy()
-            node.filter = None
-            source_expressions = node.get_source_expressions()
-            condition = When(
-                self.filter, then=Case(When(IsNull(source_expressions[0], False), then=Value(1)))
-            )
-            node.set_source_expressions([Case(condition), *source_expressions[1:]])
-            inner_expression = process_lhs(node, compiler, connection, as_expr=True)
+            try:
+                inner_expression = self.filter.as_mql(compiler, connection, as_expr=True)
+            except NotSupportedError:
+                source_expressions = self.get_source_expressions()
+                condition = When(
+                    self.filter.condition,
+                    then=Case(When(IsNull(source_expressions[0], False), then=Value(1))),
+                )
+                inner_expression = Case(condition).as_mql(compiler, connection, as_expr=True)
+            except FullResultSet:
+                inner_expression = {"$sum": 1}
+            except EmptyResultSet:
+                inner_expression = {"$sum": 0}
         else:
             lhs_mql = process_lhs(self, compiler, connection, as_expr=True)
             null_cond = {"$in": [{"$type": lhs_mql}, ["missing", "null"]]}
@@ -87,7 +92,7 @@ def string_agg(self, compiler, connection):  # noqa: ARG001
 
 def register_aggregates():
     Aggregate.as_mql_expr = aggregate
-    AggregateFilter.as_mql_expr = aggregate_filter
+    # AggregateFilter.as_mql_expr = aggregate_filter
     Count.as_mql_expr = count
     StdDev.as_mql_expr = stddev_variance
     StringAgg.as_mql_expr = string_agg