Django Performance Optimization [1]: ORM N+1 Problem

在 iCHEF 工作這幾年，其中 2019, 2020 兩年有幸碰到好幾次優化的工作，其中一些滿有趣的，覺得可以記錄下來。

打算從最常見也最簡單的問題開始講起，同時也會聊聊自己通常是怎麼處理這些效能議題的。

第一篇文章會討論的是 ORM 的 N+1 問題，以 Django 當作範例

ORM (Object Relational Mapping) 幫我們的應用程式抽象了不同的資料來源 (資料庫)，我們透過操作不同語言或框架下的 ORM API 來操作被抽象的資料源。

Error = (More Code)²

根據軟體的黃金守則，Error = (More Code)² ，我們可以知道，引入更多的中介層，代表更多意外發生的可能性。的確如此，ORM 在一些狀況下可能組合出效能有問題的 Query，需要我們對 ORM 和 SQL有跟更多的一些理解才能夠處理。

N+1 問題指的是，你有 Size 是 O(N)的資料，可以拿 1 發 SQL Query 解決的事，ORM 實際上卻幫你執行了 N+1 發 Query。

假設我現在有這樣的 Django ORM Model

class Author(models.Model):
    name = models.CharField(max_length=100)


class Book(models.Model):
    title = models.CharField(max_length=100)
    author = models.ForeignKey(Author, on_delete=models.CASCADE, related_name="books")

再看一段 Code，然後想想可能會有什麼問題

@ProfileDecorator.capture_query_count
@ProfileDecorator.do_line_profile
def retrieve_data():
    # A function to emphasize N+1 problem

    for _ in range(1000):
        arr = [random.randint(1, 1000) for _ in range(100)]
        _ = sorted(arr)

    author_queryset = Author.objects.all()

    resp = []

    for author in author_queryset:
        resp.append({
            "name": author.name,
            "books": [book.title for book in author.books.all()]
        })

    return resp

看出來了嗎，問題會出現在下面這一句

[book.title for book in author.books.all()]

# Possible Generated SQL may be like
"""
select id, title, author_id from book where author_id = 1;
select id, title, author_id from book where author_id = 2;
select id, title, author_id from book where author_id = 3;
... etc
"""

你有幾個 Author，Django 就會幫你去 Query 幾次他寫的的 Books

這個時候我們可能還會有一些問題，比如說。
1. 現實中的函式可能會更大更複雜
2. 那我是不是要先唸很多書，才有辦法從函式裡面想出可能拖慢效能的地方
3. 會不會其他地方跑得更慢

這些問題都是合理的，那我必須說，真的不要通靈，因為你根本不能確定你那天有沒有睡飽，查克拉夠不夠。最直接的辦法就是

直接去 Profile 他，證明問題在那裡

為此，我們會用到 line_profiler ，同時也計算 DB query number 作為佐證。
我先準備了 1000 個 Author，讓他們每人寫個 1 ~ 4 本書。然後會把下面兩個 Decorator 掛到上面那個有問題的函式。

import line_profiler

from django.test.utils import CaptureQueriesContext
from django.db import connection


class ProfileDecorator:
    @staticmethod
    def do_line_profile(func):

        @wraps(func)
        def wrapped(*args, **kwargs):
            prof = line_profiler.LineProfiler()
            f_wrapped = prof(func)
            resp = f_wrapped(*args, **kwargs)
            prof.print_stats()

            return resp

        return wrapped

    @staticmethod
    def capture_query_count(func):

        @wraps(func)
        def wrapped(*args, **kwargs):
            with CaptureQueriesContext(connection) as queries:
                resp = func(*args, **kwargs)
                print(f'Captured Query Count: {len(queries)}')
                return resp

        return wrapped

然後跑一下，我們就可以看到我們需要的證據

首先，你看到總時長 5.76 sec，如果對一隻 API 來說，顯然是不能接受的。
同時往下找，你可以看到 96% 的時間集中在 query books。

前面說到，這裡總共有 1000 個 Author，可以看到 query 數量剛好就是 1001 次。隨著資料量呈現性增長的 query 數量顯然是不能接受的，這就是 N+1 問題，他能夠成為常數嗎？可以，事實上，只要花費 2 發 query 就足夠了。

上面那個迭代了 1000 次的 For 迴圈是對比，嘗試模擬稍微耗時的 CPU 計算。你可看到儘管每次都呼叫 1000 次 randint，然後排序，和 DB 的 query 比起來依然微不足道。這是源於存取暫存器和記憶體的速度，和存取外部儲存裝置 (資料庫位於磁碟上的資料)完全不在同一個量級。這裡可以看到，呼叫一次 randint 計算只花了大概 150ns，但存取一次DB，卻要 5ms，兩者差了 30 多倍。應用程式效能有問題的時候，如果不是計算密集的應用，那通常出在 I/O 的機會就很高。

