1分鐘插入10億行數(shù)據(jù)!拋棄Python,寫腳本請使用Rust
最近,一位程序員表示自己急需一個「也就」十億行數(shù)據(jù)的測試數(shù)據(jù)庫,并且還得在一分鐘之內生成。
于是,他做了一個所有程序員都會做的事:寫一個Python腳本來生成數(shù)據(jù)庫。
然而,很不幸的是,這個腳本非 常慢。
于是,他又做了一個所有程序員都會做的事:進一步學習關于SQLite、Python以及不知道為什么還有Rust的知識。
項目已開源:https://github.com/avinassh/fast-sqlite3-inserts
目標
作者需要在他2019年的MacBook Pro(2.4GHz四核i5)上,一分鐘內生成一個有10億行的SQLite數(shù)據(jù)庫。

表的模式
要求:
- 生成的數(shù)據(jù)是隨機的;
- 「area」列將包含六位數(shù)的地區(qū)代碼(任何六位數(shù)都可以,不需要驗證);
- 「age」列是5、10或15中的任何一個;
- 「active」列是0或1。
不過,作者表示,對腳本的要求也不用太高,還是可以妥協(xié)的:
- 如果進程崩潰,所有的數(shù)據(jù)都丟失也沒有問題,再次運行腳本就可以了;
- 允許充分利用電腦的資源:100%的CPU,8GB的內存和剩余的SSD儲存;
- 不需要使用真正的隨機方法,來自stdlib的偽隨機方法就可以。
Python原型
在最開始的腳本中,作者試圖在一個for循環(huán)中逐一插入1000萬條記錄,而這讓用時直接達到了15分鐘。

顯然,這太慢了。
在SQLite中,每次插入都是一個事務,每個事務都保證它被寫入磁盤,作者推斷可能問題就來自這里。
于是作者開始嘗試不同規(guī)模的批量插入,發(fā)現(xiàn)10萬是一個最佳點,而運行時間減少到了10分鐘。

SQLite優(yōu)化
作者認為自己寫的代碼已經(jīng)很簡練了,并沒有什么可以優(yōu)化的空間。
于是他將下一個目標轉到了數(shù)據(jù)庫的優(yōu)化。
根據(jù)各種關于SQLite優(yōu)化的建議,作者做了一些改進。

- 關閉「journal_mode」將禁用回滾日志,也就是說,如果任何事務失敗,都無法回滾。
- 關閉「synchronous」,將使SQLite不再關心是否能可靠地寫入磁盤,而是把這個責任交給操作系統(tǒng)。也就是說,可能會出現(xiàn)SQLite并沒有成功寫入磁盤的情況。
- 「cache_size」指定了SQLite在內存中可以保留多少個內存頁。
- 當「locking_mode」為「EXCLUSIVE」模式時,SQLite鎖住的連接將永遠不會被釋放。
- 將「temp_store」設置為「MEMORY」可以讓其表現(xiàn)像一個內存數(shù)據(jù)庫。
此處作者提醒,請不要把這些操作用到生產(chǎn)上去。
重新審視Python
作者再次重寫了Python腳本,這次包括了微調的SQLite參數(shù),這次帶來了巨大的提升,運行時間大幅減少:
- 原始的for循環(huán)版本用時大約10分鐘。
- 批處理版本用時大約8.5分鐘。
PyPy
PyPy在其主頁上強調它比CPython快4倍,于是作者決定嘗試一下。
令作者有些意外的是,竟然不需要對現(xiàn)有的代碼進行任何改動,只需要在PyPy運行就可以了。
批處理版本只需要2.5分鐘,也就是速度快了接近3.5倍。
Busy Loop?
莫非是在Python的循環(huán)上耗費了太多時間?于是作者刪除了SQL指令之后再次跑了一遍代碼:
- 批處理版本在CPython中用時5.5分鐘。
- 批處理版本在PyPy中用時1.5分鐘(又是3.5倍的速度提升)。
然而用Rust重寫了相同的內容之后,循環(huán)只需要17秒。
于是,作者果斷拋棄Python,轉投Rust的懷抱。

Rust
像Python一樣,作者先寫了一個原始的Rust版本,一個循環(huán)執(zhí)行一行數(shù)據(jù)的插入。
然而,即便使用了所有SQLite的優(yōu)化,也依然消耗了大約3分鐘。于是作者進行了進一步的測試:
- 嘗試把「rusqlite」換成異步運行的「sqlx」,這讓用時直接被拉到了14分鐘。作者表示,這比自己迄今為止寫的任何一個Python迭代都要差。
- 在執(zhí)行原始SQL語句時,使用準備好的語句。這個版本的用時只有1分鐘。
最優(yōu)的版本
使用準備好的語句,以50行為一個批次插入,最終用時34.3秒。
作者又寫了一個線程版本,其中一個線程從通道接收數(shù)據(jù),還有四個線程向通道推送數(shù)據(jù)。
這個也是目前性能最好的版本,最終用時大約32.37秒。




IO時間
SQLite論壇上的網(wǎng)友提出了一個有趣的想法:測量內存數(shù)據(jù)庫所需的時間。
于是作者又跑了一遍代碼,將數(shù)據(jù)庫的位置設定為「:memory:」,rust版本完成的時間少了兩秒(29秒)。
也就是說將1億條記錄寫入到磁盤上需要2秒,這個用時似乎也是合理的。
這也說明,可能沒有更多的SQLite優(yōu)化可以以更快的方式寫入磁盤,因為99%的時間都花在生成和添加數(shù)據(jù)上。
排行榜
插入1億行數(shù)據(jù)的用時:
Rust33秒 PyPy126秒 CPython210秒
總結
- 盡可能使用SQLite PRAGMA語句
- 使用準備好的語句
- 進行分批插入
- PyPy確實比CPython快4倍
- 異步不一定更快
目前,第二快的版本是單線程運行的,而作者的電腦有4個核心,于是他在一分鐘內可以得到8億行數(shù)據(jù)。然后再經(jīng)過幾秒鐘的數(shù)據(jù)合并,時間仍然可以少于一分鐘。
網(wǎng)友評論
博主的這一番研究獲得了網(wǎng)友們的一致好評。

真的很喜歡這些觀點:
學習了更多關于PRAGMA語句。
PyPy的效率和靈活性可以通過即插即用的方式體現(xiàn)(將來一定會給它一個機會)。
文章的排版非常簡單,有適當?shù)脑创a鏈接。很有趣,很容易上手。
Rust高光時刻又來了!