Home > database >  Django - filter queryset until Sum is reached
Django - filter queryset until Sum is reached

Time:12-07

Let's imagine a model called Roll. It stores the outcome of rolling a six-sided die (D6):

class Roll(models.Model):
    outcome = models.PositiveSmallIntegerField('Roll', null=False, blank=False, default=1)

There are many rolls, for example:

print(list(Roll.objects.all().values_list('outcome', flat=True)))
>>> [1, 5, 6, 3, 5, 4, 4, 3, 2]

Now, how should I go about to get the latest N rows whose Sum('outcome') reaches an arbitrary amount, without looping in some way or another, until that amount is reached?

If we imagine the rolls:

pk    | outcome | (accumulated sum)
1     | 3       | 3
2     | 2       | 5
3     | 6       | 11
4     | 1       | 12
5     | 5       | 17
6     | 4       | 21
7     | 3       | 24
8     | 4       | 29
9     | 5       | 34
10    | 1       | 35

and the arbitrary amount 20, then the query should pick the pk 6 as the accumulated sum has now reached the amount needed.

Could something similar to the below work?:

amount = 100
Roll.objects.annotate(
    accumulated_sum=Subquery(
        Roll.objects.filter(
            pk__lte=OuterRef('pk')
        ).values('pk').order_by('pk').annotate(
            sum=Sum('outcome', distinct=True)
        ).values(
            'sum'
        )[:1]
    )
).filter(
    accumulated_sum__gte=amount
)

CodePudding user response:

This might help annotating accumuated_sum:

from django.db.models import Subquery, IntegerField

class SQSum(Subquery):
    output_field = IntegerField()
    template = "(SELECT sum(outcome) from (%(subquery)s) _sum)"

accumulated_subquery = SQSum(
    Roll.objects.filter(
        pk__lte=OuterRef('pk')
    ).values("outcome")
)

Roll.objects.annotate(accumulated_sum=accumulated_subquery).filter(accumulated_sum__gte=amount)
  • Related