Support distinct keyword.

Author: WaVEV

django_mongodb/compiler.py

@@ -45,11 +45,22 @@ def pre_sql_setup(self, with_col_aliases=False):
                         aggregation_idx += 1
                     else:
                         alias = target
-                    group[alias] = sub_expr.as_mql(self, self.connection)
+
                     column_target = expr.output_field.__class__()
                     column_target.db_column = alias
                     column_target.set_attributes_from_name(alias)
-                    replacements[sub_expr] = Col(self.collection_name, column_target)
+                    inner_column = Col(self.collection_name, column_target)
+                    if sub_expr.distinct:
+                        inner_expr = sub_expr.as_mql(self, self.connection, force_filters=True)
+                        rhs = next(iter(inner_expr.values()))
+                        group[alias] = {"$addToSet": rhs}
+                        replacing_expr = sub_expr.copy()
+                        replacing_expr.set_source_expressions([inner_column])
+                    else:
+                        group[alias] = sub_expr.as_mql(self, self.connection)
+                        replacing_expr = inner_column
+
+                    replacements[sub_expr] = replacing_expr
                 result_expr = expr.replace_expressions(replacements)
                 all_replacements.update(replacements)
             group_expressions |= set(expr.get_group_by_cols())
@@ -93,7 +104,7 @@ def _ccc(col):
                 }
             )
             pipeline = []
-            if ids is None:
+            if not ids:
                 group["_id"] = None
                 pipeline.append({"$facet": {"group": [{"$group": group}]}})
                 pipeline.append(

django_mongodb/functions.py

@@ -1,3 +1,5 @@
+from copy import deepcopy
+
 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 Case, Func, Star, Value, When
@@ -32,6 +34,8 @@
     Trim,
     Upper,
 )
+from django.db.models.lookups import Exact
+from django.db.models.query_utils import Q
 
 from .query_utils import process_lhs
 
@@ -115,21 +119,34 @@ 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:])
-        node = copy
-        cond = process_lhs(node, compiler, connection)
-    else:
-        node = self
-        lhs_mql = process_lhs(self, compiler, connection)
-        null_cond = {"$in": [{"$type": lhs_mql}, ["missing", "null"]]}
-        cond = {"$cond": {"if": null_cond, "then": 0, "else": 1}}
-    return {"$sum": cond}
+def count(self, compiler, connection, **extra_context):
+    if not self.distinct or extra_context.get("force_filters"):
+        if self.filter:
+            copy = self.copy()
+            copy.filter = None
+            source_expressions = copy.get_source_expressions()
+            filter_ = deepcopy(self.filter)
+            filter_.add(~Q(Exact(source_expressions[0], Value(None))), filter_.default)
+            condition = When(filter_, then=Value(1))
+            copy.set_source_expressions([Case(condition)] + source_expressions[1:])
+            node = copy
+            cond = process_lhs(node, compiler, connection)
+        else:
+            node = self
+            lhs_mql = process_lhs(self, compiler, connection)
+            null_cond = {"$in": [{"$type": lhs_mql}, ["missing", "null"]]}
+            cond = {
+                "$cond": {"if": null_cond, "then": None, "else": lhs_mql if self.distinct else 1}
+            }
+        operator = "$sum"
+        return {operator: cond}
+
+    operator = "$size"
+    lhs_mql = process_lhs(self, compiler, connection)
+
+    exits_null = {"$cond": {"if": {"$in": [{"$literal": None}, lhs_mql]}, "then": -1, "else": 0}}
+
+    return {"$add": [{"$size": lhs_mql}, exits_null]}
 
 
 def extract(self, compiler, connection):