Search code examples
djangodjango-modelsdjango-querysetdjango-aggregation

Annotating Django querysets with counts using multiple fields


Here is a simplified version of my model:

class Flight(models.Model):
    airline = models.CharField(max_length=100)
    origin = models.CharField(max_length=4)
    destination = models.CharField(max_length=4)

What I would like to do is group the Flight objects by their common values for certain fields and annotate the groups with the corresponding flight count, as described in the documentation here.

I know how to do this when only considering one field. For example, with

Flight.objects.values('airline').annotate(Count('id')).order_by('-id__count')

I get something like this:

[{'airline': 'First Airlines', 'id__count': 21}, 
 {'airline': 'Air Second', 'id__count': 6},
 {'airline': 'Third Airways', 'id__count': 3}, ...]

This means there are 21 Flight objects whose airline field is 'First Airlines', and so on.

But how do I generalize to count multiple fields and combine into one annotated queryset that treats values from the different fields as if they were the same field?

For example, suppose I have 3 flights from Los Angeles to New York and 2 return flights the other way (there are 3 Flight objects with origin 'LAX' and destination 'JFK', and 2 Flight objects with origin 'JFK' and destination 'LAX'). How can I get this:

[{'airport': 'LAX', 'id__count': 5}, 
 {'airport': 'JFK', 'id__count': 5}]

Solution

  • Well, you can't do it in a single query. You need to first annotate on origin, then on destination, then sum up:

    data1 = Flight.objects.values('origin').annotate(Count('id'))
    data2 = Flight.objects.values('destination').annotate(Count('id'))
    data = {}
    for airport in { x for x in data1.keys() + data2.keys() }:
        data[airport] = data1.get(airport, 0) + data2.get(airport, 0)