Zenn
👏

うわ・・・私のdbt遅すぎ!?をRustで解決したい

2024/12/22に公開
Rust
BigQuery
dbt
idea

はじめに

気軽にタイトルをつけてみて「Rustでdbtを書くのじゃ」という誤解を招かないかすごく不安になりました。
残念ながらそんなに高尚なものではないポエムなので気軽な気持ちで読んでいただければ幸いです。

これは datatech-jp Advent Calendar 2024 の22日目の記事です。

解決したい課題

普段、アナリティクスエンジニアとしてデータ分析基盤の運用を行っています。
Data Vaultを活用したデータモデリングの検討や、社内外に提供するダッシュボードのリファクタリングや開発など、日々多岐にわたる業務を行っています。

現在の会社のデータ分析基盤は、私が入社する前から数世代にわたって運用されてきました。
最大で5世代あった基盤も、先日1世代のみを残して全て整理されました。
それでもなお、現在1,500以上のモデルが存在し、dbt run や dbt compile を実行する際には30～40秒程度かかってしまう状況です。

一見「30秒くらいなら待てばいい」と思われるかもしれませんが、この時間は中途半端で、Slackを確認していると2～3分経過してdbtが完了していることもしばしばあります。
さらに、dbtのコマンドは一日に何十回も実行するため、この微妙に非効率な時間が積み重なると意外と大きなロスになると感じています。
これを何とかしたい！（部分的にでも！）というのが今回解決したい課題です。

課題のスコープ

やること

これにフォーカスします。

やらないこと

  • BigQuery以外のDWHツールのサポート
  • jinjaのコンパイル

jinjaのコンパイルまではスコープに入れていないのでやれることは限定的ですが、それでもmartを作る際のイテレーション速度を上げることには寄与してくれるのではないかと思っています。

解決イメージ

select
    user_id,
    name,
    age
from {{ ref("dim_user") }}

このようなsqlファイルがあった時に、

select
    user_id,
    name,
    age
from `my-project.my_dataset.dim_user`

このようなクエリを0.1秒くらいで出してくれるイメージです。

解決手段

「開発用適当ツール」をRustで作る。としました
参考にさせてもらった記事は以下です。

https://zenn.dev/codemountains/articles/0d3831c10c46b8

why Rust

私は静的型付け言語を大学生の頃にやったC以来まともに書いてきていなかったので、ひとつは自分のツールとして持っておきたいと思っていました。
ここ数年Rustの本を買っては途中まで読み、というのを繰り返しており全然手になじんでいませんでした。
「次はCLIを作るのがいい」とよくいろんなところで見るけど「うーん、何か自分が欲しいCLIってなんだろう」みたいに思っていてなかなか作るタイミングが掴めずにいました。
そこでこれはチャンスだと思ってRustで作ることにしました。
「なんでGolangじゃないの？」への答えとしては「まぁ、なんとなくRustが書きたいから」程度の理由です。

成果物

できたもの

fdbt model_name

のように叩くとコンパイルされたクエリが表示されます。
それをmacだったらパイプでpbcopyに渡してBigQueryのコンソールにペタって貼れます。
fdbtのfは「fastとかfuzzyとか速いけどふわっとしたdbt」くらいのノリでつけています。

速度で言うと、いつも使っている会社のdbtプロジェクトだと 38秒 かかるモデルが 0.06秒でクエリの描画されます。
生産性600倍！(と言うわけではないのは最後に書きます)

処理の流れ

  1. ~/.dbt/profiles.ymlを見てprofeleの情報をハッシュマップで持っておく
  2. プロジェクトルートにある dbt_project.yml を読み込んでその情報を元にprofileのハッシュマップからBigQueryのプロジェクト名とデータセット名を取得する
  3. コンパイル対象のモデルの {{ ref("model_name") }} の部分からmodel_nameをテーブル名として、{project_name}.{dataset_name}.{table_name} 形式に置き換える
  4. クエリを出力する

以上！
なのでjinjaの処理は全くしていないことをご了承ください。

コード

