Django - Multiple annotations with Prefetch objects

Question

I have a model test with two m2m fields: foo and bar.

I'm trying to annotate a conditional count of these related fields, i.e. to count the related objects that meet a certain condition. Retrieving this information outside of the queryset is not an option, because I need to use the annotated fields to order the result.

I've tried the following:

1. Using Prefetch objects

from django.db.models import Prefetch, Count

prefetch_foo = Prefetch('foo_set', queryset=foo.objects.filter(<some condition>))
prefetch_bar = Prefetch('bar_set', queryset=bar.objects.filter(<some condition>))
result = test.objects.prefetch_related(prefetch_foo, prefetch_bar).annotate(n_foo=Count('foo'), n_bar=Count('bar'))

This doesn't work, because prefetch_related gets applied after annotate.

2. Using conditional expressions

from django.db.models import Sum, Case, When
from django.db.models.fields import IntegerField

foo_sum = Sum(Case(When(foo__<some condition>, then=1), default=0,
                        output_field=IntegerField())))
bar_sum = Sum(Case(When(bar__<some condition>, then=1), default=0,
                        output_field=IntegerField())))
result = test.objects.annotate(n_foo=foo_sum, n_bar=bar_sum)

This doesn't work because of this bug on multiple Sum annotations: https://code.djangoproject.com/ticket/10060

3. Using RawSQL

sql = "SELECT SUM(CASE WHEN foo.<condition> "
      "THEN 1 ELSE 0 END) FROM app_test "
      "LEFT OUTER JOIN app_foo "
      "ON (app_test.id = foo.test_id) "
      "GROUP BY test.id"
result = test.objects.annotate(n_foo=RawSQL(sql, []))
# Same method for bar

I'm stuck here, because this retrieves the SUM for all rows and I can't find a way to add something like "WHERE test.id = <ID of the object the annotation corresponds to>".

Is there any way to get the right single row from the custom SQL? Or another workaround?

Answer 1

The Count function in Django now has a filter parameter that should be what you are looking for.

See documentation.

In this case:

result = test.objects.annotate(n_foo=Count('foo_set', filter=Q(<some condition>)), n_bar=Count('bar_set', filter=Q(<some condition)))

should give the expected result.