[Django]-Django: Annotate sum on variable number of columns

by

5đź‘Ť

âś…

To figure this out, it helps to picture the table first:

column_a | column b | column_c
---------+----------+----------
1        | 2        | 3
4        | 5        | 6
7        | 8        | 9

Sum is a “vertical” operation; that is, if we wanted the sum of column_a, we could do 

>>> DummyModel.objects.aggregate(total=Sum('column_a'))
{'total': 12}

As you can see, this returns 1 + 4 + 7 == 12 — so you can see why I call it a “vertical” sum. Notice we use aggregate rather than annotate: aggregate is for vertical operators.

If instead we want “horizontal” sum — the total across a row — we’d use F() and +. So to get column_a + column_b in each row, we’d use

>>> DummyModel.objects.annotate(total=F('column_a') + F('column_b')).values('total')
<QuerySet [{'total': 3}, {'total': 9}, {'total': 15}]>

Hopefully you can see why I call this a “horizontal” sum: we’re getting a sum of a and b “horizontally” in each row. And now notice we use annotate, which is for horizontal operations.

If the names of the columns aren’t known beforehand, you’d need to get tricky and use functools.reduce and operator.add to build up an expression:

>>> from functools import reduce
>>> from operator import add

>>> cols = ['column_b', 'column_c']
>>> expr = reduce(add, (F(col) for col in cols))
>>> DummyModel.objects.annotate(total=expr).values('total')
<QuerySet [{'total': 5}, {'total': 11}, {'total': 17}]>

If we want both a horizontal and a vertical sum — i.e. the sum of column_a plus the sum of column_b — we need to use Sum and F():

>>> DummyModel.objects.aggregate(total=Sum(F('column_a') + F('column_b')))
{'total': 27}

Notice: aggregate rather than annotate, since we’re ultimately going with a vertical operation; a Sum of rows. Yes, there’s a horizontal operation first, but since we ultimately Sum, we need aggregate.

So, to wrap things up, if the fields are a variable, we need to combine aggregate, Sum, and the reduce trickery from above:

>>> cols = ['column_b', 'column_c']
>>> expr = reduce(add, (F(col) for col in cols))
>>> DummyModel.objects.aggregate(total=Sum(expr))
{'total': 33}

Hope this helps!

👤jacobian

Leave a comment