Order Django queryset by annotated field and annotate it again

330 Views Asked by At

I need to order Django queryset by annotated field and then annotate result with row number.

Steps:

  1. Add complex_order field to QuerySet that defines extra logic for future ordering.
  2. Order QuerySet by new field and old fields (complex_order, order and id)
  3. Annotate sorted QuerySet with number using RowNumber() function.

Sample code that almost works

queryset.annotate(
    complex_order=Case(
        When(
            Q(foo__type=FooType.ONE.value) | Q(foo__isnull=True)
            then=Value(1),
        ),
        default=Value(0),
    ),
).order_by(
    '-complex_order', 'order', 'id'
).annotate(
    number=Window(expression=RowNumber())
)

But I have a problem with wrong number annotation. After inspecting the SQL query I noticed that annotation happens at the same time with ordering.

SELECT "table.id", ... 
CASE WHEN ("foo"."type" = ONE OR ("foo" IS NULL)) 
THEN 1 ELSE 0 END AS "complex_order", ROW_NUMBER() OVER () AS "number" 
FROM ...
ORDER BY 32 DESC, "table"."order" ASC, "table"."id" ASC

I need to annotate queryset second time, right after sorting it. Or maybe there is a better way to add iterable number for each row in queryset?

1

There are 1 best solutions below

0
Sagit Khaliullin On

The solution was a bit different than I thought.

queryset.annotate(
    complex_order=Case(
        When(
            Q(foo__type=FooType.ONE.value) | Q(foo__isnull=True)
            then=Value(1),
        ),
        default=Value(0),
    ),
).annotate(
    number=Window(expression=RowNumber()),
    order_by=[F('-complex_order'), F('order'), F('id')],
).order_by(
    'number',
)

Firstly we annotate QuerySet with complex_order, then annotate it with number using order_by inside window function. And finally we order QuerySet by number field.

Result examples:

Wrong number annotation (result of code from question)

<QuerySet [
    {'order': 1, 'id': 4387, 'complex_order': 1, 'number': 5}, 
    {'order': 2, 'id': 4388, 'complex_order': 0, 'number': 1}, 
    {'order': 3, 'id': 4389, 'complex_order': 0, 'number': 2}, 
    {'order': 4, 'id': 4390, 'complex_order': 0, 'number': 3}, 
    {'order': 5, 'id': 4391, 'complex_order': 0, 'number': 4}
]>

Correct number annotation (result of code from answer)

<QuerySet [
    {'order': 1, 'id': 4387, 'complex_order': 1, 'number': 1}, 
    {'order': 2, 'id': 4388, 'complex_order': 0, 'number': 2}, 
    {'order': 3, 'id': 4389, 'complex_order': 0, 'number': 3}, 
    {'order': 4, 'id': 4390, 'complex_order': 0, 'number': 4}, 
    {'order': 5, 'id': 4391, 'complex_order': 0, 'number': 5}
]>