CODEWITHSUNDEEP
pythondjangodjango-querysetdjango-orm
funnydman
funnydman

Reputation: 11366

How to collect results into array from annotation?

Is this possible to collect results from different Case directives and store them in an array field?

Several intances for test:

banana = Fruit.objects.create(name='banana', type='tropicals', country_of_import='Africa')
orange = Fruit.objects.create(name='orange', type='citrus', country_of_import='Spain')
apple = Fruit.objects.create(name='apple', type='apples', country_of_import='Russia')

Here is the Manager that needs to collect all results from cases. ToArray as a possible function that does this work. There is a function with a similar name ArrayAgg according to the docs:

Returns a list of values, including nulls, concatenated into an array.

It works fine with one expression, but not with several:

class CustomManager(Manager):

    def get_queryset(self):
        query = super().get_queryset().annotate(
            additional_info=ToArray(
                Case(
                    When(
                        type='tropicals',
                        then=Value('This fruit is tropical...')
                    ),
                    output_field=CharField()
                ),
                Case(
                    When(
                        country_of_import='Spain',
                        then=Value('This fruits was grown in Spain'),
                    ),
                    output_field=CharField()
                )
            )
        )
        return query

And in the end the result would be:

banana.additional_info = ['This fruit is tropical.']
orange.additional_info = ['This fruits was grown in Spain.']

Is this possible to handle with existing function or do I need to write my own?

Upvotes: 0

Views: 1760

Answers (1)

funnydman
funnydman

Reputation: 11366

Actually, there is no obvious way to implement this. The most simple solution would be to use CombinedExpression class like this:

from django.contrib.postgres.aggregates import ArrayAgg
from django.db.models import Manager, Case, When, Value, CharField

from django.db.models.expressions import CombinedExpression

class FruitManager(Manager):
    def get_queryset(self):
        query = super().get_queryset()

        query = query.annotate(
            result=(
                CombinedExpression(
                    ArrayAgg(Case(
                        When(
                            type='tropicals',
                            then=Value('This fruit is tropical...'),

                        ),
                        output_field=CharField()
                    )),
                    '||'
                    ,
                    ArrayAgg(Case(
                        When(
                            country_of_import='Africa',
                            then=Value('This fruit is citrus...'),

                        ),
                        output_field=CharField(),
                        default=Value('default value')),
                    ),
                )
            )
        )
        return query

That will be compiled to the following SQL:

SELECT "core_fruit"."id",
       (
               ARRAY_AGG(CASE
                             WHEN "core_fruit"."type" = 'tropicals'
                                 THEN 'This fruit is tropical...'
                             ELSE NULL END) ||
               ARRAY_AGG(
                       CASE
                           WHEN "core_fruit"."country_of_import" = 'Africa'
                               THEN 'This fruit is citrus...'
                           ELSE 'default value' END
                   )
           ) AS "result"
FROM "core_fruit"
GROUP BY "core_fruit"."id";

But there is one thing worth to note, CombinedExpression isn't documented because intended only for internal usage.

Upvotes: 1

Related Questions