Aggregate (and other annotated) fields in Django Rest Framework serializers

Question

I am trying to figure out the best way to add annotated fields, such as any aggregated (calculated) fields to DRF (Model)Serializers. My use case is simply a situation where an endpoint returns fields that are NOT stored in a database but calculated from a database.

Let's look at the following example:

models.py

class IceCreamCompany(models.Model):
    name = models.CharField(primary_key = True, max_length = 255)

class IceCreamTruck(models.Model):
    company = models.ForeignKey('IceCreamCompany', related_name='trucks')
    capacity = models.IntegerField()

serializers.py

class IceCreamCompanySerializer(serializers.ModelSerializer):
    class Meta:
        model = IceCreamCompany

desired JSON output:

[

    {
        "name": "Pete's Ice Cream",
        "total_trucks": 20,
        "total_capacity": 4000
    },
    ...
]

I have a couple solutions that work, but each have some issues.

Option 1: add getters to model and use SerializerMethodFields

models.py

class IceCreamCompany(models.Model):
    name = models.CharField(primary_key=True, max_length=255)

    def get_total_trucks(self):
        return self.trucks.count()

    def get_total_capacity(self):
        return self.trucks.aggregate(Sum('capacity'))['capacity__sum']

serializers.py

class IceCreamCompanySerializer(serializers.ModelSerializer):

    def get_total_trucks(self, obj):
        return obj.get_total_trucks

    def get_total_capacity(self, obj):
        return obj.get_total_capacity

    total_trucks = SerializerMethodField()
    total_capacity = SerializerMethodField()

    class Meta:
        model = IceCreamCompany
        fields = ('name', 'total_trucks', 'total_capacity')

The above code can perhaps be refactored a bit, but it won't change the fact that this option will perform 2 extra SQL queries per IceCreamCompany which is not very efficient.

Option 2: annotate in ViewSet.get_queryset

models.py as originally described.

views.py

class IceCreamCompanyViewSet(viewsets.ModelViewSet):
    queryset = IceCreamCompany.objects.all()
    serializer_class = IceCreamCompanySerializer

    def get_queryset(self):
        return IceCreamCompany.objects.annotate(
            total_trucks = Count('trucks'),
            total_capacity = Sum('trucks__capacity')
        )

This will get the aggregated fields in a single SQL query but I'm not sure how I would add them to the Serializer as DRF doesn't magically know that I've annotated these fields in the QuerySet. If I add total_trucks and total_capacity to the serializer, it will throw an error about these fields not being present on the Model.

Option 2 can be made work without a serializer by using a View but if the model contains a lot of fields, and only some are required to be in the JSON, it would be a somewhat ugly hack to build the endpoint without a serializer.

Quadrille

related: https://mcmap.net/q/242729/-how-to-add-annotate-data-in-django-rest-framework-queryset-responses – Stesha

Answer 1

Possible solution:

views.py

class IceCreamCompanyViewSet(viewsets.ModelViewSet):
    queryset = IceCreamCompany.objects.all()
    serializer_class = IceCreamCompanySerializer

    def get_queryset(self):
        return IceCreamCompany.objects.annotate(
            total_trucks=Count('trucks'),
            total_capacity=Sum('trucks__capacity')
        )

serializers.py

class IceCreamCompanySerializer(serializers.ModelSerializer):
    total_trucks = serializers.IntegerField()
    total_capacity = serializers.IntegerField()

    class Meta:
        model = IceCreamCompany
        fields = ('name', 'total_trucks', 'total_capacity')

By using Serializer fields I got a small example to work. The fields must be declared as the serializer's class attributes so DRF won't throw an error about them not existing in the IceCreamCompany model.

Quadrille

That's the recommended solution. DRF can't introspect the fields so you must specify them manually. – Barrington

I was hoping for a solution where I could modify the queryset in the serializer before the SQL actually fires off. Is there no point in the serializer lifecycle where I could insert the annotation hook? It would be nice to be able to serialize any queryset built off the IceCreamCompany model without having to manually annotate... – Housebound

I'm trying to achieve the same result but I can't figure how to use the ViewSet. Could you detail how the IceCreamCompanyViewSet is then used in your (or theoretical) code? – Ostosis

@RyanPergent in DRF, ViewSets are used by registering them with a router so they can handle requests: django-rest-framework.org/api-guide/routers – Quadrille

Hey @Coderer, any chance you arrived at a solution for adding a hook into a serializer for the incoming queryset? I'm struggling with this now, and surprised there doesn't seem to be a clear place to do this... – Luciusluck

Answer 2

I made a slight simplification of elnygreen's answer by annotating the queryset when I defined it. Then I don't need to override get_queryset().

# views.py
class IceCreamCompanyViewSet(viewsets.ModelViewSet):
    queryset = IceCreamCompany.objects.annotate(
            total_trucks=Count('trucks'),
            total_capacity=Sum('trucks__capacity'))
    serializer_class = IceCreamCompanySerializer

# serializers.py
class IceCreamCompanySerializer(serializers.ModelSerializer):
    total_trucks = serializers.IntegerField()
    total_capacity = serializers.IntegerField()

    class Meta:
        model = IceCreamCompany
        fields = ('name', 'total_trucks', 'total_capacity')

As elnygreen said, the fields must be declared as the serializer's class attributes to avoid an error about them not existing in the IceCreamCompany model.

Peh

How would you instantiate such a serializer? I am writing unit tests and IceCreamCompanySerializer(instance=ice_cream_company) is, obviously not working, – Faun

Answer 3

You can hack the ModelSerializer constructor to modify the queryset it's passed by a view or viewset.

class IceCreamCompanySerializer(serializers.ModelSerializer):
    total_trucks = serializers.IntegerField(readonly=True)
    total_capacity = serializers.IntegerField(readonly=True)

    class Meta:
        model = IceCreamCompany
        fields = ('name', 'total_trucks', 'total_capacity')

    def __new__(cls, *args, **kwargs):
        if args and isinstance(args[0], QuerySet):
              queryset = cls._build_queryset(args[0])
              args = (queryset, ) + args[1:]
        return super().__new__(cls, *args, **kwargs)

    @classmethod
    def _build_queryset(cls, queryset):
         # modify the queryset here
         return queryset.annotate(
             total_trucks=...,
             total_capacity=...,
         )

There is no significance in the name _build_queryset (it's not overriding anything), it just allows us to keep the bloat out of the constructor.

Lacewing

Please explain your answer. How does it improve the accepted one? – Debus

@Debus This allows you to keep the annotation with the serializer (rather than the view). The way it should be. See the comments on the other answer. – Honniball

@Honniball I did not mark this question as a low quality one but found it in the review queue. I still think that this answer can be improved adding more details like those that you point above. – Debus

You could probably look at the fields in Meta to get the annotations and add them too. – Concave

Answer 4

Be careful though when combining multiple aggregations (here trucks and capacity) as it will yield the wrong results because joins are used instead of sub queries.

See the django docs about that: combining multiple aggregations

In the doc, they suggest to use distinct = True. For example:

class IceCreamCompanyViewSet(viewsets.ModelViewSet):
queryset = IceCreamCompany.objects.all()
serializer_class = IceCreamCompanySerializer

def get_queryset(self):
    return IceCreamCompany.objects.annotate(
        total_trucks = Count('trucks', distinct=True),
        total_capacity = Sum('trucks__capacity', distinct=True)
    )

Tuttle