unnest を使って楽に bulk insert する

プライベートでよく PostgreSQL を使っている人の MySQL 不便シリーズ。

ActiveRecord のような高機能な ORM あるいはクエリビルダを使っていると bulk insert するクエリをライブラリがいいかんじに組み立ててくれるが、そんな高級なライブラリを使わずにもっとシンプルなもので済ませたいことがある。 そんなときに bulk insert をしたくなった場合、クエリを組み立てて変数を bind して…… というのを自分でやるのは結構面倒である。

PostgreSQL だと配列型があり unnest() で配列を行に変換できるので、これと insert into select を組み合わせると bulk insert するクエリを簡単に組み立てられる。

eagletmt=> select unnest(array[1,2,3]) as x, unnest(array[4,5,6]) as y;
 x | y
---+---
 1 | 4
 2 | 5
 3 | 6
(3 rows)

PostgreSQL だとこんなかんじで行に展開できるので、あとは insert into bulk_inserts (x, y) unnest($1), unnest($2) みたいなクエリに配列型の値を bind すれば bulk insert になる。たとえば Rust の sqlx だとこれだけ。

#[derive(Debug)]
struct Record {
    x: i32,
    y: i32,
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let records: Vec<_> = (0..1000)
        .map(|i| Record {
            x: 2 * i,
            y: 2 * i + 1,
        })
        .collect();

    let pool = sqlx::PgPool::connect("postgres://...").await?;

    let mut xs = Vec::with_capacity(records.len());
    let mut ys = Vec::with_capacity(records.len());
    for r in records {
        xs.push(r.x);
        ys.push(r.y);
    }
    sqlx::query("insert into bulk_inserts (x, y) select unnest($1), unnest($2)")
        .bind(&xs)
        .bind(&ys)
        .execute(&pool)
        .await?;

    Ok(())
}

この方法だと動的に SQL を組み立てる必要がなくて楽。なお大量のレコードをロードするなら copy を使ったほうがいいと思う。