Optimized list viewsets for annotations and large amounts of data

[sebastien-grdn]

Hello everyone,

I'd like to have feedback from the community about a performance issue (and a workaround I came up with) I've had in a list endpoint with a lot of annotations (too much for the data to be denormalized inside the model), done inside the get_queryset method, and that are used in a django-filter filterset, and of course in the serializer. I also have a lot of instances to be filtered. That led to minute long requests.

So my first question is: was I doing it wrong? I'm using GenericViewSet, ListModelMixin and the non-model fields used inside the serializer all come from annotations, they're not loaded during the serialization.

I ended up using a new strategy, overriding the get_queryset and list methods (from GenericViewSet and ListModelMixin).

I'm posting the code lower but tl;dr the idea was to:

1. Annotate only the fields that are needed for filtering and ordering (depending on the query params)
2. Using .values_list(), only retrieve ids from the database
3. Paginate the ids list, so I only get the ids of the instances the list endpoint is supposed to return
4. Query the database a second time, only for the given page of ids. In that query, all the annotations are done of course, because they're needed inside the serializer.

With this method I managed to speed up the request about 20x.

Here is the piece of code, as an example implementation of this strategy. Any feedback or questions welcome. Thanks a lot!

class ProductionFileViewSet(GenericViewSet, ListModelMixin):

    models = ProductionFile
    serializer_class = ProductionFileListSerializer
    queryset = ProductionFile.objects.all()
    permission_classes = [IsKattyUser]

    # Filters
    filter_backends = [filters.SearchFilter, DjangoFilterBackend, filters.OrderingFilter]
    # For rest_framework (SearchFilter) -> "search" parameter
    search_fields = [
        "description",
        "company_name",
        "brand_names",
        "category_names",
        "product_names",
        "supplier_name",
        "operator_name",
        "tag_labels",
        "gencods",
        "state",
        "translation_state",
        "request_type",
        "total_time",
    ]
    # For django-filter (DjangoFilterBackend)
    filterset_class = ProductionFileFilter
    # For rest_framework (OrderingFilter)
    ordering_fields = [
        "id",
        "description",
        "company_name",
        "brand_names",
        "product_names",
        "supplier_name",
        "operator_name",
        "tag_labels",
        "gencods",
        "state",
        "translation_state",
        "request_type",
        "production_file_number",
        "order_id",
        "start_date",
    ]
    # This will be used as the default ordering
    ordering = ["-id"]

    # All the annotations are in a class attribute, and they'll be used only if needed
    annotations = {
        "production_file_number": Count("order__production_files"),
        "company_name": F("order__company__name"),
        "brand_names": ArrayAgg(
            "products__brand__name",
            distinct=True,
            filter=Q(products__isnull=False),
        ),
        "product_names": ArrayAgg("products__name", distinct=True, filter=Q(products__isnull=False)),
        "gencods": ArrayAgg("products__gencod", distinct=True, filter=Q(products__isnull=False)),
        "supplier_name": F("supplier__name"),
        "operator_name": Concat(F("operator__first_name"), Value(" "), F("operator__last_name")),
        "tag_labels": ArrayAgg("tags__name", distinct=True, filter=Q(tags__isnull=False)),
        "total_time": Coalesce(Sum("interventions__time"), Value(0)),
        "category_names": ArrayAgg(
            "products__category__name",
            distinct=True,
            filter=Q(products__isnull=False),
        ),
        "translation_state": F("translation_request__state"),
        "billing_slip_id": F("billing_slip__id"),
    }
    # This attribute is used to store the used annotations
    pre_annotations = None

    def get_queryset(self):
        if self.action == "list":
            # Annotate only the fields that will be filtered
            # The rest of the annotations will be done in the list method, after the pagination
            self.pre_annotations = []
            annotations = {}
            for query_param in self.request.query_params:
                if query_param in self.annotations:
                    annotations[query_param] = self.annotations[query_param]
                    self.pre_annotations.append(query_param)

            # If the ordering is on a field that is not annotated, we need to annotate it
            ordering = self.request.query_params.get("ordering", "").lstrip("-")
            if ordering and ordering in self.annotations and ordering not in self.pre_annotations:
                annotations[ordering] = self.annotations[ordering]
                self.pre_annotations.append(ordering)

            # If global search is used, annotate all the fields that are searched (from self.search_fields)
            if "search" in self.request.query_params:
                for key in self.search_fields:
                    if key in self.annotations and key not in self.pre_annotations:
                        annotations[key] = self.annotations[key]
                        self.pre_annotations.append(key)

            return self.queryset.annotate(**annotations)
        return self.queryset

    def list(self, request, *args, **kwargs):
        queryset = self.filter_queryset(self.get_queryset())

        # Extract ids from the queryset
        queryset_ids = [id for id in queryset.values_list("id", flat=True)]
        page_ids = self.paginate_queryset(queryset_ids)

        if page_ids is not None:
            # Filter the queryset on page_ids and order it by the page_ids order
            page = ProductionFile.objects.filter(id__in=page_ids).order_by(
                Case(*[When(pk=pk, then=pos) for pos, pk in enumerate(page_ids)])
            )
            # Add all annotations (since they were lost in the new queryset)
            page = page.annotate(**self.annotations)
            serializer = self.get_serializer(page, many=True)
            return self.get_paginated_response(serializer.data)

        # In the case page is None (should never happen), add missing annotations
        annotations = {}
        for key, value in self.annotations.items():
            if key not in self.pre_annotations:
                annotations[key] = value

        queryset = queryset.annotate(**annotations)
        serializer = self.get_serializer(queryset, many=True)
        return Response(serializer.data)

[auvipy]

as every projects need is different, hence DRF let us to override its built in methods as needed. I think your approach is OK.

[andriijas]

I think you can also speed up by using prefetch_related and select_related

Im considering doing something similar where it only annotates counts when ordering is specified in query string but I know I will get exception bc it wont be possible to put it in ordering_fields

@method_decorator(cache_control(public=True, max_age=30), name="list")
@method_decorator(cache_control(private=True, max_age=0), name="retrieve")
@method_decorator(cache_control(private=True, max_age=0), name="my")
@method_decorator(cache_control(public=True, max_age=600), name="stats")
class SomeThingsViewSet(ModelViewSet):
    queryset = (
        SomeThing.objects.prefetch_related(
            "attachments",
            "kpis__kpi_type",
            "implementations",
            "implementations__user",
            "tags",
            "tags__tag",
            "likes",
        )
        .select_related("user")
        .annotate(
            likes_count=Count("likes", distinct=True),
            implementations_count=Count("implementations", distinct=True),
            comments_count=Count("comments", distinct=True),
        )
        .all()
    )

    filterset_class = SomeThingsFilter
    serializer_class = SomeThingSerializer
    permission_classes = [SomeThingsPermissions]
    search_fields = [
        "title",
        "subtitle",
        "business_problem",
        "additional_information",
    ]

    lookup_field = "slug"

    ordering_fields = [
        "created_at",
        "published_at",
        "updated_at",
        "likes_count",
        "implementations_count",
        "comments_count",
    ]
    ordering = ["-published_at"]

    def get_queryset(self):
        queryset = super().get_queryset()
        if self.action == "list" and not self.request.query_params.get("statuses"):
            return queryset.filter(current_status=SomeThing.Status.PUBLISHED)
        if self.action == "my":
            return queryset.filter(user=self.request.user)
        return queryset