在SQLite中插入10億條Python VS Rust

作者：蟲蟲搜奇 2021-09-07 10:55:36

寫腳本來進行數據處理，比如說給數據庫導入導出數據，這種任務一般來說最方便的方法是用python腳本，但是如果數據量比較大時候(比如上億條)時候Python就會超級慢，看到無法忍受。

在實際生活中，市場有這樣的案例：寫腳本來進行數據處理，比如說給數據庫導入導出數據，這種任務一般來說最方便的方法是用python腳本，但是如果數據量比較大時候(比如上億條)時候Python就會超級慢，看到無法忍受。在這種案例時候該怎么做呢，有一個外國老哥分享了自己的實踐經歷，并且對比了Python和Rust語言給SQLite插入十一條數據的情況，最后用Rust實現了在一分鐘來完成任務。我們在此分享一下該實踐過程，希望能對大家有所啟迪，大家也可以嘗試自己最拿手方法來實現該例子，并對比一下具體性能。

概述

案例中的任務是SQLite數據庫插入10億條的數據。表(user)數據結構和約束如下：

create table IF NOT EXISTS user 
( 
id INTEGER not null primary key, 
area CHAR(6), 
age INTEGER not null, 
active INTEGER not null 
);

隨機生成數據。其中are列為六位數的區號(任何六位數字)。age將是5、10 或15中的一個數字。Active為0或1。

實驗環境硬件配置為：MacBook Pro，2019(2.4 GHz 四核i5，8GB內存，256GB SSD硬盤，Big Sur 11.1)。
任務前提：任務無需保持程序穩健性，如果進程崩潰并且所有數據都丟失了也沒關系。可以再次運行腳本。
需要充分利用我的機器資源：100% CPU、8GB 內存和千兆字節的SSD空間。

無需使用真正的隨機方法，stdlib偽隨機方法即可。

Python

首先是原始版本的Python方法。Python標準庫提供了一個SQLite模塊，首先使用它編寫了第一個版本。代碼如下：

import sqlite3 
from commons import get_random_age, get_random_active, get_random_bool, get_random_area_code, create_table 
DB_NAME = "naive.db" 
def faker(con: sqlite3.Connection, count=100_000): 
for _ in range(count): 
age = get_random_age() 
active = get_random_active() 
# switch for area code 
if get_random_bool(): 
# random 6 digit number 
area = get_random_area_code() 
con.execute('INSERT INTO user VALUES (NULL,?,?,?)', (area, age, active)) 
else: 
con.execute('INSERT INTO user VALUES (NULL,NULL,?,?)', (age, active)) 
con.commit() 
def main(): 
con = sqlite3.connect(DB_NAME, isolation_level=None) 
con.execute('PRAGMA journal_mode = WAL;') 
create_table(con) 
faker(con, count=10_000_000) 
if __name__ == '__main__': 
main()

在該腳本中，通for循環中一一插入1000萬條數據。執行花了將近15分鐘。基于此進行優化迭代，提高性能。

SQLite中，每次插入都是原子性的并且為一個事務。每個事務都需要保證寫入磁盤(涉及IO操作)，因此可能會很慢。為了優化，可以嘗試通過不同大小的批量插入，對比發現，100000是最佳選擇。通過這個簡單的更改，運行時間減少到了10分鐘，優化了3分之一，但是仍然非常耗時。優化后，批量插入版本源碼：

SQLite庫優化

除了在代碼層優化外，如果對于單純的數據寫入，對數據庫本身搞的優化也是非常重要的。對于SQLite優化，可以做如下配置：

PRAGMA journal_mode = OFF; 
PRAGMA synchronous = 0; 
PRAGMA cache_size = 1000000; 
PRAGMA locking_mode = EXCLUSIVE; 
PRAGMA temp_store = MEMORY;

具體解釋：

首先，journal_mode設置為OFF，將會關閉回滾日志，禁用 SQLite 的原子提交和回滾功能，這樣在事務失敗情況下，無法恢復，基于例子實例穩健性要求可以設置，但是嚴禁在生產環境中使用。

其次，關閉synchronous，SQLite可以不再校驗磁盤寫入的數據可靠性。寫入SQLite可能并不意味著它已刷新到磁盤。同樣，嚴禁在生產環境中啟用。

cache_size用戶指定SQLite允許在內存中保留多少內存頁。不要在生產中分配太高的的數值。

使用在EXCLUSIVE鎖定模式，SQLite連接持有的鎖永遠不會被釋放。

設置temp_store到MEMOR將使其表現得像一個內存數據庫。

優化性能

對上面的兩個腳本，添加 SQLite優化參數，然后重新運行：

def main():     
con = sqlite3.connect(DB_NAME, isolation_level=None)     
con.execute('PRAGMA journal_mode = OFF;')     
con.execute('PRAGMA synchronous = 0;')     
con.execute('PRAGMA cache_size = 1000000;') # give it a GB     
con.execute('PRAGMA locking_mode = EXCLUSIVE;')     
con.execute('PRAGMA temp_store = MEMORY;')     
create_table(con)

faker(con, count=100_000_000)

優化后版本，原始版本，插入1億行數據，大概花了10分鐘;對比批量插入版本大概花了8.5分鐘。

pypy版本

對比CPython PyPy在數據處理中可以提高性能，據說可以提高4倍以上的性能。本實驗中也嘗試編譯PyPy解釋器，運行腳本(代碼無需修改)。

使用pypy解釋器，批處理版本，插入1億行數據只需2.5分鐘。性能大概是Cpython的3.5倍，可見傳說的4倍性能提高確實是真的，誠不我欺也!。同時，為了測試在純循環插入中消耗的時間，在腳本中刪除SQL指令并運行：

以上腳本在CPython中耗時5.5分鐘。PyPy執行耗時1.5分鐘(同樣提高了3.5倍)。

Rust

在完成Python各種優化折騰。又嘗試了Rust版本的插入，對比也有個原始版本和批量插入版本。原始版本，也是每行插入：

use rusqlite::{params, Connection}; 
mod common; 
fn faker(mut conn: Connection, count: i64) { 
let tx = conn.transaction().unwrap(); 
for _ in 0..count { 
let with_area = common::get_random_bool(); 
let age = common::get_random_age(); 
let is_active = common::get_random_active(); 
if with_area { 
let area_code = common::get_random_area_code(); 
tx.execute( 
"INSERT INTO user VALUES (NULL, ?, ?, ?)", 
params![area_code, age, is_active], 
) 
.unwrap(); 
} else { 
tx.execute( 
"INSERT INTO user VALUES (NULL, NULL, ?, ?)", 
params![age, is_active], 
) 
.unwrap(); 
} 
} 
tx.commit().unwrap(); 
} 
fn main() { 
let conn = Connection::open("basic.db").unwrap(); 
conn.execute_batch( 
"PRAGMA journal_mode = OFF; 
PRAGMA synchronous = 0; 
PRAGMA cache_size = 1000000; 
PRAGMA locking_mode = EXCLUSIVE; 
PRAGMA temp_store = MEMORY;", 
) 
.expect("PRAGMA"); 
conn.execute( 
"CREATE TABLE IF NOT EXISTS user ( 
id INTEGER not null primary key, 
area CHAR(6), 
age INTEGER not null, 
active INTEGER not null)", 
[], 
) 
.unwrap(); 
faker(conn, 100_000_000) 
}

該版執行，大概用時3分鐘。然后我做了進一步的實驗：

將rusqlite，換成sqlx異步運行。

use std::str::FromStr;     
 
use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqliteSynchronous};     
use sqlx::{ConnectOptions, Connection, Executor, SqliteConnection, Statement};     
 
mod common;     
 
async fn faker(mut conn: SqliteConnection, count: i64) -> Result<(), sqlx::Error> {     
let mut tx = conn.begin().await?;     
let stmt_with_area = tx     
.prepare("INSERT INTO user VALUES (NULL, ?, ?, ?)")     
.await?;     
let stmt = tx     
.prepare("INSERT INTO user VALUES (NULL, NULL, ?, ?)")     
.await?;     
for _ in 0..count {     
let with_area = common::get_random_bool();     
let age = common::get_random_age();     
let is_active = common::get_random_active();     
if with_area {     
let area_code = common::get_random_area_code();     
stmt_with_area     
.query()     
.bind(area_code)     
.bind(age)     
.bind(is_active)     
.execute(&mut tx)     
.await?;     
} else {     
stmt.query()     
.bind(age)     
.bind(is_active)     
.execute(&mut tx)     
.await?;     
}     
}     
tx.commit().await?;     
Ok(())     
}     
 
#[tokio::main]     
async fn main() -> Result<(), sqlx::Error> {     
let mut conn = SqliteConnectOptions::from_str("basic_async.db")     
.unwrap()     
.create_if_missing(true)     
.journal_mode(SqliteJournalMode::Off)     
.synchronous(SqliteSynchronous::Off)     
.connect()     
.await?;     
conn.execute("PRAGMA cache_size = 1000000;").await?;     
conn.execute("PRAGMA locking_mode = EXCLUSIVE;").await?;     
conn.execute("PRAGMA temp_store = MEMORY;").await?;     
conn.execute(     
"CREATE TABLE IF NOT EXISTS user (     
id INTEGER not null primary key,     
area CHAR(6),     
age INTEGER not null,     
active INTEGER not null);",     
)     
.await?;     
faker(conn, 100_000_000).await?;     
Ok(())     
}

這個版本花了大約14分鐘。性能反而下降下降了。比Python版本還要差(原因值得深析)。

對執行的原始SQL語句，切換到準備好的語句并在循環中插入行，但重用了準備好的語句。該版本只用了大約一分鐘。

使用準備好的語句并將它們插入到50行的批次中，插入10億條，耗時34.3 秒。

use rusqlite::{Connection, ToSql, Transaction}; 
mod common; 
fn faker_wrapper(mut conn: Connection, count: i64) { 
let tx = conn.transaction().unwrap(); 
faker(&tx, count); 
tx.commit().unwrap(); 
} 
fn faker(tx: &Transaction, count: i64) { 
// that is, we will batch 50 inserts of rows at once 
let min_batch_size: i64 = 50; 
if count < min_batch_size { 
panic!("count cant be less than min batch size"); 
} 
// jeez, refactor this! 
let mut with_area_params = " (NULL, ?, ?, ?),".repeat(min_batch_size as usize); 
with_area_params.pop(); 
let with_area_paramswith_area_params = with_area_params.as_str(); 
let mut without_area_params = " (NULL, NULL, ?, ?),".repeat(min_batch_size as usize); 
without_area_params.pop(); 
let without_area_paramswithout_area_params = without_area_params.as_str(); 
let st1 = format!("INSERT INTO user VALUES {}", with_area_params); 
let st2 = format!("INSERT INTO user VALUES {}", without_area_params); 
let mut stmt_with_area = tx.prepare_cached(st1.as_str()).unwrap(); 
let mut stmt = tx.prepare_cached(st2.as_str()).unwrap(); 
for _ in 0..(count / min_batch_size) { 
let with_area = common::get_random_bool(); 
let age = common::get_random_age(); 
let is_active = common::get_random_active(); 
let mut param_values: Vec<_> = Vec::new(); 
if with_area { 
// lets prepare the batch 
let mut vector = Vec::<(String, i8, i8)>::new(); 
for _ in 0..min_batch_size { 
let area_code = common::get_random_area_code(); 
vector.push((area_code, age, is_active)); 
} 
for batch in vector.iter() { 
param_values.push(&batch.0 as &dyn ToSql); 
param_values.push(&batch.1 as &dyn ToSql); 
param_values.push(&batch.2 as &dyn ToSql); 
} 
stmt_with_area.execute(&*param_values).unwrap(); 
} else { 
// lets prepare the batch 
let mut vector = Vec::<(i8, i8)>::new(); 
for _ in 0..min_batch_size { 
vector.push((age, is_active)); 
} 
for batch in vector.iter() { 
param_values.push(&batch.0 as &dyn ToSql); 
param_values.push(&batch.1 as &dyn ToSql); 
} 
stmt.execute(&*param_values).unwrap(); 
} 
} 
} 
fn main() { 
let conn = Connection::open("basic_batched.db").unwrap(); 
conn.execute_batch( 
"PRAGMA journal_mode = OFF; 
PRAGMA synchronous = 0; 
PRAGMA cache_size = 1000000; 
PRAGMA locking_mode = EXCLUSIVE; 
PRAGMA temp_store = MEMORY;", 
) 
.expect("PRAGMA"); 
conn.execute( 
"CREATE TABLE IF NOT EXISTS user ( 
id INTEGER not null primary key, 
area CHAR(6), 
age INTEGER not null, 
active INTEGER not null)", 
[], 
) 
.unwrap(); 
faker_wrapper(conn, 100_000_000) 
} 
創建了一個線程版本，其中有一個從通道接收數據的寫入線程和四個將數據推送到管道其他線程。 
use rusqlite::{Connection, ToSql}; 
use std::sync::mpsc; 
use std::sync::mpsc::{Receiver, Sender}; 
use std::thread; 
mod common; 
static MIN_BATCH_SIZE: i64 = 50; 
enum ParamValues { 
WithArea(Vec<(String, i8, i8)>), 
WithoutArea(Vec<(i8, i8)>), 
} 
fn consumer(rx: Receiver<ParamValues>) { 
let mut conn = Connection::open("threaded_batched.db").unwrap(); 
conn.execute_batch( 
"PRAGMA journal_mode = OFF; 
PRAGMA synchronous = 0; 
PRAGMA cache_size = 1000000; 
PRAGMA locking_mode = EXCLUSIVE; 
PRAGMA temp_store = MEMORY;", 
) 
.expect("PRAGMA"); 
conn.execute( 
"CREATE TABLE IF NOT EXISTS user ( 
id INTEGER not null primary key, 
area CHAR(6), 
age INTEGER not null, 
active INTEGER not null)", 
[], 
) 
.unwrap(); 
let tx = conn.transaction().unwrap(); 
{ 
// jeez, refactor this! 
let mut with_area_params = " (NULL, ?, ?, ?),".repeat(MIN_BATCH_SIZE as usize); 
with_area_params.pop(); 
let with_area_paramswith_area_params = with_area_params.as_str(); 
let mut without_area_params = " (NULL, NULL, ?, ?),".repeat(MIN_BATCH_SIZE as usize); 
without_area_params.pop(); 
let without_area_paramswithout_area_params = without_area_params.as_str(); 
let st1 = format!("INSERT INTO user VALUES {}", with_area_params); 
let st2 = format!("INSERT INTO user VALUES {}", without_area_params); 
let mut stmt_with_area = tx.prepare_cached(st1.as_str()).unwrap(); 
let mut stmt_without_area = tx.prepare_cached(st2.as_str()).unwrap(); 
for param_values in rx { 
let mut row_values: Vec<&dyn ToSql> = Vec::new(); 
match param_values { 
ParamValues::WithArea(values) => { 
for batch in values.iter() { 
row_values.push(&batch.0 as &dyn ToSql); 
row_values.push(&batch.1 as &dyn ToSql); 
row_values.push(&batch.2 as &dyn ToSql); 
} 
stmt_with_area.execute(&*row_values).unwrap(); 
} 
ParamValues::WithoutArea(values) => { 
for batch in values.iter() { 
row_values.push(&batch.0 as &dyn ToSql); 
row_values.push(&batch.1 as &dyn ToSql); 
} 
stmt_without_area.execute(&*row_values).unwrap(); 
} 
} 
} 
} 
tx.commit().unwrap(); 
} 
fn producer(tx: Sender<ParamValues>, count: i64) { 
if count < MIN_BATCH_SIZE { 
panic!("count cant be less than min batch size"); 
} 
for _ in 0..(count / MIN_BATCH_SIZE) { 
let with_area = common::get_random_bool(); 
let age = common::get_random_age(); 
let is_active = common::get_random_active(); 
let mut param_values: Vec<_> = Vec::new(); 
if with_area { 
// lets prepare the batch 
let mut vector = Vec::<(String, i8, i8)>::new(); 
for _ in 0..MIN_BATCH_SIZE { 
let area_code = common::get_random_area_code(); 
vector.push((area_code, age, is_active)); 
} 
for batch in vector.iter() { 
param_values.push(&batch.0 as &dyn ToSql); 
param_values.push(&batch.1 as &dyn ToSql); 
param_values.push(&batch.2 as &dyn ToSql); 
} 
// send the values 
tx.send(ParamValues::WithArea(vector)).unwrap(); 
} else { 
// lets prepare the batch 
let mut vector = Vec::<(i8, i8)>::new(); 
for _ in 0..MIN_BATCH_SIZE { 
vector.push((age, is_active)); 
} 
for batch in vector.iter() { 
param_values.push(&batch.0 as &dyn ToSql); 
param_values.push(&batch.1 as &dyn ToSql); 
} 
// send the values 
tx.send(ParamValues::WithoutArea(vector)).unwrap(); 
} 
} 
} 
fn main() { 
// setup the DB and tables 
let (tx, rx): (Sender<ParamValues>, Receiver<ParamValues>) = mpsc::channel(); 
// lets launch the consumer 
let consumer_handle = thread::spawn(|| consumer(rx)); 
let cpu_count = num_cpus::get(); 
let total_rows = 100_000_000; 
let each_producer_count = (total_rows / cpu_count) as i64; 
let mut handles = Vec::with_capacity(cpu_count); 
for _ in 0..cpu_count { 
let thread_tx = tx.clone(); 
handles.push(thread::spawn(move || { 
producer(thread_tx, each_producer_count.clone()) 
})) 
} 
for t in handles { 
t.join().unwrap(); 
} 
drop(tx); 
// wait till consumer is exited 
consumer_handle.join().unwrap(); 
}

這是性能最好的版本，耗時約32.37秒。

基準測試對比：

總結

通過案例不同任務實驗，總體上可以得到：

通過SQLite PRAGMA語句優化設置可以提高插入性能。
使用準備好的語句可以提高性能
進行批量插入可以提高性能。
PyPy 實際上比CPython快4倍
線程/異步不一定能提高性能。

責任編輯：趙寧寧來源：蟲蟲搜奇