@@ -44,10 +44,8 @@ def _get_group_alias_column(self, expr, annotation_group_idx):
4444
4545 def _get_column_from_expression (self , expr , alias ):
4646 """
47- Create a new column with the specified output type and alias to hold the aggregate value.
48-
49- This function generates a column target from the given expression, setting the column's
50- output type and assigning the provided alias to the column.
47+ Create a column named `alias` from the given expression to hold the
48+ aggregate value.
5149 """
5250 column_target = expr .output_field .__class__ ()
5351 column_target .db_column = alias
@@ -58,17 +56,18 @@ def _prepare_expressions_for_pipeline(self, expression, target, annotation_group
5856 """
5957 Prepare expressions for the aggregation pipeline.
6058
61- This function handles the computation of aggregation functions used by various expressions.
62- It separates and creates intermediate columns, and replaces nodes
63- to simulate a group by operation.
64-
65- In MongoDB, the `$group` stage does not allow operations over the aggregator
66- (e.g., `COALESCE(AVG(field), 3)`). However, it does support operations inside
67- the aggregation (e.g., `AVG(number * 2)`).
68- This function manages first cases by splitting the computation into stages:
69- it computes the aggregation function first, then applies additional operations
70- in a subsequent stage by replacing the aggregate expressions
71- with new columns prefixed by `__aggregation`.
59+ Handle the computation of aggregation functions used by various
60+ expressions. Separate and create intermediate columns, and replace
61+ nodes to simulate a group by operation.
62+
63+ MongoDB's $group stage doesn't allow operations over the aggregator,
64+ e.g. COALESCE(AVG(field), 3). However, it supports operations inside
65+ the aggregation, e.g. AVG(number * 2).
66+
67+ Handle the first case by splitting the computation into stages: compute
68+ the aggregation first, then applies additional operations in a
69+ subsequent stage by replacing the aggregate expressions with new
70+ columns prefixed by `__aggregation`.
7271 """
7372 replacements = {}
7473 group = {}
@@ -81,6 +80,8 @@ def _prepare_expressions_for_pipeline(self, expression, target, annotation_group
8180 column_target .set_attributes_from_name (alias )
8281 inner_column = Col (self .collection_name , column_target )
8382 if sub_expr .distinct :
83+ # If the expression should return distinct values, use
84+ # $addToSet to deduplicate.
8485 rhs = sub_expr .as_mql (self , self .connection , resolve_inner_expression = True )
8586 group [alias ] = {"$addToSet" : rhs }
8687 replacing_expr = sub_expr .copy ()
@@ -98,7 +99,7 @@ def _prepare_expressions_for_pipeline(self, expression, target, annotation_group
9899 return replacements , group
99100
100101 def _prepare_annotations_for_aggregation_pipeline (self ):
101- """Prepare annotations for the MongoDB aggregation pipeline."""
102+ """Prepare annotations for the aggregation pipeline."""
102103 replacements = {}
103104 group = {}
104105 annotation_group_idx = itertools .count (start = 1 )
@@ -109,7 +110,6 @@ def _prepare_annotations_for_aggregation_pipeline(self):
109110 )
110111 replacements .update (new_replacements )
111112 group .update (expr_group )
112-
113113 having_replacements , having_group = self ._prepare_expressions_for_pipeline (
114114 self .having , None , annotation_group_idx
115115 )
@@ -124,18 +124,15 @@ def _get_group_id_expressions(self, order_by):
124124 for expr , (_ , _ , is_ref ) in order_by :
125125 if not is_ref :
126126 group_expressions |= set (expr .get_group_by_cols ())
127-
128127 for expr , * _ in self .select :
129128 group_expressions |= set (expr .get_group_by_cols ())
130-
131129 having_group_by = self .having .get_group_by_cols () if self .having else ()
132130 for expr in having_group_by :
133131 group_expressions .add (expr )
134132 if isinstance (self .query .group_by , tuple | list ):
135133 group_expressions |= set (self .query .group_by )
136134 elif self .query .group_by is None :
137135 group_expressions = set ()
138-
139136 if not group_expressions :
140137 ids = None
141138 else :
@@ -151,7 +148,6 @@ def _get_group_id_expressions(self, order_by):
151148 ids [alias ] = Value (True ).as_mql (self , self .connection )
152149 if replacement is not None :
153150 replacements [col ] = replacement
154-
155151 return ids , replacements
156152
157153 def _build_aggregation_pipeline (self , ids , group ):
@@ -180,15 +176,13 @@ def _build_aggregation_pipeline(self, ids, group):
180176 for key in ids :
181177 value = f"$_id.{ key }
182178 if self .GROUP_SEPARATOR in key :
183- subtable , field = key .split (self .GROUP_SEPARATOR )
184- projected_fields [subtable ][field ] = value
179+ table , field = key .split (self .GROUP_SEPARATOR )
180+ projected_fields [table ][field ] = value
185181 else :
186182 projected_fields [key ] = value
187-
188183 pipeline .append ({"$addFields" : projected_fields })
189184 if "_id" not in projected_fields :
190185 pipeline .append ({"$unset" : "_id" })
191-
192186 return pipeline
193187
194188 def pre_sql_setup (self , with_col_aliases = False ):
@@ -213,7 +207,6 @@ def pre_sql_setup(self, with_col_aliases=False):
213207 }
214208 )
215209 self .aggregation_pipeline = pipeline
216-
217210 self .annotations = {
218211 target : expr .replace_expressions (all_replacements )
219212 for target , expr in self .query .annotation_select .items ()
0 commit comments