@@ -459,12 +459,8 @@ def get_combinator_queries(self):
459
459
)
460
460
)
461
461
compiler_ .pre_sql_setup ()
462
- # Standardize columns as main query required.
463
462
columns = compiler_ .get_columns ()
464
- if self .query .values_select :
465
- _ , exprs = zip (* columns , strict = True )
466
- columns = tuple (zip (main_query_fields , exprs , strict = True ))
467
- parts .append ((compiler_ .build_query (columns ), compiler_ .collection_name ))
463
+ parts .append ((compiler_ .build_query (columns ), compiler_ , columns ))
468
464
except EmptyResultSet :
469
465
# Omit the empty queryset with UNION.
470
466
if self .query .combinator == "union" :
@@ -473,24 +469,39 @@ def get_combinator_queries(self):
473
469
# Raise EmptyResultSet if all the combinator queries are empty.
474
470
if not parts :
475
471
raise EmptyResultSet
476
- combinator_pipeline = parts .pop (0 )[0 ].get_pipeline ()
477
- for part , collection in parts :
478
- combinator_pipeline .append (
479
- {"$unionWith" : {"coll" : collection , "pipeline" : part .get_pipeline ()}}
480
- )
481
- if not self .query .combinator_all :
482
- ids = {}
483
- for alias , expr in main_query_columns :
484
- # Unfold foreign fields.
485
- if isinstance (expr , Col ) and expr .alias != self .collection_name :
486
- ids [self ._unfold_column (expr )] = expr .as_mql (self , self .connection )
487
- else :
488
- ids [alias ] = f"${ alias } "
489
- combinator_pipeline .append ({"$group" : {"_id" : ids }})
490
- projected_fields = self ._fold_columns (ids )
491
- combinator_pipeline .append ({"$addFields" : projected_fields })
492
- if "_id" not in projected_fields :
493
- combinator_pipeline .append ({"$unset" : "_id" })
472
+ combinator_pipeline = []
473
+ for part , compiler_ , columns in parts :
474
+ inner_pipeline = part .get_pipeline ()
475
+ # Standardize result fields.
476
+ fields = {}
477
+ # When a .count() is called, the main_query_field has length 1
478
+ # otherwise it has the same length as columns.
479
+ for alias , (ref , expr ) in zip (main_query_fields , columns , strict = False ):
480
+ if isinstance (expr , Col ) and expr .alias != compiler_ .collection_name :
481
+ fields [expr .alias ] = 1
482
+ else :
483
+ fields [alias ] = f"${ ref } " if alias != ref else 1
484
+ inner_pipeline .append ({"$project" : fields })
485
+ # Combine query with the current combinator pipeline.
486
+ if combinator_pipeline :
487
+ combinator_pipeline .append (
488
+ {"$unionWith" : {"coll" : compiler_ .collection_name , "pipeline" : inner_pipeline }}
489
+ )
490
+ else :
491
+ combinator_pipeline = inner_pipeline
492
+ if not self .query .combinator_all :
493
+ ids = {}
494
+ for alias , expr in main_query_columns :
495
+ # Unfold foreign fields.
496
+ if isinstance (expr , Col ) and expr .alias != self .collection_name :
497
+ ids [self ._unfold_column (expr )] = expr .as_mql (self , self .connection )
498
+ else :
499
+ ids [alias ] = f"${ alias } "
500
+ combinator_pipeline .append ({"$group" : {"_id" : ids }})
501
+ projected_fields = self ._fold_columns (ids )
502
+ combinator_pipeline .append ({"$addFields" : projected_fields })
503
+ if "_id" not in projected_fields :
504
+ combinator_pipeline .append ({"$unset" : "_id" })
494
505
return combinator_pipeline
495
506
496
507
def get_lookup_pipeline (self ):
0 commit comments