Fix columns order in aggregation queries

Author: WaVEV

django_mongodb_backend/compiler.py

@@ -403,12 +403,6 @@ def columns(self):
         columns = (
             self.get_default_columns(select_mask) if self.query.default_cols else self.query.select
         )
-        # Populate QuerySet.select_related() data.
-        related_columns = []
-        if self.query.select_related:
-            self.get_related_selections(related_columns, select_mask)
-            if related_columns:
-                related_columns, _ = zip(*related_columns, strict=True)
 
         annotation_idx = 1
 
@@ -427,11 +421,28 @@ def project_field(column):
                 annotation_idx += 1
             return target, column
 
-        return (
-            tuple(map(project_field, columns))
-            + tuple(self.annotations.items())
-            + tuple(map(project_field, related_columns))
-        )
+        selected = []
+        if self.query.selected is None:
+            selected = [
+                *(project_field(col) for col in columns),
+                *self.annotations.items(),
+            ]
+        else:
+            for expression in self.query.selected.values():
+                # Reference to an annotation.
+                if isinstance(expression, str):
+                    alias, expression = expression, self.annotations[expression]
+                # Reference to a column.
+                elif isinstance(expression, int):
+                    alias, expression = project_field(columns[expression])
+                selected.append((alias, expression))
+        # Populate QuerySet.select_related() data.
+        related_columns = []
+        if self.query.select_related:
+            self.get_related_selections(related_columns, select_mask)
+            if related_columns:
+                related_columns, _ = zip(*related_columns, strict=True)
+        return tuple(selected) + tuple(map(project_field, related_columns))
 
     @cached_property
     def base_table(self):
@@ -478,7 +489,11 @@ def get_combinator_queries(self):
                 # If the columns list is limited, then all combined queries
                 # must have the same columns list. Set the selects defined on
                 # the query on all combined queries, if not already set.
-                if not compiler_.query.values_select and self.query.values_select:
+                selected = self.query.selected
+                if selected is not None and compiler_.query.selected is None:
+                    compiler_.query = compiler_.query.clone()
+                    compiler_.query.set_values(selected)
+                elif not compiler_.query.values_select and self.query.values_select:
                     compiler_.query = compiler_.query.clone()
                     compiler_.query.set_values(
                         (

django_mongodb_backend/expressions.py

@@ -150,7 +150,11 @@ def ref(self, compiler, connection):  # noqa: ARG001
         if isinstance(self.source, Col) and self.source.alias != compiler.collection_name
         else ""
     )
-    return f"${prefix}{self.refs}"
+    if hasattr(self, "ordinal"):
+        refs, _ = compiler.columns[self.ordinal - 1]
+    else:
+        refs = self.refs
+    return f"${prefix}{refs}"
 
 
 def star(self, compiler, connection):  # noqa: ARG001

tests/indexes_/test_condition.py

@@ -78,7 +78,7 @@ def test_composite_index(self):
                 {
                     "$and": [
                         {"number": {"$gte": 3}},
-                        {"$or": [{"body": {"$gt": "test1"}}, {"body": {"$in": ["A", "B"]}}]},
+                        {"$or": [{"body": {"$gt": "test1"}}, {"body": {"$in": ("A", "B")}}]},
                     ]
                 },
             )