Django: Annotate model with a prefix sum

Question

I have a model TimeTrial that aggregates an ordered list of Leg s.

class TimeTrial(models.Model):
    objects = TimeTrialManager()

class Leg(models.Model):
    timetrial = models.ForeignKey(TimeTrial)
    duration = models.FloatField()
    order = models.PositiveIntegerField(default=0)

    @property
    def duration_prefix_sum(self):
        qs = Leg.objects.filter(timetrial=self.timetrial,
                                order__lte=self.order)
        return qs.aggregate(Sum('duration'))

class TimeTrialManager(models.Manager):
    def get_queryset(self):
        qs = super(TimeTrialManager, self).get_queryset()
        qs = qs.annotate(leg_count=Count('leg'))
        qs = qs.annotate(duration=Sum('leg__duration'))
        return qs

(Stripped-down version of the actual models .)

The TimeTrialManager annotates each TimeTrial with the number of Leg s and the sum of Leg.duration . I would like to compute the duration_prefix_sum using Django's annotate / aggregate in a similar way, so I don't incur a database hit when I display every leg's duration_prefix_sum in timetrial.leg_set.all() .

I know I can write Python code to compute it:

timetrial.leg_set_annotated = []
s = 0
for l in timetrial.leg_set.all():
    s += l.duration
    l.duration_prefix_sum_computed = s
    timetrial.leg_set_annotated.append(l)

But I would rather use annotate/aggregate/Case/When of Django.

Answer 1

I've found a solution, but I'm not sure it's entirely robust. I've implemented it in my project .

We add the filter timetrial__leg__order__lte=F('order') to add a join containing all Legs with order less or equal to the current. Then, we add the annotation duration_prefix_sum as Sum('timetrial__leg__duration') , which computes the prefix sum.

from django.db.models import Sum, F
qs = Leg.objects.filter(timetrial__leg__order__lte=F('order'))
qs = qs.annotate(duration_prefix_sum=Sum('timetrial__leg__duration'))

I think this will fail if the code that uses this queryset makes another annotation/filter on 'timetrial__leg' , so additions to make this more robust are welcome.