Improvements.

Author: WaVEV

django_mongodb/compiler.py

@@ -7,7 +7,7 @@
 from django.db import IntegrityError, NotSupportedError
 from django.db.models import Count
 from django.db.models.aggregates import Aggregate, Variance
-from django.db.models.expressions import Case, Col, Ref, Value, When
+from django.db.models.expressions import Case, Col, OrderBy, 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
@@ -381,8 +381,10 @@ def get_columns(self):
         which should be loaded by the query.
         """
         select_mask = self.query.get_select_mask()
-        columns = (
-            self.get_default_columns(select_mask) if self.query.default_cols else self.query.select
+        columns = filter(
+            # The extra order by columns are handled by order_by_objs variables.
+            lambda col: not isinstance(col, OrderBy),
+            self.get_default_columns(select_mask) if self.query.default_cols else self.query.select,
         )
         # Populate QuerySet.select_related() data.
         related_columns = []
@@ -439,13 +441,10 @@ def get_combinator_queries(self):
                             *self.query.annotation_select,
                         )
                     )
-                compiler_.pre_sql_setup(with_col_aliases=False)
-                # Avoid $project (columns=None) if unneeded.
-                columns = (
-                    compiler_.get_columns()
-                    if compiler_.query.annotations or not compiler_.query.default_cols
-                    else None
-                )
+                compiler_.pre_sql_setup()
+                # Standardize columns as main query required.
+                _, exprs = zip(*compiler_.get_columns(), strict=True)
+                columns = tuple(zip(self.query.values_select, exprs, strict=True))
                 parts.append((compiler_.build_query(columns), compiler_.collection_name))
 
             except EmptyResultSet:
@@ -454,7 +453,9 @@ def get_combinator_queries(self):
                 if self.query.combinator == "union":
                     continue
                 raise
-
+        # Raise EmptyResultSet if all the combinator queries are empty.
+        if not parts:
+            raise EmptyResultSet
         combinator_pipeline = parts.pop(0)[0].get_pipeline() if parts else None
         if self.query.combinator == "union":
             for part, collection in parts:
@@ -463,13 +464,26 @@ def get_combinator_queries(self):
                 )
                 if not self.query.combinator_all:
                     ids = {}
-                    annotation_group_idx = itertools.count(start=1)
-                    for _, expr in self.get_columns():
-                        alias, replacement = self._get_group_alias_column(
-                            expr, annotation_group_idx
+                    for alias, expr in self.get_columns():
+                        ids[alias] = (
+                            expr.as_mql(self, self.connection)
+                            if isinstance(expr, Col | Ref)
+                            else f"${alias}"
                         )
-                        ids[alias] = expr.as_mql(self, self.connection)
                     combinator_pipeline.append({"$group": {"_id": ids}})
+                    projected_fields = defaultdict(dict)
+                    for key in ids:
+                        value = f"$_id.{key}"
+                        if self.GROUP_SEPARATOR in key:
+                            table, field = key.split(self.GROUP_SEPARATOR)
+                            projected_fields[table][field] = value
+                        else:
+                            projected_fields[key] = value
+                    # Convert defaultdict to dict so it doesn't appear as
+                    # "defaultdict(<CLASS 'dict'>, ..." in query logging.
+                    combinator_pipeline.append({"$addFields": dict(projected_fields)})
+                    if "_id" not in projected_fields:
+                        combinator_pipeline.append({"$unset": "_id"})
         else:
             raise NotSupportedError(f"Combinator {self.query.combinator} isn't supported.")
         return combinator_pipeline