Fixed #35586 -- Added support for set-returning database functions.

Aggregation optimization didn't account for not referenced set-returning annotations on Postgres.

Co-authored-by: Simon Charette <charette.s@gmail.com>

6 files changed, 54 insertions, 0 deletions

diff --git a/django/db/models/expressions.py b/django/db/models/expressions.py
index dc09e43fda..eee0eafc83 100644
--- a/django/db/models/expressions.py
+++ b/django/db/models/expressions.py
@@ -182,6 +182,8 @@ class BaseExpression:
     allowed_default = False
     # Can the expression be used during a constraint validation?
     constraint_validation_compatible = True
+    # Does the expression possibly return more than one row?
+    set_returning = False
 
     def __init__(self, output_field=None):
         if output_field is not None:
diff --git a/django/db/models/sql/query.py b/django/db/models/sql/query.py
index c1e2fc1d4f..aef3f48f10 100644
--- a/django/db/models/sql/query.py
+++ b/django/db/models/sql/query.py
@@ -491,6 +491,11 @@ class Query(BaseExpression):
             )
             or having
         )
+        set_returning_annotations = {
+            alias
+            for alias, annotation in self.annotation_select.items()
+            if getattr(annotation, "set_returning", False)
+        }
         # Decide if we need to use a subquery.
         #
         # Existing aggregations would cause incorrect results as
@@ -510,6 +515,7 @@ class Query(BaseExpression):
             or qualify
             or self.distinct
             or self.combinator
+            or set_returning_annotations
         ):
             from django.db.models.sql.subqueries import AggregateQuery
 
@@ -551,6 +557,9 @@ class Query(BaseExpression):
                         if annotation.get_group_by_cols():
                             annotation_mask.add(annotation_alias)
                     inner_query.set_annotation_mask(annotation_mask)
+                    # Annotations that possibly return multiple rows cannot
+                    # be masked as they might have an incidence on the query.
+                    annotation_mask |= set_returning_annotations
 
             # Add aggregates to the outer AggregateQuery. This requires making
             # sure all columns referenced by the aggregates are selected in the
diff --git a/docs/ref/models/expressions.txt b/docs/ref/models/expressions.txt
index 1b6a208d01..7833580ac1 100644
--- a/docs/ref/models/expressions.txt
+++ b/docs/ref/models/expressions.txt
@@ -1095,6 +1095,16 @@ calling the appropriate methods on the wrapped expression.
         :py:data:`NotImplemented` which forces the expression to be computed on
         the database.
 
+    .. attribute:: set_returning
+
+    .. versionadded:: 5.2
+
+        Tells Django that this expression contains a set-returning function,
+        enforcing subquery evaluation. It's used, for example, to allow some
+        Postgres set-returning functions (e.g. ``JSONB_PATH_QUERY``,
+        ``UNNEST``, etc.) to skip optimization and be properly evaluated when
+        annotations spawn rows themselves. Defaults to ``False``.
+
     .. method:: resolve_expression(query=None, allow_joins=True, reuse=None, summarize=False, for_save=False)
 
         Provides the chance to do any preprocessing or validation of
diff --git a/docs/releases/5.2.txt b/docs/releases/5.2.txt
index 02a068e5af..a15e669205 100644
--- a/docs/releases/5.2.txt
+++ b/docs/releases/5.2.txt
@@ -218,6 +218,10 @@ Models
 * Added support for validation of model constraints which use a
   :class:`~django.db.models.GeneratedField`.
 
+* The new :attr:`.Expression.set_returning` attribute specifies that the
+  expression contains a set-returning function, enforcing subquery evaluation.
+  This is necessary for many Postgres set-returning functions.
+
 Requests and Responses
 ~~~~~~~~~~~~~~~~~~~~~~
 
diff --git a/tests/annotations/models.py b/tests/annotations/models.py
index fbb9ca6988..914770d2fe 100644
--- a/tests/annotations/models.py
+++ b/tests/annotations/models.py
@@ -58,3 +58,11 @@ class Company(models.Model):
 class Ticket(models.Model):
     active_at = models.DateTimeField()
     duration = models.DurationField()
+
+
+class JsonModel(models.Model):
+    data = models.JSONField(default=dict, blank=True)
+    id = models.IntegerField(primary_key=True)
+
+    class Meta:
+        required_db_features = {"supports_json_field"}
diff --git a/tests/annotations/tests.py b/tests/annotations/tests.py
index 703847e1dd..29660a827e 100644
--- a/tests/annotations/tests.py
+++ b/tests/annotations/tests.py
@@ -1,7 +1,9 @@
 import datetime
 from decimal import Decimal
+from unittest import skipUnless
 
 from django.core.exceptions import FieldDoesNotExist, FieldError
+from django.db import connection
 from django.db.models import (
     BooleanField,
     Case,
@@ -15,6 +17,7 @@ from django.db.models import (
     FloatField,
     Func,
     IntegerField,
+    JSONField,
     Max,
     OuterRef,
     Q,
@@ -43,6 +46,7 @@ from .models import (
     Company,
     DepartmentStore,
     Employee,
+    JsonModel,
     Publisher,
     Store,
     Ticket,
@@ -1167,6 +1171,23 @@ class NonAggregateAnnotationTestCase(TestCase):
                 with self.assertRaisesMessage(ValueError, msg):
                     Book.objects.annotate(**{crafted_alias: Value(1)})
 
+    @skipUnless(connection.vendor == "postgresql", "PostgreSQL tests")
+    @skipUnlessDBFeature("supports_json_field")
+    def test_set_returning_functions(self):
+        class JSONBPathQuery(Func):
+            function = "jsonb_path_query"
+            output_field = JSONField()
+            set_returning = True
+
+        test_model = JsonModel.objects.create(
+            data={"key": [{"id": 1, "name": "test1"}, {"id": 2, "name": "test2"}]}, id=1
+        )
+        qs = JsonModel.objects.annotate(
+            table_element=JSONBPathQuery("data", Value("$.key[*]"))
+        ).filter(pk=test_model.pk)
+
+        self.assertEqual(qs.count(), len(qs))
+
 
 class AliasTests(TestCase):
     @classmethod