Fixed #32776 -- Added support for Array subqueries on PostgreSQL.

5 files changed, 119 insertions, 1 deletions

diff --git a/django/contrib/postgres/expressions.py b/django/contrib/postgres/expressions.py
new file mode 100644
index 0000000000..ea7cbe038d
--- /dev/null
+++ b/django/contrib/postgres/expressions.py
@@ -0,0 +1,14 @@
+from django.contrib.postgres.fields import ArrayField
+from django.db.models import Subquery
+from django.utils.functional import cached_property
+
+
+class ArraySubquery(Subquery):
+    template = 'ARRAY(%(subquery)s)'
+
+    def __init__(self, queryset, **kwargs):
+        super().__init__(queryset, **kwargs)
+
+    @cached_property
+    def output_field(self):
+        return ArrayField(self.query.output_field)
diff --git a/docs/ref/contrib/postgres/expressions.txt b/docs/ref/contrib/postgres/expressions.txt
new file mode 100644
index 0000000000..79590c2393
--- /dev/null
+++ b/docs/ref/contrib/postgres/expressions.txt
@@ -0,0 +1,37 @@
+=====================================
+PostgreSQL specific query expressions
+=====================================
+
+.. module:: django.contrib.postgres.expressions
+   :synopsis: PostgreSQL specific query expressions
+
+These expressions are available from the
+``django.contrib.postgres.expressions`` module.
+
+``ArraySubquery()`` expressions
+===============================
+
+.. class:: ArraySubquery(queryset)
+
+.. versionadded:: 4.0
+
+``ArraySubquery`` is a :class:`~django.db.models.Subquery` that uses the
+PostgreSQL ``ARRAY`` constructor to build a list of values from the queryset,
+which must use :meth:`.QuerySet.values` to return only a single column.
+
+This class differs from :class:`~django.contrib.postgres.aggregates.ArrayAgg`
+in the way that it does not act as an aggregate function and does not require
+an SQL ``GROUP BY`` clause to build the list of values.
+
+For example, if you want to annotate all related books to an author as JSON
+objects::
+
+    >>> from django.db.models import OuterRef
+    >>> from django.db.models.functions import JSONObject
+    >>> from django.contrib.postgres.expressions import ArraySubquery
+    >>> books = Book.objects.filter(author=OuterRef('pk')).values(
+    ...     json=JSONObject(title='title', pages='pages')
+    ... )
+    >>> author = Author.objects.annotate(books=ArraySubquery(books)).first()
+    >>> author.books
+    [{'title': 'Solaris', 'pages': 204}, {'title': 'The Cyberiad', 'pages': 295}]
diff --git a/docs/ref/contrib/postgres/index.txt b/docs/ref/contrib/postgres/index.txt
index 03ff6da1e0..7627043037 100644
--- a/docs/ref/contrib/postgres/index.txt
+++ b/docs/ref/contrib/postgres/index.txt
@@ -30,6 +30,7 @@ a number of PostgreSQL specific data types.
 
     aggregates
     constraints
+    expressions
     fields
     forms
     functions
diff --git a/docs/releases/4.0.txt b/docs/releases/4.0.txt
index 54d0fd4c1d..bc8e1639c8 100644
--- a/docs/releases/4.0.txt
+++ b/docs/releases/4.0.txt
@@ -131,6 +131,11 @@ Minor features
   :class:`~django.contrib.postgres.operations.AddConstraintNotValid` on
   PostgreSQL.
 
+* The new
+  :class:`ArraySubquery() <django.contrib.postgres.expressions.ArraySubquery>`
+  expression allows using subqueries to construct lists of values on
+  PostgreSQL.
+
 :mod:`django.contrib.redirects`
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
diff --git a/tests/postgres_tests/test_array.py b/tests/postgres_tests/test_array.py
index 6fb34b3d92..62a495d5a7 100644
--- a/tests/postgres_tests/test_array.py
+++ b/tests/postgres_tests/test_array.py
@@ -10,7 +10,7 @@ from django.core.exceptions import FieldError
 from django.core.management import call_command
 from django.db import IntegrityError, connection, models
 from django.db.models.expressions import Exists, OuterRef, RawSQL, Value
-from django.db.models.functions import Cast, Upper
+from django.db.models.functions import Cast, JSONObject, Upper
 from django.test import TransactionTestCase, modify_settings, override_settings
 from django.test.utils import isolate_apps
 from django.utils import timezone
@@ -28,6 +28,7 @@ try:
     from psycopg2.extras import NumericRange
 
     from django.contrib.postgres.aggregates import ArrayAgg
+    from django.contrib.postgres.expressions import ArraySubquery
     from django.contrib.postgres.fields import ArrayField
     from django.contrib.postgres.fields.array import (
         IndexTransform, SliceTransform,
@@ -551,6 +552,66 @@ class TestQuerying(PostgreSQLTestCase):
             1,
         )
 
+    def test_filter_by_array_subquery(self):
+        inner_qs = NullableIntegerArrayModel.objects.filter(
+            field__len=models.OuterRef('field__len'),
+        ).values('field')
+        self.assertSequenceEqual(
+            NullableIntegerArrayModel.objects.alias(
+                same_sized_fields=ArraySubquery(inner_qs),
+            ).filter(same_sized_fields__len__gt=1),
+            self.objs[0:2],
+        )
+
+    def test_annotated_array_subquery(self):
+        inner_qs = NullableIntegerArrayModel.objects.exclude(
+            pk=models.OuterRef('pk')
+        ).values('order')
+        self.assertSequenceEqual(
+            NullableIntegerArrayModel.objects.annotate(
+                sibling_ids=ArraySubquery(inner_qs),
+            ).get(order=1).sibling_ids,
+            [2, 3, 4, 5],
+        )
+
+    def test_group_by_with_annotated_array_subquery(self):
+        inner_qs = NullableIntegerArrayModel.objects.exclude(
+            pk=models.OuterRef('pk')
+        ).values('order')
+        self.assertSequenceEqual(
+            NullableIntegerArrayModel.objects.annotate(
+                sibling_ids=ArraySubquery(inner_qs),
+                sibling_count=models.Max('sibling_ids__len'),
+            ).values_list('sibling_count', flat=True),
+            [len(self.objs) - 1] * len(self.objs),
+        )
+
+    def test_annotated_ordered_array_subquery(self):
+        inner_qs = NullableIntegerArrayModel.objects.order_by('-order').values('order')
+        self.assertSequenceEqual(
+            NullableIntegerArrayModel.objects.annotate(
+                ids=ArraySubquery(inner_qs),
+            ).first().ids,
+            [5, 4, 3, 2, 1],
+        )
+
+    def test_annotated_array_subquery_with_json_objects(self):
+        inner_qs = NullableIntegerArrayModel.objects.exclude(
+            pk=models.OuterRef('pk')
+        ).values(json=JSONObject(order='order', field='field'))
+        siblings_json = NullableIntegerArrayModel.objects.annotate(
+            siblings_json=ArraySubquery(inner_qs),
+        ).values_list('siblings_json', flat=True).get(order=1)
+        self.assertSequenceEqual(
+            siblings_json,
+            [
+                {'field': [2], 'order': 2},
+                {'field': [2, 3], 'order': 3},
+                {'field': [20, 30, 40], 'order': 4},
+                {'field': None, 'order': 5},
+            ],
+        )
+
 
 class TestDateTimeExactQuerying(PostgreSQLTestCase):