DuckDBを使ってS3とTiDBの間でデータを移動する

## はじめに

S3を使っていると、parquet形式のファイルをしばしば使います（Aurora Snapshotとか)。このデータをTiDBに入れたり、逆にTiDBのデータをS3にエクスポートしたいというニーズはまぁある（個人の感想です）のですが、そのためにETLジョブ組むのも面倒ですよね？

[DuckDB](https://duckdb.org/)はまさにそういう作業にうってつけです。DuckDBはparquetファイルの扱いにとても優れており、更にs3を簡単に、個人的にはglueやAthenaより遥かに容易に扱えます。
更に、MySQLにも接続でき、それらの間のデータ移動もできます。

本記事ではDuckDBを使ってS3、TiDBそれぞれに接続してみて、更にTiDBからS3、S3からTiDB間のデータの受け渡しを行います。環境はMacOSです。

## DuckDBのインストール

いくつか方法はあると思いますが、Homebrewでインストールしました。

```bash
> brew install duckdb
> duckdb --version
v1.1.3 19864453f7
```

## DuckDBからS3への接続

https://duckdb.org/docs/extensions/httpfs/s3api

DuckDBはS3を標準でサポートしており、認証方式も非常に多彩です。特に、SSO(AWS IAM Identity Center)をサポートしているのが便利です。ここではSSOを使った認証を行っていきます。

1. はじめにBucketを作って適当なparquetのファイルを置きます。この記事では [TimeStored](https://www.timestored.com/data/sample/parquet) の`titanic.parquet`を使いました。

2. 次にaws cliを使ってssoにログインしておきます。
```bash
> aws sso login 
 # [ブラウザでログインして許可しておく]
```

3. duckDBを起動して、s3のシークレットを作成します。
```sql
CREATE SECRET secret_s3 (
    TYPE S3,
    PROVIDER CREDENTIAL_CHAIN,
    CHAIN 'sso'
);
┌─────────┐
│ Success │
│ boolean │
├─────────┤
│ true    │
└─────────┘
```

4. SQLでクエリします。parquetの読み取りは`read_parquet`関数を使う必要がありますが、ファイルパスの末尾に`.parquet`がある場合は勝手に使ってくれます。すごく便利です。

```sql
select * from 's3://seki-duckdb/titanic.parquet';
┌─────────────┬──────────┬────────┬───────────────────────────┬─────────┬────────┬───────┬───────┬──────────────────┬─────────┬─────────────┬──────────┐
│ PassengerId │ Survived │ Pclass │           Name            │   Sex   │  Age   │ SibSp │ Parch │      Ticket      │  Fare   │    Cabin    │ Embarked │
│    int64    │  int64   │ int64  │          varchar          │ varchar │ double │ int64 │ int64 │     varchar      │ double  │   varchar   │ varchar  │
├─────────────┼──────────┼────────┼───────────────────────────┼─────────┼────────┼───────┼───────┼──────────────────┼─────────┼─────────────┼──────────┤
│           1 │        0 │      3 │ Braund, Mr. Owen Harris   │ male    │   22.0 │     1 │     0 │ A/5 21171        │    7.25 │             │ S        │
│           2 │        1 │      1 │ Cumings, Mrs. John Brad…  │ female  │   38.0 │     1 │     0 │ PC 17599         │ 71.2833 │ C85         │ C        │
│           3 │        1 │      3 │ Heikkinen, Miss. Laina    │ female  │   26.0 │     0 │     0 │ STON/O2. 3101282 │   7.925 │             │ S        │
│           4 │        1 │      1 │ Futrelle, Mrs. Jacques …  │ female  │   35.0 │     1 │     0 │ 113803           │    53.1 │ C123        │ S        │
│           5 │        0 │      3 │ Allen, Mr. William Henry  │ male    │   35.0 │     0 │     0 │ 373450           │    8.05 │             │ S        │
...

```

これだけです。めちゃくちゃ簡単じゃないですか？

## DuckDBからTiDB(MySQL)への接続

### TiDB Cloud Serverlessの準備
サーバレスクラスタと適当なデータを用意しましょう。

1. 適当なServerlessクラスタを作成します
2. ヘルプ（右下の?マーク) から `Guided Tour of SQL Editor` を選択して、githubのデータをimportします。数分かかると思います
![スクリーンショット 2024-12-04 10.46.49.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/12495/f6c4330c-dfcc-f0ed-a67b-93601d70703a.png)
3. インポートが完了したら閉じて構いません
4. 接続情報を取得します。Overview画面の右上から、`Connect` ボタンを押して接続画面を開きます。この情報をDuckDBから接続する際に利用します
![screenshot-20241204-115640.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/12495/7a38b0e5-b537-2e1d-14b4-51d6deba2681.png)

### DuckDBからTiDBへの接続

https://duckdb.org/docs/extensions/mysql.html

DuckDBのMySQL extentionを使います。といっても拡張機能をINSTALL, LOADで呼び出すだけで、別にインストールが必要なわけではありません。

拡張機能を呼び出し、先程の接続情報を元にシークレットを作成します。

```sql
INSTALL mysql;
LOAD mysql;

CREATE SECRET  (
    TYPE MYSQL,
    HOST 'gateway01.ap-northeast-1.prod.aws.tidbcloud.com',
    PORT 4000,
    DATABASE github_sample,
    USER '<Serverlessのユーザー名>',
    PASSWORD '<Serverlessのパスワード>',
    SSL_MODE 'verify_identity',
    SSL_CA '/etc/ssl/cert.pem'
);

```

MySQLのデータベースをアタッチして、DuckDBのデータベースのように利用できます。

```sql
ATTACH '' AS github_sample (TYPE MYSQL);
select * from github_sample.github_repos limit 1;
┌─────────┬──────────────────────┬──────────────────┬───┬──────────────────────┬──────────────────────┬─────────────────────┬─────────────────────┐
│ repo_id │      repo_name       │ primary_language │ … │     description      │      embedding       │     created_at      │     updated_at      │
│  int32  │       varchar        │     varchar      │   │       varchar        │       varchar        │      timestamp      │      timestamp      │
├─────────┼──────────────────────┼──────────────────┼───┼──────────────────────┼──────────────────────┼─────────────────────┼─────────────────────┤
│    1227 │ mislav/will_paginate │ Ruby             │ … │ Pagination library…  │ [-0.022141477,0.01…  │ 2008-02-25 20:21:40 │ 2024-06-10 09:07:56 │
├─────────┴──────────────────────┴──────────────────┴───┴──────────────────────┴──────────────────────┴─────────────────────┴─────────────────────┤
│ 1 rows                                                                                                                     11 columns (7 shown) │
└─────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┘

```

## S3(Parquet) ⇔ TiDB

さて以上でDuckDBでS3、TiDBそれぞれを扱えるようになりました。DuckDBがすごいのはここからで、これらのデータベース間でデータのやり取りができます。

たとえば、TiDBから読み出したデータをparquet形式でS3に保存するには`COPY`を使います。

```sql
-- TiDBからS3に保存
COPY github_sample.github_repos TO 's3://seki-duckdb/github_repos.parquet';
100% ▕████████████████████████████████████████████████████████████▏

-- 保存されたファイルをクエリ
select * from read_parquet('s3://seki-duckdb/github_repos.parquet') limit 1;

┌─────────┬──────────────────────┬──────────────────┬───┬──────────────────────┬──────────────────────┬─────────────────────┬─────────────────────┐
│ repo_id │      repo_name       │ primary_language │ … │     description      │      embedding       │     created_at      │     updated_at      │
│  int32  │       varchar        │     varchar      │   │       varchar        │       varchar        │      timestamp      │      timestamp      │
├─────────┼──────────────────────┼──────────────────┼───┼──────────────────────┼──────────────────────┼─────────────────────┼─────────────────────┤
│    1227 │ mislav/will_paginate │ Ruby             │ … │ Pagination library…  │ [-0.022141477,0.01…  │ 2008-02-25 20:21:40 │ 2024-06-10 09:07:56 │
├─────────┴──────────────────────┴──────────────────┴───┴──────────────────────┴──────────────────────┴─────────────────────┴─────────────────────┤
│ 1 rows                                                                                                                     11 columns (7 shown) │
└─────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┘

```

parquetファイルにはワイルドカードが使えるので、日次でエクスポートしてまとめて分析みたいなことも可能です。

```sql

select count(*) from read_parquet('s3://seki-duckdb/github_repos.parquet') ;
┌──────────────┐
│ count_star() │
│    int64     │
├──────────────┤
│         5607 │
└──────────────┘

-- 例として、同じデータを別名で保存してみる
COPY github_sample.github_repos TO 's3://seki-duckdb/github_repos2.parquet';
100% ▕████████████████████████████████████████████████████████████▏

-- ワイルドカードを使って複数のparquetファイルをまとめて分析する
select count(*) from read_parquet('s3://seki-duckdb/github_repos*.parquet');

┌──────────────┐
│ count_star() │
│    int64     │
├──────────────┤
│        11214 │
└──────────────┘
```

逆にparquetのデータをTiDBに保存してみましょう。parquetはテーブル定義を持っているので、`CREATE TABLE ... AS SELECT ...` を使って、テーブルを定義することなしにTiDB側にImportできるのです。

```sql
-- s3からTiDBにテーブルを作成する
CREATE TABLE github_sample.titanic as
    select * from read_parquet('s3://seki-duckdb/titanic.parquet');

-- 作成されたテーブルをクエリ
select * from github_sample.titanic limit 1;
┌─────────────┬──────────┬────────┬─────────────────────────┬─────────┬────────┬───────┬───────┬───────────┬────────┬─────────┬──────────┐
│ PassengerId │ Survived │ Pclass │          Name           │   Sex   │  Age   │ SibSp │ Parch │  Ticket   │  Fare  │  Cabin  │ Embarked │
│    int64    │  int64   │ int64  │         varchar         │ varchar │ double │ int64 │ int64 │  varchar  │ double │ varchar │ varchar  │
├─────────────┼──────────┼────────┼─────────────────────────┼─────────┼────────┼───────┼───────┼───────────┼────────┼─────────┼──────────┤
│           1 │        0 │      3 │ Braund, Mr. Owen Harris │ male    │   22.0 │     1 │     0 │ A/5 21171 │   7.25 │         │ S        │
└─────────────┴──────────┴────────┴─────────────────────────┴─────────┴────────┴───────┴───────┴───────────┴────────┴─────────┴──────────┘
```

## さいごに

DuckDBはparquetの他にもCSVやJSONにも対応し、またPostgreSQLやsqliteなど多くのDBにも対応しています。それらのデータベース間で今回の記事のようにデータを移動したり加工したりすることができます。非常に軽量ですが便利なデータベースなので活用してみてください！