add support for ordering by nulls first & last

Author: WaVEV

README.md

@@ -122,8 +122,6 @@ Migrations for 'admin':
 - `Subquery`, `Exists`, and using a `QuerySet` in `QuerySet.annotate()` aren't
   supported.
 
-* Ordering a `QuerySet` by `nulls_first` or `nulls_last` isn't supported.
-
 - `DateTimeField` doesn't support microsecond precision, and correspondingly,
   `DurationField` stores milliseconds rather than microseconds.
 

django_mongodb/compiler.py

@@ -6,9 +6,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
@@ -351,8 +352,13 @@ def build_query(self, columns=None):
         ordering_fields, sort_ordering, extra_fields = self._get_ordering()
         query.project_fields = self.get_project_fields(columns, ordering_fields)
         query.ordering = sort_ordering
-        if extra_fields and columns is None:
-            query.extra_fields = self.get_project_fields(extra_fields)
+        # extra_fields may contain references to other columns or expressions.
+        if columns is None:
+            extra_fields += ordering_fields
+        if extra_fields:
+            query.extra_fields = {
+                field_name: expr.as_mql(self, self.connection) for field_name, expr in extra_fields
+            }
         where = self.get_where()
         try:
             expr = where.as_mql(self, self.connection) if where else {}
@@ -465,13 +471,24 @@ def _get_ordering(self):
         idx = itertools.count(start=1)
         for order in self.order_by_expressions or []:
             if isinstance(order.expression, Ref):
+                # TODO change?
+                # field_name = order.expression.as_mql(self, self.connection).removeprefix("$")
+                # extra_fields[field_name] = order.expression
                 field_name = order.expression.refs
             elif isinstance(order.expression, Col):
                 field_name = order.expression.as_mql(self, self.connection).removeprefix("$")
+                fields[field_name] = order.expression
             else:
                 # The expression must be added to extra_fields with an alias.
                 field_name = f"__order{next(idx)}"
                 extra_fields[field_name] = order.expression
+            # If the expression is ordered by NULLS FIRST or NULLS LAST,
+            # add a field for sorting that's 1 if null or 0 if not.
+            if order.nulls_first or order.nulls_last:
+                null_fieldname = f"__order{next(idx)}"
+                condition = When(IsNull(order.expression, True), then=Value(1))
+                extra_fields[null_fieldname] = Case(condition, default=Value(0))
+                sort_ordering[null_fieldname] = DESCENDING if order.nulls_first else ASCENDING
             fields[field_name] = order.expression
             sort_ordering[field_name] = DESCENDING if order.descending else ASCENDING
         return tuple(fields.items()), sort_ordering, tuple(extra_fields.items())

django_mongodb/features.py

@@ -376,6 +376,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",
@@ -419,11 +420,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

@@ -138,7 +138,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
 
@@ -140,11 +140,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: