INTPYTHON-804 Prevent QuerySet.union() queries from duplicating $project

Author: WaVEV

django_mongodb_backend/compiler.py

@@ -625,7 +625,10 @@ def get_combinator_queries(self):
                     fields[expr.alias] = 1
                 else:
                     fields[alias] = f"${ref}" if alias != ref else 1
-            inner_pipeline.append({"$project": fields})
+            # Avoid duplicating the same $project stage when reusing subquery
+            # projections.
+            if not inner_pipeline or inner_pipeline[-1] != {"$project": fields}:
+                inner_pipeline.append({"$project": fields})
             # Combine query with the current combinator pipeline.
             if combinator_pipeline:
                 combinator_pipeline.append(

docs/releases/5.2.x.rst

@@ -15,7 +15,8 @@ New features
 Bug fixes
 ---------
 
-- ...
+- Prevented ``QuerySet.union()`` queries from duplicating the ``$project``
+  pipeline.
 
 5.2.2
 =====

tests/queries_/models.py

@@ -13,6 +13,7 @@ def __str__(self):
 class Book(models.Model):
     title = models.CharField(max_length=10)
     author = models.ForeignKey(Author, models.CASCADE)
+    isbn = models.CharField(max_length=13)
 
     def __str__(self):
         return self.title

tests/queries_/test_qs_combinators.py

@@ -0,0 +1,71 @@
+from django.test import TestCase
+
+from django_mongodb_backend.test import MongoTestCaseMixin
+
+from .models import Book
+
+
+class UnionTests(MongoTestCaseMixin, TestCase):
+    def test_union_simple_conditions(self):
+        with self.assertNumQueries(1) as ctx:
+            list(Book.objects.filter(title="star wars").union(Book.objects.filter(isbn__in="1234")))
+        self.assertAggregateQuery(
+            ctx.captured_queries[0]["sql"],
+            "queries__book",
+            [
+                {"$match": {"title": "star wars"}},
+                {"$project": {"_id": 1, "author_id": 1, "title": 1, "isbn": 1}},
+                {
+                    "$unionWith": {
+                        "coll": "queries__book",
+                        "pipeline": [
+                            {"$match": {"isbn": {"$in": ("1", "2", "3", "4")}}},
+                            {"$project": {"_id": 1, "author_id": 1, "title": 1, "isbn": 1}},
+                        ],
+                    }
+                },
+                {
+                    "$group": {
+                        "_id": {
+                            "_id": "$_id",
+                            "author_id": "$author_id",
+                            "title": "$title",
+                            "isbn": "$isbn",
+                        }
+                    }
+                },
+                {
+                    "$addFields": {
+                        "_id": "$_id._id",
+                        "author_id": "$_id.author_id",
+                        "title": "$_id.title",
+                        "isbn": "$_id.isbn",
+                    }
+                },
+            ],
+        )
+
+    def test_union_all_simple_conditions(self):
+        with self.assertNumQueries(1) as ctx:
+            list(
+                Book.objects.filter(title="star wars").union(
+                    Book.objects.filter(isbn="1234"), all=True
+                )
+            )
+        self.assertAggregateQuery(
+            ctx.captured_queries[0]["sql"],
+            "queries__book",
+            [
+                {"$match": {"title": "star wars"}},
+                {"$project": {"_id": 1, "author_id": 1, "title": 1, "isbn": 1}},
+                {
+                    "$unionWith": {
+                        "coll": "queries__book",
+                        "pipeline": [
+                            {"$match": {"isbn": "1234"}},
+                            {"$project": {"_id": 1, "author_id": 1, "title": 1, "isbn": 1}},
+                        ],
+                    }
+                },
+            ],
+        )