こちらは Fusic Advent Calendar 2019 - Qiita 6日目の記事です。
昨日は @tsukabo による 社内基幹システムのバージョンアップに挑んだ話 でした。
バージョンアップはエンジニアの宿命でもありますよね。私も直近、バージョンアップ案件が控えているので参考になります。
今回は、Railsの大容量のCSVインポート処理をRustで書いてみた件です。
ただただ、Rustを書いてみたい
Tech系のブログや記事を見てると、Rustを目にする機会が多くなってきました。書いてみたくなったので、CSVをインポートしてDBに挿入する処理をRustで書くというお題目を設定しました。
前提
Railsで、Articleというテーブルに対して、約2万行のCSVファイルを読み込んで挿入するというタスクがあります。
すでにPostgreslでDBがあり、Articleというテーブルも存在しています。
Article は以下のschemaです。
create_table "articles", force: :cascade do |t|
t.string "title" # タイトル
t.string "body" # 本文
end
今回はその処理をRustで書いてみます。
CSVインポートして、ただ表示する
まずはCSVインポートですが、Cargoのパッケージがあるのでそれを使います。
[dependencies]
csv = "1.1"
次に、CSVを読み込んで表示する処理です。（https://docs.rs/csv/ のチュートリアルの通り）
extern crate rust;
use self::rust::*;
use std::io;
use std::process;
fn read() -> Result<(), Box<dyn Error>> {
let mut rdr = csv::Reader::from_reader(io::stdin());
for result in rdr.records() {
let record = result?;
let title = &record[0];
let body = &record[1];
println!("title: {}", &title);
println!("body: {}", &body);
}
Ok(())
}
fn main() {
if let Err(err) = read() {
println!("error running read: {}", err);
process::exit(1);
}
}
実行
$ cargo run --bin show_csv < sample.csv
title: AAAAAA
body: BBBBB
Dieselを使ってDBへアクセスする
次に、DBへアクセスします。
RustからDBへアクセスする方法はいろいろありますが、今回は少しでもRailsに近い感じにするため、ORMの Diesel を使いました。
以下はしばらく、 https://diesel.rs/guides/getting-started/ に書かれている通りの設定が続きます。
[dependencies]
diesel = { version = "1.0.0", features = ["postgres"] }
データベースへの接続設定を環境変数にセットし、セットアップします。
$ echo DATABASE_URL=postgres://username:password@localhost/yama_sample > .env
$ diesel setup
使用するテーブルのmigrateを行います。
$ diesel migration generate create_articles
するとmigrationsフォルダ以下に、テーブル作成と削除用のSQLファイルが生成されるので、適宜設定します。
CREATE TABLE articles (
id SERIAL PRIMARY KEY,
title TEXT,
body TEXT
)
Modelを定義します。
use super::schema::articles;
#[derive(Queryable)]
pub struct Article {
pub id: i64,
pub title: String,
pub body: String,
}
チュートリアルでは、 id: i32 になっていましたが、Railsで作ったテーブルはBigIntだったため、 i64 へ変更しています。
Schemaを定義します。
table! {
articles (id) {
id -> BigInt,
title -> Text,
body -> Text,
}
}
こちらもidを IntegerからBigIntへ変更しています。
これで最低限の設定は完了です。
CSVで読み込んだデータをテーブルに挿入
いよいよ本丸です。
まず、Insert用の構造体を定義します。
#[derive(Insertable)]
#[table_name="articles"]
pub struct NewArticle<'a> {
pub title: &'a str,
pub body: &'a str,
}
次に、Insert用の関数を定義します。
pub fn create_article<'a>(conn: &PgConnection, title: &'a str, body: &'a str) -> Article {
use schema::articles;
let new_article = NewArticle {
title: title,
body: body,
};
diesel::insert_into(articles::table)
.values(&new_article)
.get_result(conn)
.expect("Error saving new article")
}
最後に、CSVを読み込み、そのデータをDBへInsertするロジックを追加します。
extern crate rust;
extern crate diesel;
use self::rust::*;
use std::error::Error;
use std::io;
use std::process;
fn create() -> Result<(), Box<dyn Error>> {
let connection = establish_connection();
let mut rdr = csv::Reader::from_reader(io::stdin());
for result in rdr.records() {
let record = result?;
let title = &record[0];
let body = &record[1];
let article = create_article(&connection, &title, &body);
}
Ok(())
}
fn main() {
if let Err(err) = create() {
println!("error running create: {}", err);
process::exit(1);
}
}
実行すると、CSVのデータが1行ずつテーブルにInsertされます。
Rustに任せるのも一手
詳細に検証すると数値がずれるとは思いますが、速報値として、
- 同じ処理をRailsのままでやる: 2分7秒
- Rust: 30秒
という結果でした。もちろん、Railsの方でbulk insertを使うなどしてチューニングすることはできますし、単純に比較するものでもないと思います。
今回は、Rustでやったらどうなるかという好奇心から実装しましたが、Dieselという強力なORMもあることですし、この処理をRustに任せるというのは選択肢としてありうるのかなと思います。
次は @kawano-fusic の「S3×Lambda×Cloudwatch Eventsで、バッチ処理の監視機構を簡単に導入する」です。
はりきってどうぞ！