Фильтр по оконной функции или как сделать SQL-запрос из QuerySet Django

20 октября 2020 г.

Олег Я.

API Django ORM прекрасен и удобен в 90% случаев. Кроме того, его можно расширять кастомными QuerySet и Manager'ами. Но когда встает необходимость составить SQL-запрос, к которому Django не приспособлен, ORM оказывается не помощником, а препятствием. К счастью, Django любезно уступает дорогу, позволяя нам писать сырой SQL.

Один из видов запросов, которые Django не позволяет выполнить - это фильтр по оконной функции (window function). Его можно получить, сделав SELECT из подзапроса, что Django не поддерживает.

В этом посте я опишу проблему и покажу, как можно сделать сырой SELECT из готового QuerySet.

Весь код был протестирован на Django 3.1.

TLDR; Краткое содержание

Если у нас есть QuerySet с оконной функцией:

from django.db import models
from django.db.models.functions import window

from catalog.models import Product

products = Product.objects.all().annotate(row_number=models.Window(
    expression=window.RowNumber(),
    partition_by=[models.F('category_id')],
    order_by=[models.F('price').asc()],
)).order_by('price')

Мы можем выполнить фильтр по ней так:

sql, params = products.query.sql_with_params()

count_per_category = 3
count = 10
products_filtered = Product.objects.raw("""
        SELECT * FROM ({}) products_with_row_numbers
        WHERE row_number <= %s
        LIMIT %s
    """.format(sql),
    [*params, count_per_category, count],
)

selected_products = list(products_filtered)

Далее подробно.

Проблема: фильтр по оконной функции

Допустим в нашем интернет-магазине есть модель категории:

class Category(models.Model):
    name = models.CharField(max_length=255)

И модель товара:

class Product(models.Model):
    category = models.ForeignKey(Category, on_delete=models.CASCADE)
    name = models.CharField(max_length=255)
    price = models.DecimalField(max_digits=12, decimal_places=2)

И перед нами стоит стоит задача сделать выборку из 10 самых недорогих товаров магазина. Не вопрос! Такой запрос можно составить в два счета:

products = Product.objects.all().order_by('price')[:10]

Но вот загвоздка: выборка должна представлять весь разнообразный ассортимент магазина. Будет не хорошо, если все 10 полученных товаров будут из категории аксессуаров, нам нужно показать, какие у нас продаются смартфоны и планшеты. Поэтому перед нами ставят еще одно условие: в полученной выборке должно быть не больше 3 товаров из одной и той же категории.

Если вы используете подходящую СУБД, вы можете решить эту задачу с помощью оконной функции. Оконные функции - полезная вещь, но объяснение их работы не входит в тему в этого поста, поэтому я просто продемонстрирую, что получилось:

from django.db import models
from django.db.models.functions import window

products = Product.objects.all().annotate(row_number=models.Window(
    expression=window.RowNumber(),
    partition_by=[models.F('category_id')],
    order_by=[models.F('price').asc()],
)).order_by('price')

На PostgreSQL этот запрос выглядит так:

SELECT
    "catalog_product"."id",
    "catalog_product"."category_id",
    "catalog_product"."name",
    "catalog_product"."price",
    ROW_NUMBER() OVER (
        PARTITION BY "catalog_product"."category_id"
        ORDER BY "catalog_product"."price" ASC
    ) AS "row_number"
FROM
    "catalog_product"
ORDER BY
    "catalog_product"."price" ASC

С помощью этого запроса мы дописываем в каждый товар атрибут row_number - это порядковый номер товара среди самых дешевых товаров своей категории.

То есть у самого первого смартфона в списке row_number будет равен 1. У следующего - равен 2. В то же время у первого планшета row_number тоже будет 1, у второго - 2, и т.д.

Кажется, что проблема решена, да? Осталось только добавить фильтр по этому полю, чтобы в выборку не попадали 4-е, 5-е и так далее товары. Наивно и самонадеянно мы выполняем запрос:

products = products.filter(row_number__lte=3)

...и получаем исключение:

django.db.utils.NotSupportedError: Window is disallowed in the filter clause.

Оказывается, ни Django, ни даже PostgreSQL не позволяют фильтровать строки по результату оконной функции!

Однако в PostgreSQL можно обойти это ограничение, обернув запрос в новый SELECT:

SELECT *
FROM
    (
        SELECT
            "catalog_product"."id",
            "catalog_product"."category_id",
            "catalog_product"."name",
            "catalog_product"."price",
            ROW_NUMBER() OVER (
                PARTITION BY "catalog_product"."category_id"
                ORDER BY "catalog_product"."price" ASC
            ) AS "row_number"
        FROM
            "catalog_product"
        ORDER BY
            "catalog_product"."price" ASC
    ) subquery
WHERE subquery.row_number < 3

Обратите внимание, что мы вычисляем колонку row_number во внутреннем запросе, а фильтр по этой колонке добавляем во внешний запрос.

Как теперь составить подобный запрос на Django? К сожалению, никак. Точно не средствами ORM. Мне было бы интересно увидеть, как такое можно провернуть без сырого SQL, но я не знаю такого способа.

Django предоставляет Subquery, но мы не можем использовать его тут. Ведь нам придется дописать во внутренний запрос .filter(pk=OuterRef('pk')), и тогда все row_number приравняются к единице, потому что в каждом подзапросе будет одна строка.

И вот тут перед нам встает необходимость использовать сырой SQL. При этом мы бы не хотели писать запрос на SQL целиком, ведь так мы не сможем интегрировать другие QuerySet в этот запрос, не редактируя SQL.

Пишем фильтр с минимумом SQL

Если вы раньше отлаживали запросы Django, возможно вы узнали про атрибут QuerySet.query, когда пытались напечатать SQL в консоль:

print(products.query)

Получается, мы можем взять значение этого атрибута и обернуть в свой SQL, да? Нет, этого делать не стоит.

QuerySet.query - это объект django.db.models.sql.Query. Печатая его, мы преобразуем его в строку. Смотрим документацию метода Query.__str__:

def __str__(self):
    """
    Return the query as a string of SQL with the parameter values
    substituted in (use sql_with_params() to see the unsubstituted string).

    Parameter values won't necessarily be quoted correctly, since that is
    done by the database interface at execution time.
    """

То есть str(products.query) не всегда возвращает валидный SQL.

Вместо этого следует использовать метод products.query.sql_with_params. Он возвращает текст SQL отдельно от списка его параметров:

sql, params = products.query.sql_with_params()

Мы можем использовать текст SQL для формирования своего запроса, а параметры передать Django во время его выполнения.

products_filtered = Product.objects.raw("""
        SELECT * FROM ({}) products_with_row_numbers
        WHERE row_number <= 3
        LIMIT 10
    """.format(sql),
    params,
)

3 и 10 можно вынести в параметры нашего запроса, важно только положить их в правильном порядке после params:

count_per_category = 3
count = 10
products_filtered = Product.objects.raw("""
        SELECT * FROM ({}) products_with_row_numbers
        WHERE row_number <= %s
        LIMIT %s
    """.format(sql),
    [*params, count_per_category, count],
)

Обратите внимание, что подзапрос мы вставляем в наш SQL с помощью маркера {} и метода str.format (потому что это не параметр запроса, а часть его текста), а параметры записываем с помощью %s, как этого требует Django.

На этом все. products_filtered будет содержать перечисление из товаров, подходящих нашим критериям.

Полный код выглядит так:

from django.db import models
from django.db.models.functions import window

from catalog.models import Product

products = Product.objects.all().annotate(row_number=models.Window(
    expression=window.RowNumber(),
    partition_by=[models.F('category_id')],
    order_by=[models.F('price').asc()],
)).order_by('price')

sql, params = products.query.sql_with_params()

count_per_category = 3
count = 10
products_filtered = Product.objects.raw("""
        SELECT * FROM ({}) products_with_row_numbers
        WHERE row_number <= %s
        LIMIT %s
    """.format(sql),
    [*params, count_per_category, count],
)

selected_products = list(products_filtered)

Бонус: Как программно соблюсти одинаковый порядок в QuerySet и в оконной функции

Возможно, вы обратили внимание, что в этом запросе мы указываем порядок товаров дважды: в QuerySet с помощью метода order_by, и в оконной функции в параметре order_by:

products = Product.objects.all().annotate(row_number=models.Window(
    expression=window.RowNumber(),
    partition_by=[models.F('category_id')],
    order_by=[models.F('price').asc()],
)).order_by('price')

Но что если мы получили QuerySet products на вход, и мы не знаем, какой ему был задан порядок? При этом мы хотим, чтобы оконная функция соблюдала его.

Это реализуемо. Порядок QuerySet'а записывается в атрибут QuerySet.query.order_by - это тюпл из строк, которые мы передали методу QuerySet.order_by:

products = Product.objects.all().order_by('price')
print(products.query.order_by)  # ('price',)

products = Product.objects.all().order_by('-price', 'name')
print(products.query.order_by)  # ('-price', 'name')

products = Product.objects.all()
# пустой тюпл - порядок не задан
print(products.query.order_by)  # ()

Нужно преобразовать этот тюпл в список, который поддерживает models.Window. Важно учесть разницу в нотации: убывающий порядок в query.order_by записывается как '-price', но в models.Window нам надо передать models.F('price').desc() - без дефиса.

Код преобразования выглядит так:

order_by = []
for field in products.query.order_by:
    if field.startswith('-'):
        # Убывающий порядок
        desc = True
        field = field[1:]  # убираем дефис
    else:
        # Возрастающий порядок
        desc = False

    order_field = models.F(field)
    if desc:
        order_field = order_field.desc()
    else:
        order_field = order_field.asc()

    order_by.append(order_field)

Получившийся список можно передать в models.Window:

products = products.annotate(row_number=models.Window(
    expression=window.RowNumber(),
    partition_by=[models.F('category_id')],
    order_by=order_by,
))