Refactor folding and unfolding foreign columns in group stages

Author: WaVEV

django_mongodb/compiler.py

@@ -32,6 +32,33 @@ def __init__(self, *args, **kwargs):
         # A list of OrderBy objects for this query.
         self.order_by_objs = None
 
+    def _unfold_column(self, col):
+        """
+        Flattens a field by returning its target or by replacing dots with GROUP_SEPARATOR
+        for foreign fields.
+        """
+        if self.collection_name == col.alias:
+            return col.target.column
+        # If this is a foreign field, replace the normal dot (.) with
+        # GROUP_SEPARATOR since FieldPath field names may not contain '.'.
+        return f"{col.alias}{self.GROUP_SEPARATOR}{col.target.column}"
+
+    def _fold_columns(self, unfold_columns):
+        """
+        Converts flat columns into a nested dictionary, grouping fields by table names.
+        """
+        result = defaultdict(dict)
+        for key in unfold_columns:
+            value = f"$_id.{key}"
+            if self.GROUP_SEPARATOR in key:
+                table, field = key.split(self.GROUP_SEPARATOR)
+                result[table][field] = value
+            else:
+                result[key] = value
+        # Convert defaultdict to dict so it doesn't appear as
+        # "defaultdict(<CLASS 'dict'>, ..." in query logging.
+        return dict(result)
+
     def _get_group_alias_column(self, expr, annotation_group_idx):
         """Generate a dummy field for use in the ids fields in $group."""
         replacement = None
@@ -42,11 +69,7 @@ def _get_group_alias_column(self, expr, annotation_group_idx):
             alias = f"__annotation_group{next(annotation_group_idx)}"
             col = self._get_column_from_expression(expr, alias)
             replacement = col
-        if self.collection_name == col.alias:
-            return col.target.column, replacement
-        # If this is a foreign field, replace the normal dot (.) with
-        # GROUP_SEPARATOR since FieldPath field names may not contain '.'.
-        return f"{col.alias}{self.GROUP_SEPARATOR}{col.target.column}", replacement
+        return self._unfold_column(col), replacement
 
     def _get_column_from_expression(self, expr, alias):
         """
@@ -186,17 +209,8 @@ def _build_aggregation_pipeline(self, ids, group):
         else:
             group["_id"] = ids
             pipeline.append({"$group": group})
-            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.
-            pipeline.append({"$addFields": dict(projected_fields)})
+            projected_fields = self._fold_columns(ids)
+            pipeline.append({"$addFields": projected_fields})
             if "_id" not in projected_fields:
                 pipeline.append({"$unset": "_id"})
         return pipeline
@@ -453,8 +467,7 @@ def get_combinator_queries(self):
                 parts.append((compiler_.build_query(columns), compiler_.collection_name))
 
             except EmptyResultSet:
-                # Omit the empty queryset with UNION and with DIFFERENCE if the
-                # first queryset is nonempty.
+                # Omit the empty queryset with UNION.
                 if self.query.combinator == "union":
                     continue
                 raise
@@ -470,25 +483,14 @@ def get_combinator_queries(self):
                 if not self.query.combinator_all:
                     ids = {}
                     for alias, expr in main_query_columns:
-                        collection = expr.alias if isinstance(expr, Col) else None
-                        if collection and collection != self.collection_name:
-                            ids[
-                                f"{expr.alias}{self.GROUP_SEPARATOR}{expr.target.column}"
-                            ] = expr.as_mql(self, self.connection)
+                        # Unfold foreign fields.
+                        if isinstance(expr, Col) and expr.alias != self.collection_name:
+                            ids[self._unfold_column(expr)] = expr.as_mql(self, self.connection)
                         else:
                             ids[alias] = f"${alias}"
                     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)})
+                    projected_fields = self._fold_columns(ids)
+                    combinator_pipeline.append({"$addFields": projected_fields})
                     if "_id" not in projected_fields:
                         combinator_pipeline.append({"$unset": "_id"})
         else: