[Rust] arrow-rsでCSVをParquetに変換してみる

こんにちは、福岡オフィスのyoshihitohです。

最近業務でGlue/SparkやAthenaといったビッグデータ関連のツールを使う機会が増えてきました。Spark関連の動画セッションや解説記事を見るとカラムナ形式が推奨されていることが多いのですが、カラム形式に馴染みがなくなぜ早くなるのかイメージが掴めませんでした。

とりあえずどういった形式なのか実データを確認するため、CSVからParquetに変換してみることにしました。データの内部構造と読み書き処理を見比べたかったので、今回は arrow-rs を使って変換してみました。

検証環境

• macOS : Big Sur 11.4
• Rust : 1.53.0

試すこと

以下のCSVデータをParquetに変換してみます。

data/input.csv

id,title
1,"Classmethod, Inc."
2,"Happy Anniversary!!"

arrow-rs?

arrow-rsはApache ArrowのRust向け実装です。

Apache Arrowはインメモリの列指向データ形式です。 ツールやフレームワーク間のデータ転送を効率化するために利用されています。 詳細については公式サイトやブログ投稿を参照してください。

やってみる

プロジェクト作成

まずプロジェクトを作成します

$ cargo new csv-to-parquet
    Created binary (application) `csv-to-parquet` package

続いてCargo.tomlに依存クレートを設定します。今回は arrow と parquet の2つです。

Cargo.toml

[dependencies]
arrow = "4.4"
parquet = "4.4"

スキーマ作成

CSVとParquetの入出力で使用するスキーマを定義します。

use std::sync::Arc;
use arrow::datatypes::{DataType, Field, Schema};

// NOTE: 以降の入出力で必要になるためArcで包む
let schema = Arc::new(Schema::new(vec![
        Field::new("id", DataType::UInt64, false),
        Field::new("title", DataType::Utf8, false),
    ]));

CSVを読み込む

arrowの csv モジュールを使って読み込みます。

use std::fs::File;
use arrow::csv;

let input = File::open("./data/input.csv")?;
let mut csv = csv::ReaderBuilder::new()
        .has_header(true)
        .with_quote(b'"')
        .with_schema(Arc::clone(&schema))
        .build(input)?;

Parquetを出力する

parquetの arrow モジュールを使って出力します。

let output = File::create("./data/output.parquet")?;
let mut writer = parquet::arrow::ArrowWriter::try_new(output, schema,None)?;
while letSome(batch) = csv.next() { // CSVのデータ(RecordBatch)を読み込む
    let batch = batch?;
    writer.write(&batch)?;
}
writer.close()?;

ここまでで入出力の実装は完了です。

以下、コードの全体像です。

src/main.rs

use std::fs::File;
use std::sync::Arc;

use arrow::csv;
use arrow::datatypes::{DataType, Field, Schema};
use parquet::arrow::ArrowWriter;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let schema = Arc::new(Schema::new(vec![
        Field::new("id", DataType::UInt64, false),
        Field::new("title", DataType::Utf8, false),
    ]));
    let input = File::open("./data/input.csv")?;
    let mut csv = csv::ReaderBuilder::new()
        .has_header(true)
        .with_quote(b'"')
        .with_schema(Arc::clone(&schema))
        .build(input)?;

    let output = File::create("./data/output.parquet")?;
    let mut writer = ArrowWriter::try_new(output, schema,None)?;
    while letSome(batch) = csv.next() {
        let batch = batch?;
        writer.write(&batch)?;
    }
    writer.close()?;

    Ok(())
}

動作確認

早速実行してみましょう。 ./data/output.parquet が出力されればOKです。

$ cargo run
   Compiling csv-to-parquet v0.1.0 (/Users/yoshihitoh/workspace/projects/blog/rust/csv-to-parquet)
    Finished dev [unoptimized + debuginfo] target(s) in 3.26s
     Running `target/debug/csv-to-parquet`

中身を確認してみましょう。中身の閲覧にはparquet-toolsを使います。ツールについてはこちらの記事を参照してください。

$ parquet-tools cat ./data/output.parquet
id = 1
title = Classmethod, Inc.

id = 2
title = Happy Anniversary!!

CSVと同じ内容が表示されました！ばっちり変換できてそうです！

終わりに

arrowとparquetクレートを利用することで手軽にCSVからParquetに変換できることがわかりました。今後はカラムナ形式の内部構造を確認しながら、なぜ分散処理システムと相性が良いのか確認していきたいと思います！