Pythonは本当に必要か？Rustで作るデータパイプライン

はじめに

昨今のデータ基盤の流行りはELTだと思いますが、つまりT処理をBigQueryやSnowflakeといったDWHサービス上で実行してしまうケースが多く、ELはAPIやほかサービスからデータを引っこ抜いてDWHサービスに生データに近い形で登録しておきたいといったことがだいぶ多いのではないかと思っています。

そして最近はfivetran, troccoのような便利なSaaSもありその部分の実装を行う機会はだいぶ減ったとは思うのですが、それでもSaaSだと高い、オープンソースだと要件に合わないようなことも多くあり自前Pythonで実装だ！となることもまだまだあるのかと思います。

Pythonはデータ分野で超便利な言語で、チームのケイパビリティなどを考慮しても採用しやすいと思いますが、一方でしばしばシンプルな処理なのにメモリを食いつぶしツライ・・・みたいことが起こりがちだと思っています。

今回はデータパイプラインのEL処理をRustで実装することで、

どの程度性能面のメリットを得られる可能性があるのか？
実装はシンプルにできるのか？

というところに焦点を当てて検証してみようと思います。

※ Rustはこのために勉強した感じなので、色々とアレなところがあったらすみません(自衛)

パフォーマンス比較

まずシナリオとして以下を想定して実験してみます

ローカルにあるCSVをBigQueryにinsert
BQ Loadではなく、メモリ上でCSVを解釈し逐次insertする
- BQにロードするための前処理をすることが多いという想定(ロードした時刻の付与とか、UTF変換とか)

とりあえず少しまとまったデータが欲しかったので、雑にPythonに作らせます(ChatGPTに作ってもらいました)

データ生成用コード(Python)

 1import csv
 2import random
 3import time
 4
 5# Configuration
 6num_rows = 1_000_000
 7output_file = "large_data.csv"
 8
 9# Define sample data
10names = ["Alice", "Bob", "Charlie", "David", "Eve"]
11countries = ["USA", "Canada", "UK", "Australia", "Germany", "Japan"]
12
13# Generate CSV
14start_time = time.time()
15with open(output_file, mode="w", newline="") as file:
16    writer = csv.writer(file)
17    # Write header
18    writer.writerow(["id", "name", "age", "country", "signup_date"])
19
20    # Write rows
21    for i in range(1, num_rows + 1):
22        writer.writerow([
23            i,  # id
24            random.choice(names),  # name
25            random.randint(18, 99),  # age
26            random.choice(countries),  # country
27            f"2024-{random.randint(1, 12):02d}-{random.randint(1, 28):02d}",  # signup_date
28        ])
29
30end_time = time.time()
31print(f"Generated {num_rows} rows in {end_time - start_time:.2f} seconds. Output: {output_file}")

比較用のコードです。どちらも1000件ずつinsertしていくという作りにしてあります。

Python

 1import asyncio
 2import concurrent.futures
 3import time
 4
 5import pandas as pd
 6from google.cloud import bigquery
 7
 8
 9def insert_batch(batch, table_ref):
10    client.insert_rows_json(table_ref, batch)
11
12def insert_to_bigquery_python(csv_file, dataset_id, table_id):
13    table_ref = client.dataset(dataset_id).table(table_id)
14
15    start_time = time.time()
16
17    df = pd.read_csv(csv_file)
18
19    rows = df.to_dict(orient="records")
20
21    async def process_batches():
22        with concurrent.futures.ThreadPoolExecutor() as executor:
23            loop = asyncio.get_event_loop()
24            tasks = []
25            for i in range(0, len(rows), 1000):
26                batch = rows[i:i+1000]
27                tasks.append(loop.run_in_executor(executor, insert_batch, batch, table_ref))
28            await asyncio.gather(*tasks)
29
30    asyncio.run(process_batches())
31
32    print(f"Python Insert Time: {time.time() - start_time} seconds")
33
34if __name__ == "__main__":
35    client = bigquery.Client()
36    insert_to_bigquery_python("large_data.csv", "DATASET_ID", "TABLE_ID")

Rust

 1use csv::Reader;
 2use serde_json::json;
 3use reqwest::Client;
 4use tokio::task;
 5use std::fs::File;
 6use std::io;
 7use std::time::Instant;
 8use std::env;
 9
10#[tokio::main]
11async fn main() -> io::Result<()> {
12    let csv_file = "large_data.csv";
13    let dataset_id = "DATASET_ID";
14    let table_id = "TABLE_ID";
15    let bigquery_url = format!(
16        "https://bigquery.googleapis.com/bigquery/v2/projects/{project_id}/datasets/{dataset_id}/tables/{table_id}/insertAll",
17        project_id = "GCP_PROJECT_ID",
18        dataset_id = dataset_id,
19        table_id = table_id
20    );
21
22    let client = Client::new();
23    let mut rdr = Reader::from_reader(File::open(csv_file)?);
24
25    let headers = rdr.headers()?.clone(); // Clone headers to avoid borrowing `rdr`
26    let mut rows = Vec::new();
27    let start_time = Instant::now();
28
29    let mut tasks = vec![];
30
31    for result in rdr.records() {
32        let record = result?;
33        let mut json_record = json!({});
34        for (header, value) in headers.iter().zip(record.iter()) {
35            json_record[header] = json!(value);
36        }
37        rows.push(json_record);
38
39        if rows.len() >= 1000 {
40            let batch = rows.clone();
41            rows.clear();
42            let client = client.clone();
43            let url = bigquery_url.clone();
44
45            tasks.push(task::spawn(async move {
46                insert_to_bigquery(&client, &url, &batch).await;
47            }));
48        }
49    }
50
51    if !rows.is_empty() {
52        let batch = rows;
53        let client = client.clone();
54        let url = bigquery_url.clone();
55
56        tasks.push(task::spawn(async move {
57            insert_to_bigquery(&client, &url, &batch).await;
58        }));
59    }
60
61    for task in tasks {
62        task.await.unwrap();
63    }
64
65    println!("Rust Insert Time: {:?} seconds", start_time.elapsed().as_secs());
66    Ok(())
67}
68
69async fn insert_to_bigquery(client: &Client, url: &str, rows: &[serde_json::Value]) {
70    let access_token = env::var("GOOGLE_ACCESS_TOKEN").expect("GOOGLE_ACCESS_TOKEN environment variable is not set");
71
72    let payload = json!({ "rows": rows.iter().map(|r| json!({ "json": r })).collect::<Vec<_>>() });
73
74    let response = client.post(url)
75        .header("Authorization", format!("Bearer {}", access_token))
76        .json(&payload)
77        .send()
78        .await;
79
80    match response {
81        Ok(resp) if resp.status().is_success() => println!("Batch inserted successfully"),
82        Ok(resp) => eprintln!("Failed to insert batch: {:?}", resp.text().await),
83        Err(err) => eprintln!("Error sending request: {:?}", err),
84    }
85}

3回ほど実行してみました。環境はそれぞれDockerコンテナで固めて以下で計算しています。

実行時間はプログラムからの出力
メモリはdocker statsで取得し、実行期間中の最大値 - 開始時

以下が結果です。かなり大きな差がついたように感じます。

	Python実行時間(sec)	Pythonメモリ使用量最大値(MB)	Rust実行時間(sec)	Rustメモリ使用量最大値(MB)
1回目	57.13560914993286	379	8	221.79
2回目	58.565985918045044	379	7	211.259
3回目	58.839012145996094	379.5	7	204.959
平均	58.1802024	379.1666667	7.333333333	212.6693333

並列化しないときはほぼ変わらない結果(どちらも大体400秒程度)だったのですが、Rustでは非同期処理の効率の高さがかなり際立っていそうです。

実際のデータ量がもっと多かったり並列数が増えたりしてくるとこの差が顕著になってきそうな気もしますね。

実装のシンプルさ

見た目の話であれば慣れや好みの問題も多分に含んでしまうので、ここではRustで得られるメリットの1つである型安全性からもたらされるシンプルさについて書いてみようと思います。

私としてはPythonのduck typingも悪くないと思っているのですが、プロダクションに乗せるコードだともう少しいい感じにしたいと思うときも多くあります。

例えば、以下のコードは本当にageがint値であることの保証はどこにもなく、例えば入力がユーザの作ったスプレッドシートのようなものを想定すると変な値が入ることは十分にあり、このコードが実行可能な保証はありません。

1import pandas as pd
2
3df = pd.read_json("data.json")
4print(df["age"] + 1)  # エラーの可能性がある

これを回避しようとすると、一つずつ型のチェックをするコードを書いてOKの場合処理すると言った形で、非常に見慣れたとてつもなく冗長なコードになると思います。

これがrustだと以下のように書けます。 (動かす場合、JSON文字列にあるコメントは削除してください)

 1use serde::Deserialize;
 2use serde_json::Value;
 3
 4#[derive(Deserialize, Debug)]
 5struct Record {
 6    age: u32,
 7    name: String,
 8}
 9
10fn main() {
11    let json_data = r#"
12        [
13            {"age": 25, "name": "Alice"},
14            {"age": "invalid", "name": "Bob"}, // 壊れたデータ
15            {"age": 30, "name": "Charlie"}
16        ]
17    "#;
18
19    let records: Vec<Result<Record, serde_json::Error>> =
20        serde_json::from_str::<Vec<Value>>(json_data)
21            .unwrap()
22            .into_iter()
23            .map(|item| serde_json::from_value::<Record>(item))
24            .collect();
25
26    for record in records {
27        match record {
28            Ok(valid_record) => println!("Valid: {:?}", valid_record),
29            Err(e) => eprintln!("Invalid record: {}", e),
30        }
31    }
32}

まず入力されるデータは型としてしっかり定義し、読み取った結果が正常か異常かによって処理を簡単に分けられるので、実行時にクラッシュする可能性を大きく下げられます。また型として明確なので、何らかの集計処理などの演算を行いたい場合も型に基づいた処理が行われるので、コードの意味も自明になるのが良いですね。

まとめ

データ界隈のエンジニアの共通スキルセットがPythonと言えそうな現状を考えると、Rustの導入は慎重になるべきというのは前提としてあると思っています。何よりPythonはデータに関して優れたエコシステムがあるため、特に作りはじめの段階においてはそれを手放してまで初手Rustというのは中々選択しにくいでしょう。

ただ、Rustを採用することで高速かつ安全なものが作れそうというのは実験を通して十分わかりました。

そのため、

処理時間が伸びていてなんとか短縮したい
CloudRunのようなCPUやメモリ課金の従量課金部分を削減したい
データパイプラインの安定性が特に求められる

のような、ある程度規模が大きかったり、データ活用が進んでいってクリティカルな利用が増えたシーンであれば学習コストを割いてでも採用してみても良さそうに思います。

Rustの学習コストが高い部分ですが、Web上のリソースだと断片的になりがちだったり、公式の資料が基本的に英語で辛かったので、私はRustの学習にあたって以下の本を読みました

Rustプログラミング完全ガイド他言語との比較で違いが分かる！

よくまとまっており、基本を押さえるのには大変よかったです。

どこかで本番投入できる機会があったら是非やりたいなぁと思いながら、本記事は締めようと思います。。

はじめに#

パフォーマンス比較#

実装のシンプルさ#

まとめ#

はじめに

パフォーマンス比較

実装のシンプルさ

まとめ