我們現在看修復版本的程式

@ProfileDecorator.capture_query_count
@ProfileDecorator.do_line_profile
def retrieve_data_fixed_version():
    # A version that fix N+1 problem 
    #   by using prefetch_related to preload data into memory
    for _ in range(1000):
        arr = [random.randint(1, 1000) for _ in range(100)]
        _ = sorted(arr)

    # 藉由 prefetch_related，我們把相關的 Books 預先用一發 query 載入記憶體
    author_queryset = Author.objects.prefetch_related("books").all()

    resp = []

    for author in author_queryset:
        resp.append({
            "name": author.name,
            "books": [book.title for book in author.books.all()]
        })

    return resp

Django prefetch_related Document
透過預先載入 Books 到 ORM 的快取中，我們可以把 Query 控制在剛好 2 個

# 藉由 prefetch_related，我們把相關的 Books 預先用一發 query 載入記憶體
author_queryset = Author.objects.prefetch_related("books").all()

# Possible Generated SQL just 2
"""
select id, name from author;
select id, title, author_id from book where author_id in (...);
"""

然後我們再來 Profile 他一次

Ya ! 現在他變成 0.35 sec了，大概快了 13 倍。你可以看到現在 DB Query所花的時間已經大抵和 CPU 運算差不多，他只戳了 DB 2 次 !

當然 0.35 可能還是有點慢，但這是另一個課題了，你可以把 profile 的 overhead 拿掉再跑一次，可能會快一些。

拿掉 profiler 跑，又快了一倍，現在比較能夠接受了。

另一個方向的 N+1

記得我們剛剛搜尋的方向嗎，是從 Author -> Book，而 Foreign Key 是從 Book -> Author。如果搜尋的方向變成 Book -> Author，也就是順著 Foreign Key 的方向，那麼理論上 1 發帶有 join 的 Query 就能搞定。

book_query = Book.objects.filter(name__contain="The Ring")
resp = [
    {
        "title": book.name,
        "auothor_name": book.author.name
    }
    for book in book_query
]

# 然而實際上，我們看到的可能是
# Select title, author_id from book where name like "%The Ring%"
# Select name from author where id = 1
# Select name from author where id = 2
# Select name from author where id = 3

對的，我們看到 ORM 又打出了多餘的 Query 了，他不知道要 join，所以這邊我們提醒他要 join

Django select_related Document

book_query = Book.objects.select_related().filter(name__contain="The Ring")
resp = [
    {
        "title": book.name,
        "auothor_name": book.author.name
    }
    for book in book_query
]

# Select book.title, book.author, author.name 
# from book
# join author where author.id = book.author_id
# where book.title like "%The Ring%"

預防勝於治療

如果我們能在側是階段就防守問題的發生，不只能為未來省下許多麻煩 (跟同事來來回回還有 profiling 看問題)，也能在一開始就提交性能更可靠的程式碼。

當然不是每個程式片段都去測試效能，我想這是不切實際的。

我自己的習慣上，如果一個函式要從 DB 載入幾百上千 row 以上的資料，並且掃描過 2 個以上的 table，那可能就有好好為他加上測試的需要。當然，實際情況可以取決團隊的文化和共識。

預先測試你的程式片段會打出幾發 Query 可能是個好主意。

from django.test.utils import CaptureQueriesContext
from django.db import connections

expected_count = 2

with CaptureQueriesContext(connections["my_datavase_name"]) as db_context:

  tested_func() 

  sql_count = len(db_context.captured_queries)
  assert sql_count < expected_count, f"DB query more than {expected_count}"

善用手邊的工具

Django 這樣的框架提供了很棒的 toolbox，如果你很確定在看 SQL 的問題，也可以很快從這邊得到值得參考的資料，他也有附上 backtrace。

結語

到這邊就差不多了，希望你會喜歡，也有學到東西。
總結就是應用程式的效能問題上，通常出現在 I/O 的機會很高，今天介紹了一種問題的形式，叫做 N+1。
就跟你在寫 leetcode 的時候，常常順順的就可以解出 O(N) 的作法，但是其實要求要做到 O(1) 的感覺很像 (並沒有)。
下一篇會聊聊另外一個效能議題，我們來探討 Python 下沒有被 GC 即時回收的 object 可能造成什麼問題