開発用適当ツールとしてレポジトリの公開もしていない(する予定もない)ものかつ、Rust初心者が書いたものなのでお粗末なものですが、お粗末ゆえにmain.rsしかないのでサクッと動かせるかもと思い一応貼っておきます。
恥ずかしいのでトグルで隠しておきます。

main.rs
use regex::Regex;

use clap::Parser;
use dirs::home_dir;
use std::fs;

#[derive(Parser, Debug)]
#[command(version, about, long_about = None)]
struct Args {
    model_name: String,
    #[arg(long, default_value_t = String::from("dev"))]
    target: String,
}

fn enum_files(dir: &str, model_name: &str) -> Vec<String> {
    let mut files = vec![];
    let entries = std::fs::read_dir(dir).unwrap();
    for entry in entries {
        let entry = entry.unwrap();
        let path = entry.path();
        if path.is_dir() { 
                let mut subfiles = enum_files(&path.to_string_lossy(),model_name);
                files.append(&mut subfiles);
            continue;
        }
        let fname = entry.file_name().to_string_lossy().to_string();
        // TODO: .sqlが入っていても通るようにする
        if format!("{}.sql", &model_name) == fname {
            files.push(entry.path().to_string_lossy().to_string());
        }
    }
    files
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let args = Args::parse();

    // working directory直下のdbt_project.ymlを読み、profile名を取得する
    let dbt_project_content = fs::read_to_string("dbt_project.yml")?;
    let dbt_project: serde_yaml::Value = serde_yaml::from_str(&dbt_project_content).unwrap();
    let profile_name = &dbt_project["profile"];

    // ~/.dbt/profiles.ymlからdbt_project.ymlで指定されたprofile名のプロジェクト名とデータセット名を取得する
    let mut profiles_path = home_dir().ok_or("ホームディレクトリが見つかりません")?;
    profiles_path.push(".dbt/profiles.yml");
    let yaml_content = fs::read_to_string(profiles_path)?;
    let profiles: serde_yaml::Value = serde_yaml::from_str(&yaml_content).unwrap();
    let project_name = &profiles[profile_name]["outputs"][&args.target]["project"]
        .as_str()
        .ok_or("projectが見つかりませんでした")?;
    let dataset_name = &profiles[profile_name]["outputs"][&args.target]["dataset"]
        .as_str()
        .ok_or("datasetが見つかりませんでした")?;

    // TODO: modelsのところもdbt_project.ymlから取得できるようにする
    let files = enum_files("./models", &args.model_name);
    for file_name in files{
        let target_sql_content = fs::read_to_string(file_name)?;

        let re = Regex::new(r#"\{\{\s*ref\(['\"](.+?)['\"]\)\s*\}\}"#).unwrap();
        // ref() の中身を BigQuery のテーブル名に変換
        let processed_query = re.replace_all(&target_sql_content, |caps: &regex::Captures| {
            format!("`{}.{}.{}`", project_name, dataset_name, &caps[1])
        });

        println!("{}", processed_query);
    }
    Ok(())
}

最後に

jinjaのコンパイルを行っているわではないので使えるモデルは限定的です。
なので速度が600倍だから生産性も600倍！とはもちろんならず、むしろ「結構いろんなところでjinja使ってるな」というのが最近感じているところです。
でもこれを最初の一歩として、自分の仕事の生産性を上げるための開発ツールは今後も少しづつ太らせていきたいと思っています。
「最近dbt遅いな」と文句を言っていた時に、Rustで書かれたdbtの代わりになるようなものも調べていたのですが、dbtのエコシステムが強すぎて乗り換えるのは当分先かなと思っています。(特にうちはautomate_dvもすごく使っているわけで。)
それよりも今は静的型付け言語をちゃんと習得したい。今年ずっとSQL書いてたのでモダンなプログラミング言語触りたい。の気持ちが強いのでここに対してはあまりwhyを投げかけず趣味でやっていきたいと思っています。
あれ？なんのアドカレだっけ？dbtのアドカレじゃないしいいか。
最後まで読んでいただきありがとうございました。

Discussion