Help us understand the problem. What is going on with this article?

Rustでcsvを読み込みDBにインポートする

こちらは 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のパッケージがあるのでそれを使います。

Cargo.toml
[dependencies]
csv = "1.1"

次に、CSVを読み込んで表示する処理です。(https://docs.rs/csv/ のチュートリアルの通り)

src/bin/show_csv.rs
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/ に書かれている通りの設定が続きます。

Cargo.toml
[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ファイルが生成されるので、適宜設定します。

migrations/2019-12-05-213319_create_articles/up.sql
CREATE TABLE articles (
  id SERIAL PRIMARY KEY,
  title TEXT,
  body TEXT
)

Modelを定義します。

src/models.rs
use super::schema::articles;

#[derive(Queryable)]
pub struct Article {
    pub id: i64,
    pub title: String,
    pub body: String,
}

チュートリアルでは、 id: i32 になっていましたが、Railsで作ったテーブルはBigIntだったため、 i64 へ変更しています。

Schemaを定義します。

src/schema.rs
table! {
    articles (id) {
        id -> BigInt,
        title -> Text,
        body -> Text,
    }
}

こちらもidを IntegerからBigIntへ変更しています。

これで最低限の設定は完了です。

CSVで読み込んだデータをテーブルに挿入

いよいよ本丸です。
まず、Insert用の構造体を定義します。

src/models.rs
#[derive(Insertable)]
#[table_name="articles"]
pub struct NewArticle<'a> {
    pub title: &'a str,
    pub body: &'a str,
}

次に、Insert用の関数を定義します。

src/lib.rs
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するロジックを追加します。

src/bin/update_article.rs
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で、バッチ処理の監視機構を簡単に導入する」です。
はりきってどうぞ!

Why not register and get more from Qiita?
  1. We will deliver articles that match you
    By following users and tags, you can catch up information on technical fields that you are interested in as a whole
  2. you can read useful information later efficiently
    By "stocking" the articles you like, you can search right away
Comments
No comments
Sign up for free and join this conversation.
If you already have a Qiita account
Why do not you register as a user and use Qiita more conveniently?
You need to log in to use this function. Qiita can be used more conveniently after logging in.
You seem to be reading articles frequently this month. Qiita can be used more conveniently after logging in.
  1. We will deliver articles that match you
    By following users and tags, you can catch up information on technical fields that you are interested in as a whole
  2. you can read useful information later efficiently
    By "stocking" the articles you like, you can search right away
ユーザーは見つかりませんでした