如何在Django ORM中基于带LIMIT的子查询使用窗口函数分区？

现有Django模型类 #

class A(models.Model):
    a: models.CharField(max_length=10)

需要复现的SQL逻辑 #

SELECT "id"
FROM (
        SELECT id, a
        FROM a_a
        LIMIT 100000
    ) inner
ORDER BY RANK() OVER (PARTITION BY a);

问题说明 #

已知窗口函数的使用方法，但无法通过Django ORM实现基于带LIMIT限制的子查询进行分区。由于数据量过大，不能对全表计算排名，必须通过子查询的LIMIT控制查询的最坏成本，目前使用原生SQL实现，想确认Django ORM是否支持该场景。

解答 #

Django ORM支持这个场景，可通过子查询+窗口函数的组合实现，具体代码如下：

from django.db.models import Window, F
from django.db.models.functions import Rank

# 先构造带LIMIT的子查询，获取前100000条数据的id和a字段
inner_query = A.objects.values('id', 'a')[:100000]

# 基于子查询添加窗口函数排名，并按排名排序，最终只取id字段
result = inner_query.annotate(
    rank=Window(
        expression=Rank(),
        partition_by=[F('a')],
    )
).order_by('rank').values('id')

上述代码生成的SQL会和目标逻辑完全一致：先通过LIMIT截取前100000条数据，再对这部分数据按a字段分区计算排名，最后按排名排序并返回id。

注意：需使用Django 2.0及以上版本，该版本开始正式支持窗口函数特性。

内容的提问来源于stack exchange，提问作者Nuschk