1. 程式人生 > 實用技巧 >減輕記憶體負擔,在 pymysql 中使用 SSCursor 查詢結果集較大的 SQL

減輕記憶體負擔,在 pymysql 中使用 SSCursor 查詢結果集較大的 SQL

前言

預設情況下,使用 pymysql 查詢資料使用的遊標類是 Cursor,比如:

import pymysql.cursors

# 連線資料庫
connection = pymysql.connect(host='localhost',
                             user='user',
                             password='passwd',
                             db='db',
                             charset='utf8mb4')

try:
    with connection.cursor() as cursor:
        # 讀取所有資料
        sql = "SELECT `id`, `password` FROM `users` WHERE `email`=%s"
        cursor.execute(sql, ('[email protected]',))
        result = cursor.fetchall()
        print(result)
finally:
    connection.close()

這種寫法會將查詢到的所有資料寫入記憶體中,若在結果較大的情況下,會對記憶體造成很大的壓力,所幸 pymysql 實現了一種 SSCursor 遊標類,它允許將查詢結果按需返回,而不是一次性全部返回導致記憶體使用量飆升。

SSCursor

官方文件的解釋為:

Unbuffered Cursor, mainly useful for queries that return a lot of data,
or for connections to remote servers over a slow network.

Instead of copying every row of data into a buffer, this will fetch
rows as needed. The upside of this is the client uses much less memory,
and rows are returned much faster when traveling over a slow network
or if the result set is very big.

There are limitations, though. The MySQL protocol doesn't support
returning the total number of rows, so the only way to tell how many rows
there are is to iterate over every row returned. Also, it currently isn't
possible to scroll backwards, as only the current row is held in memory.

大致翻譯為:無快取的遊標,主要用於查詢大量結果集或網路連線較慢的情況。不同於普通的遊標類將每一行資料寫入快取的操作,該遊標類會按需讀取資料,這樣的好處是客戶端消耗的記憶體較小,而在網路連線較慢或結果集較大的情況下,資料的返回也會更快。當然,缺點就是它不支援返回結果的行數(也就是呼叫 rowcount

屬性將不會得到正確的結果,一共有多少行資料則需要全部迭代完成才能知道),當然它也不支援往回讀取資料(這也很好理解,畢竟是生成器嘛)。

它的寫法如下:

from pymysql.cursors import SSCursor

connection = pymysql.connect(host='localhost',
                             user='user',
                             password='passwd',
                             db='db',
                             charset='utf8mb4')

# 建立遊標
cur = connection.cursor(SScursor)
cur.execute('SELECT * FROM test_table')

# 讀取資料
# 此時的 cur 對記憶體消耗相對 Cursor 類來說簡直微不足道
for data in cur:
    print(data)

本質上對所有遊標類的迭代都是在不斷的呼叫 fetchone 方法,不同的是 SSCursor 對 fetchone 方法的實現不同罷了。這一點檢視原始碼即可發現:
Cursor 類 fetchone 方法原始碼(可見它是在根據下標獲取列表中的某條資料):

SSCursor 類 fetchone 方法原始碼(讀取資料並不做快取):

跳坑

當然,如果沒有坑就沒必要為此寫一篇文章了,開開心心的用著不香嗎。經過多次使用,發現在使用 SSCursor 遊標類(以及其子類 SSDictCursor)時,需要特別注意以下兩個問題:

1. 讀取資料間隔問題

每條資料間的讀取間隔若超過 60s,可能會造成異常,這是由於 MySQL 的 NET_WRITE_TIMEOUT 設定引起的錯誤(該設定值預設為 60),如果讀取的資料有處理時間較長的情況,那麼則需要考慮更改 MySQL 的相關設定了。(tips: 使用 sql SET NET_WRITE_TIMEOUT = xx 更改該設定或修改 MySQL配置檔案)

2. 讀取資料時對資料庫的其它操作行為

因為 SSCursor 是沒有快取的,只要結果集沒有被讀取完成,就不能使用該遊標繫結的連線進行其它資料庫操作(包括生成新的遊標物件),如果需要做其它操作,應該使用新的連線。比如:

from pymysql.cursors import SSCursor

def connect():
    connection = pymysql.connect(host='localhost',
                                 user='user',
                                 password='passwd',
                                 db='db',
                                 charset='utf8mb4')
    return connection

conn1 = connect()
conn2 = connect()

cur1 = conn1.cursor(SScursor)
cur2 = conn1.cursor()

with conn1.cursor(SSCursor) as ss_cur, conn2.cursor() as cur:
	try:
        ss_cur.execute('SELECT id, name FROM test_table')

        for data in ss_cur:
            # 使用 conn2 的遊標更新資料
            if data[0] == 15:
                cur.execute('UPDATE tset_table SET name="kingron" WHERE id=%s', args=[data[0])

            print(data)
    finally:
    	conn1.close()
    	conn2.close()

參考

  1. Cursor Objects — PyMySQL 0.7.2 documentation
  2. Using SSCursor (streaming cursor) to solve Python using pymysql to query large amounts of data leads to memory usage is too high