Search code examples
pythondjangodjango-modelsdjango-querysetdjango-annotate

Django querysets optimization - preventing selection of annotated fields


Let's say I have following models:

class Invoice(models.Model):
    ...

class Note(models.Model):
    invoice = models.ForeignKey(Invoice, related_name='notes', on_delete=models.CASCADE)
    text = models.TextField()

and I want to select Invoices that have some notes. I would write it using annotate/Exists like this:

Invoice.objects.annotate(
    has_notes=Exists(Note.objects.filter(invoice_id=OuterRef('pk')))
).filter(has_notes=True)

This works well enough, filters only Invoices with notes. However, this method results in the field being present in the query result, which I don't need and means worse performance (SQL has to execute the subquery 2 times).

I realize I could write this using extra(where=) like this:

Invoice.objects.extra(where=['EXISTS(SELECT 1 FROM note WHERE invoice_id=invoice.id)'])

which would result in the ideal SQL, but in general it is discouraged to use extra / raw SQL. Is there a better way to do this?


Solution

  • You can remove annotations from the SELECT clause using .values() query set method. The trouble with .values() is that you have to enumerate all names you want to keep instead of names you want to skip, and .values() returns dictionaries instead of model instances.

    Django internaly keeps the track of removed annotations in QuerySet.query.annotation_select_mask. So you can use it to tell Django, which annotations to skip even wihout .values():

    class YourQuerySet(QuerySet):
        def mask_annotations(self, *names):
            if self.query.annotation_select_mask is None:
                self.query.set_annotation_mask(set(self.query.annotations.keys()) - set(names))
            else:
                self.query.set_annotation_mask(self.query.annotation_select_mask - set(names))
            return self
    

    Then you can write:

    invoices = (Invoice.objects
      .annotate(has_notes=Exists(Note.objects.filter(invoice_id=OuterRef('pk'))))
      .filter(has_notes=True)
      .mask_annotations('has_notes')
    )
    

    to skip has_notes from the SELECT clause and still geting filtered invoice instances. The resulting SQL query will be something like:

    SELECT invoice.id, invoice.foo FROM invoice
    WHERE EXISTS(SELECT note.id, note.bar FROM notes WHERE note.invoice_id = invoice.id) = True
    

    Just note that annotation_select_mask is internal Django API that can change in future versions without a warning.