1分鐘插入10億行數(shù)據(jù)！拋棄Python，寫腳本請使用Rust
最近，一位程序員表示自己急需一個「也就」十億行數(shù)據(jù)的測試數(shù)據(jù)庫，并且還得在一分鐘之內(nèi)生成。
于是，他做了一個所有程序員都會做的事：寫一個Python腳本來生成數(shù)據(jù)庫。
然而，很不幸的是，這個腳本非 常慢。
于是，他又做了一個所有程序員都會做的事：進一步學習關(guān)于SQLite、Python以及不知道為什么還有Rust的知識。
項目已開源：https://github.com/avinassh/fast-sqlite3-inserts
目標
作者需要在他2019年的MacBook Pro（2.4GHz四核i5）上，一分鐘內(nèi)生成一個有10億行的SQLite數(shù)據(jù)庫。
表的模式
要求：
- 生成的數(shù)據(jù)是隨機的；
- 「area」列將包含六位數(shù)的地區(qū)代碼（任何六位數(shù)都可以，不需要驗證）；
- 「age」列是5、10或15中的任何一個；
- 「active」列是0或1。
不過，作者表示，對腳本的要求也不用太高，還是可以妥協(xié)的：
- 如果進程崩潰，所有的數(shù)據(jù)都丟失也沒有問題，再次運行腳本就可以了；
- 允許充分利用電腦的資源：100%的CPU，8GB的內(nèi)存和剩余的SSD儲存；
- 不需要使用真正的隨機方法，來自stdlib的偽隨機方法就可以。
Python原型
在最開始的腳本中，作者試圖在一個for循環(huán)中逐一插入1000萬條記錄，而這讓用時直接達到了15分鐘。
顯然，這太慢了。
在SQLite中，每次插入都是一個事務(wù)，每個事務(wù)都保證它被寫入磁盤，作者推斷可能問題就來自這里。
于是作者開始嘗試不同規(guī)模的批量插入，發(fā)現(xiàn)10萬是一個最佳點，而運行時間減少到了10分鐘。
SQLite優(yōu)化
作者認為自己寫的代碼已經(jīng)很簡練了，并沒有什么可以優(yōu)化的空間。
于是他將下一個目標轉(zhuǎn)到了數(shù)據(jù)庫的優(yōu)化。
根據(jù)各種關(guān)于SQLite優(yōu)化的建議，作者做了一些改進。
- 關(guān)閉「journal_mode」將禁用回滾日志，也就是說，如果任何事務(wù)失敗，都無法回滾。
- 關(guān)閉「synchronous」，將使SQLite不再關(guān)心是否能可靠地寫入磁盤，而是把這個責任交給操作系統(tǒng)。也就是說，可能會出現(xiàn)SQLite并沒有成功寫入磁盤的情況。
- 「cache_size」指定了SQLite在內(nèi)存中可以保留多少個內(nèi)存頁。
- 當「locking_mode」為「EXCLUSIVE」模式時，SQLite鎖住的連接將永遠不會被釋放。
- 將「temp_store」設(shè)置為「MEMORY」可以讓其表現(xiàn)像一個內(nèi)存數(shù)據(jù)庫。
此處作者提醒，請不要把這些操作用到生產(chǎn)上去。
重新審視Python
作者再次重寫了Python腳本，這次包括了微調(diào)的SQLite參數(shù)，這次帶來了巨大的提升，運行時間大幅減少：
- 原始的for循環(huán)版本用時大約10分鐘。
- 批處理版本用時大約8.5分鐘。
PyPy
PyPy在其主頁上強調(diào)它比CPython快4倍，于是作者決定嘗試一下。
令作者有些意外的是，竟然不需要對現(xiàn)有的代碼進行任何改動，只需要在PyPy運行就可以了。
批處理版本只需要2.5分鐘，也就是速度快了接近3.5倍。
Busy Loop？
莫非是在Python的循環(huán)上耗費了太多時間？于是作者刪除了SQL指令之后再次跑了一遍代碼：
- 批處理版本在CPython中用時5.5分鐘。
- 批處理版本在PyPy中用時1.5分鐘（又是3.5倍的速度提升）。
然而用Rust重寫了相同的內(nèi)容之后，循環(huán)只需要17秒。
于是，作者果斷拋棄Python，轉(zhuǎn)投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)友提出了一個有趣的想法：測量內(nèi)存數(shù)據(jù)庫所需的時間。
于是作者又跑了一遍代碼，將數(shù)據(jù)庫的位置設(shè)定為「:memory:」，rust版本完成的時間少了兩秒（29秒）。
也就是說將1億條記錄寫入到磁盤上需要2秒，這個用時似乎也是合理的。
這也說明，可能沒有更多的SQLite優(yōu)化可以以更快的方式寫入磁盤，因為99%的時間都花在生成和添加數(shù)據(jù)上。
排行榜
插入1億行數(shù)據(jù)的用時：
Rust33秒 PyPy126秒 CPython210秒
總結(jié)
- 盡可能使用SQLite PRAGMA語句
- 使用準備好的語句
- 進行分批插入
- PyPy確實比CPython快4倍
- 異步不一定更快
目前，第二快的版本是單線程運行的，而作者的電腦有4個核心，于是他在一分鐘內(nèi)可以得到8億行數(shù)據(jù)。然后再經(jīng)過幾秒鐘的數(shù)據(jù)合并，時間仍然可以少于一分鐘。
網(wǎng)友評論
博主的這一番研究獲得了網(wǎng)友們的一致好評。
真的很喜歡這些觀點：
學習了更多關(guān)于PRAGMA語句。
PyPy的效率和靈活性可以通過即插即用的方式體現(xiàn)（將來一定會給它一個機會）。
文章的排版非常簡單，有適當?shù)脑创a鏈接。很有趣，很容易上手。
Rust高光時刻又來了！
