Support ordering with null first / null last.

Author: WaVEV

django_mongodb/compiler.py

@@ -5,9 +5,10 @@
 from django.db import DatabaseError, IntegrityError, NotSupportedError
 from django.db.models import Count, Expression
 from django.db.models.aggregates import Aggregate, Variance
-from django.db.models.expressions import Col, Ref, Value
+from django.db.models.expressions import Case, Col, Ref, Value, When
 from django.db.models.functions.comparison import Coalesce
 from django.db.models.functions.math import Power
+from django.db.models.lookups import IsNull
 from django.db.models.sql import compiler
 from django.db.models.sql.constants import GET_ITERATOR_CHUNK_SIZE, MULTI, SINGLE
 from django.utils.functional import cached_property
@@ -333,13 +334,14 @@ def build_query(self, columns=None):
         query = self.query_class(self)
         query.aggregation_pipeline = self.get_aggregation_pipeline()
         query.lookup_pipeline = self.get_lookup_pipeline()
-        ordering_fields, order, need_extra_fields = self.preprocess_orderby()
+        orderby_annotations, ordering_fields, order = self.preprocess_orderby()
         query.project_fields = self.get_project_fields(columns, ordering_fields)
         query.ordering = order
-        if need_extra_fields and columns is None:
-            query.extra_fields = self.get_project_fields(
-                ((name, field) for name, field in ordering_fields if name.startswith("__order"))
-            )
+
+        # Post pipeline fields, some of them need some refs to be compted, so we add this fields
+        # after the main part of the pipeline has finished.
+        if orderby_annotations:
+            query.extra_fields = self.get_project_fields(orderby_annotations, add_fields=True)
         try:
             where = getattr(self, "where", self.query.where)
             query.mongo_query = (
@@ -412,7 +414,7 @@ def _get_aggregate_expressions(self, expr):
     def get_aggregation_pipeline(self):
         return self._group_pipeline
 
-    def get_project_fields(self, columns=None, ordering=None):
+    def get_project_fields(self, columns=None, ordering=None, add_fields=False):
         fields = {}
         for name, expr in columns or []:
             try:
@@ -430,7 +432,7 @@ def get_project_fields(self, columns=None, ordering=None):
                 # another column.
                 fields[name] = 1 if name == column else f"${column}"
 
-        if fields:
+        if fields and not add_fields:
             # Add related fields.
             for alias in self.query.alias_map:
                 if self.query.alias_refcount[alias] and self.collection_name != alias:
@@ -446,20 +448,25 @@ def get_project_fields(self, columns=None, ordering=None):
 
     def preprocess_orderby(self):
         fields = {}
+        orderby_annotations = {}
         result = SON()
-        need_extra_fields = False
         idx = itertools.count(start=1)
         for order in self._order_by or []:
-            if isinstance(order.expression, Ref):
-                fieldname = order.expression.refs
-            elif isinstance(order.expression, Col):
+            if isinstance(order.expression, Col | Ref):
                 fieldname = order.expression.as_mql(self, self.connection).removeprefix("$")
+                fields[fieldname] = order.expression
             else:
                 fieldname = f"__order{next(idx)}"
-                need_extra_fields = True
-            fields[fieldname] = order.expression
+                orderby_annotations[fieldname] = order.expression
+
+            if order.nulls_first or order.nulls_last:
+                null_fieldname = f"__order{next(idx)}"
+                condition = When(IsNull(order.expression, True), then=Value(1))
+                orderby_annotations[null_fieldname] = Case(condition, default=Value(0))
+                result[null_fieldname] = DESCENDING if order.nulls_first else ASCENDING
+
             result[fieldname] = DESCENDING if order.descending else ASCENDING
-        return tuple(fields.items()), result, need_extra_fields
+        return tuple(orderby_annotations.items()), tuple(fields.items()), result
 
 
 class SQLInsertCompiler(SQLCompiler):

django_mongodb/features.py

@@ -377,6 +377,7 @@ def django_test_expected_failures(self):
         "QuerySet.distinct() is not supported.": {
             "aggregation.tests.AggregateTestCase.test_sum_distinct_aggregate",
             "lookup.tests.LookupTests.test_lookup_collision_distinct",
+            "ordering.tests.OrderingTests.test_orders_nulls_first_on_filtered_subquery"
             "queries.tests.ExcludeTest17600.test_exclude_plain_distinct",
             "queries.tests.ExcludeTest17600.test_exclude_with_q_is_equal_to_plain_exclude",
             "queries.tests.ExcludeTest17600.test_exclude_with_q_object_distinct",
@@ -420,11 +421,6 @@ def django_test_expected_failures(self):
             "queries.tests.ValuesQuerysetTests.test_named_values_list_without_fields",
             "select_related.tests.SelectRelatedTests.test_select_related_with_extra",
         },
-        "Ordering a QuerySet by null_first/nulls_last is not supported on MongoDB.": {
-            "ordering.tests.OrderingTests.test_order_by_nulls_first",
-            "ordering.tests.OrderingTests.test_order_by_nulls_last",
-            "ordering.tests.OrderingTests.test_orders_nulls_first_on_filtered_subquery",
-        },
         "QuerySet.update() crash: Unrecognized expression '$count'": {
             "update.tests.AdvancedTests.test_update_annotated_multi_table_queryset",
         },

django_mongodb/functions.py

@@ -205,7 +205,7 @@ def wrapped(self, compiler, connection):
         lhs_mql = process_lhs(self, compiler, connection)
         return {
             "$cond": {
-                "if": {"$eq": [lhs_mql, None]},
+                "if": connection.mongo_operators["isnull"](lhs_mql, True),
                 "then": None,
                 "else": {f"${operator}": lhs_mql},
             }

django_mongodb/operations.py

@@ -6,9 +6,9 @@
 
 from bson.decimal128 import Decimal128
 from django.conf import settings
-from django.db import DataError, NotSupportedError
+from django.db import DataError
 from django.db.backends.base.operations import BaseDatabaseOperations
-from django.db.models.expressions import Combinable, OrderBy
+from django.db.models.expressions import Combinable
 from django.utils import timezone
 from django.utils.regex_helper import _lazy_re_compile
 
@@ -158,11 +158,6 @@ def convert_uuidfield_value(self, value, expression, connection):
             value = uuid.UUID(value)
         return value
 
-    def check_expression_support(self, expression):
-        if isinstance(expression, OrderBy) and (expression.nulls_first or expression.nulls_last):
-            option = "null_first" if expression.nulls_first else "nulls_last"
-            raise NotSupportedError(f"Ordering a QuerySet by {option} is not supported on MongoDB.")
-
     def combine_expression(self, connector, sub_expressions):
         lhs, rhs = sub_expressions
         if connector == Combinable.BITLEFTSHIFT:

django_mongodb/query.py

@@ -86,7 +86,7 @@ def get_pipeline(self):
         if self.extra_fields:
             pipeline.append({"$addFields": self.extra_fields})
         if self.ordering:
-            pipeline.append({"$sort": dict(self.ordering)})
+            pipeline.append({"$sort": self.ordering})
         if self.query.low_mark > 0:
             pipeline.append({"$skip": self.query.low_mark})
         if self.query.high_mark is not None: