mizzsugar’s blog

Pythonで学んだことや読書録を書きます。

DjangoORMでウィンドウ関数を使おう

この記事はDjango Advent Calendar 2019の記事です。

Django Advent Calendar 2019 - Qiita

そして私の初のアドベントカレンダーです!

最近仕事でBigQueryで分析関数を使うことが多いので、

そのなかでもウィンドウ関数をDjangoのORMでも使えないかなと思い調べてみました。

ウィンドウ関数とは

データベース製品よっては分析関数と呼ぶところもあればウィンドウ関数と呼ぶところもあります。

この記事ではPostgreSQLを使います。PostgreSQLでは分析関数をウィンドウ関数と呼んでいるのでウィンドウ関数で統一します。

PostgreSQLの公式ドキュメントには下記のように説明されています。

ウィンドウ関数は現在の行に何らかとも関係するテーブル行の集合に渡って計算を行います。 これは集約関数により行われる計算の形式と似たようなものです。 とは言っても、通常の集約関数とは異なり、ウィンドウ関数の使用は単一出力行に行をグループ化しません。 行はそれぞれ個別の身元を維持します。 裏側では、ウィンドウ関数は問い合わせ結果による現在行だけでなく、それ以上の行にアクセスすることができます。

3.5. ウィンドウ関数

例えば、このようなテーブルがある場合。

テーブル定義

    Column     |  Type   | Collation | Nullable |               Default                
---------------+---------+-----------+----------+--------------------------------------
 id            | integer |           | not null | nextval('salaries_id_seq'::regclass)
 department_id | integer |           | not null | 
 salary        | integer |           | not null | 
 employee_id   | integer |           | not null | 
Indexes:
    "salaries_pkey" PRIMARY KEY, btree (id)
    "salaries_departm_f58eef_idx" btree (department_id)
    "salaries_employe_a009a1_idx" btree (employee_id)


通常の集約関数で部署ごとの平均給与を出すとこのようになります。

SELECT
    department_id,
    ROUND(AVG(salary), 0) as avg_salary
FROM
    public.salaries
GROUP BY
    department_id
ORDER BY
    avg_salary DESC


出力結果例

 department_id | avg_salary 
---------------+------------
            16 |    9737374
            12 |    9694656
            15 |    9647917
            11 |    9588957
             5 |    9338559
            14 |    9337349
             6 |    9123184
             8 |    9070103
            13 |    9063918
            10 |    9020534
             9 |    8948148
            20 |    8934615
             4 |    8929108
            18 |    8712230
            17 |    8695489
            19 |    8669903
             2 |    8596131
             3 |    8578168
             7 |    8551880
             1 |    6790198


平均給与を出すだけなら良いですが、一人ひとりの社員の給与とその社員がその部署の中で何番目に給与が高いのかも出したいとします。

出力結果例

 department_id | employee_id |  salary  | rank 
---------------+-------------+----------+------
             1 |         414 | 40000000 |    1
             1 |         255 | 30000000 |    2
             1 |         316 | 20000000 |    3
             2 |         342 | 60000000 |    1
             2 |         673 | 50000000 |    2
             2 |         359 | 40000000 |    3
             3 |         409 | 60000000 |    1
             3 |         476 | 20000000 |    2
             3 |         410 | 10000000 |    3


このような結果を出すには上記のような単純なクエリでは出力できないので工夫が必要です。

そういう要望がある時に、ウィンドウ関数を知っていると、ウィンドウ関数を使わない場合よりもシンプルなクエリで実現できるまたはクエリのコストを抑えられることが多々あります。

ウィンドウ関数を使うとこのようなクエリになります。


SELECT
    department_id,
    employee_id,
    salary,
    rank() OVER (PARTITION BY department_id ORDER BY salary DESC)
FROM public.salaries
;


ウィンドウ関数なしだとこうまります。ウィンドウ関数を使ったクエリに比べてやや複雑です。

SELECT
    salaries.department_id,
    salaries.employee_id,
    salaries.salary,
    (
        SELECT COUNT(*)
        FROM public.salaries as rank
        WHERE
            rank.department_id = salaries.department_id
            AND rank.salary > salaries.salary
    ) + 1,
    rank() OVER (PARTITION BY department_id ORDER BY salary DESC)
FROM public.salaries
;


また、コストもレコード数が10000件のテーブルに関して

ウィンドウ関数を使ったクエリだと1024だったのに対し

ウィンドウ関数なしのクエリだと2072289でした。およそ206倍です。


rank()を説明しますと、PARTITION BYで部署ごとに分け、ORDER BYで部署の中でsalaryが多い順に順番を振り分けます。

ウィンドウ関数のPARTITION BYやORDER BYでどのようにテーブルを振り分けているかについては、下記のBigQueryのドキュメントの図1がわかりやすいです。

BigQueryなので分析関数と書いていますが、PostgreSQLのウィンドウ関数も同じ仕組みです。

cloud.google.com


DjangoORMでWindow関数を使うには

実行環境

* Python 3.8.0
* PostgreSQL 11.4
* Django 2.2.6
* psycopg2-binary 2.8.4


今回使用するモデル(テーブル定義は先程記述したテーブルと同じになります)

models.py

from django.db import models


class Salary(models.Model):
    # 面倒がってFK貼っていませんが実運用ではemployees  departmentsテーブル作ってFK貼ると思います
    employee_id = models.IntegerField()
    department_id = models.IntegerField()
    salary = models.IntegerField()

    class Meta:
        db_table = 'salaries'
        indexes = [
            models.Index(fields=['employee_id']),
            models.Index(fields=['department_id']),
        ]
    
    def __str__(self):
        return f'{self.department_id}_{self.employee_id}'


ウィンドウ関数をORMで表現するには、annotateを使います。

先ほどのクエリをORMで表現すると、下記のようになります。

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

from employee.models import Salary

window =  {
     'partition_by': [F('department_id')],
     'order_by': F('salary').desc()
}


Salary.objects.annotate(
    rank=Window(expression=Rank(), **window)
).values(
    'department_id',
    'employee_id',
    'salary',
    'rank'
)


F()は、モデルのフィールドを表すオブジェクトです。この例の場合、partition_byのdepartment_idとorder_byのsalaryです。

F()式のいいところは、実際にデータベースから値を取り出してPythonのメモリに格納しなくてもモデルのフィールドを参照できるところです。

https://docs.djangoproject.com/ja/2.2/ref/models/expressions/#f-expressions


annotateでウィンドウ関数でどのような項目を出力したいかしていします。引数expressionにDenseRank()を入れると、変数windowで指定したpartition byとorder byの通りにランキングされます。


この例では、rankという名前のアトリビュートにランキングを出力するように書いています。

https://docs.djangoproject.com/ja/2.2/ref/models/database-functions/#denserank


valuesで出力するアトリビュートを指定します。

valuesを指定しないと、id, department_id, employee_id, salary, rankの全てが出力されます。


Django Debug Toolbarをインストールし、

python manage.py debugsqlshellで上記のSQLを実行すると、

ウィンドウ関数が利用されたことを確認できました!

SELECT "salaries"."id",
       "salaries"."employee_id",
       "salaries"."department_id",
       "salaries"."salary",
       RANK() OVER (PARTITION BY "salaries"."department_id"
                    ORDER BY "salaries"."salary" DESC) AS "rank"
FROM "salaries"

https://docs.djangoproject.com/ja/2.2/ref/models/expressions/#window-functions


おまけ

ウィンドウ関数を紹介しましたが、場合によってはウィンドウ関数で表現できるけれどもウィンドウ関数を使わないほうが良い場合があります。


例えば、社員の給与と一緒にその社員が所属している部署の平均給与を表示したい場合です。

下記のような出力結果を期待します。

 department_id | employee_id | salary  | avg_salary 
---------------+-------------+---------+------------
             4 |        1707 | 8999649 |    5492152
             9 |        8296 | 8999531 |    5575559
             1 |        4641 | 8999048 |    5468614
            17 |         222 | 8998863 |    5376142
            19 |        4686 | 8997529 |    5449444
            10 |        1768 | 8994513 |    5588093
            17 |        7013 | 8994161 |    5376142
            15 |        6940 | 8994098 |    5535015
            11 |         577 | 8992925 |    5566160
            10 |        8139 | 8992773 |    5588093
             3 |        7718 | 8992511 |    5595843


ウィンドウ関数を使うと、下記のようなSQLになります。


SQLクエリ

SELECT
    department_id,
    employee_id,
    salary,
    rank() OVER (PARTITION BY department_id ORDER BY salary DESC)
FROM public.salaries
;


ORMクエリ

from django.db.models import F, Window, Avg

from employee.models import Salary


window =  {
    'partition_by': [F('department_id')],
    'order_by': F('salary').desc()
}


Salary.objects.annotate(
    avg_salary=Window(expression=Avg('salary'), **window)
).values(
    'department_id',
    'employee_id',
    'salary',
    'avg_salary'
)


ウィンドウ関数を使わないと、こうなります。

SQLクエリ

SELECT
    salaries.department_id,
    salaries.employee_id,
    salaries.salary,
    salary_averages.avg_salary
FROM
    public.salaries
JOIN
    (
    SELECT
        department_id,
        ROUND(AVG(salary), 0) as avg_salary
    FROM
        public.salaries
    GROUP BY department_id
) as salary_averages
ON
    salaries.department_id = salary_averages.department_id
ORDER BY
    salaries.salary DESC
;


ORMクエリは諦めました...

こういうの見つけたのですが、普通にSQL書いたほうが楽だと思ってしまいました。↓

Self join with django ORM - Stack Overflow


部署ごとの平均給与を出すクエリまではできました。給与が高い順です。

from employee.models import Salary


Salary.objects.values('department_id').annotate(avg_salary=models.Avg('salary')).order_by('-salary')


10000件レコードを入れてEXPLAINでコストを比較したところ、

①のウィンドウ関数を使ったクエリだと1024になったのに対し ②のJOINで対応したクエリだと393になりました。

ウィンドウ関数は便利ですが、その時々で出力したい結果に応じて使い分けようという話でした!


次回は、xKxAxKxさんです!

よろしくお願いします。