dbtで引き出すSnowflake MLの真の力
やること
この記事では、dbtを使ってSnowflake MLを最大限活用する方法をご紹介します。
単なるデモに留まらず、運用を見据えたEnd-to-EndのMLパイプライン構築を目指します。
使用したコードは全てGitHubのリポジトリにて公開しています。
使用するデータ
Snowflakeのデモデータとして提供されている、tasty bytesを使用します。
こちらのチュートリアルからコードを一部拝借し、Snowflakeにdbt projectの環境を作ります。
tasty bytesは世界各地でキッチンカーを運営するフランチャイズの本部です。
各キッチンカーごとのPOSデータが利用できます。
今回はキッチンカーが属する地域別の商品需要予測をすることにしました。
実行環境のセットアップ
今回は全てのコードを dbt projects on Snowflake (Workspace) で実行します。
まずはsetupに入っているコードを実行していきます。
初めてdbt projectを触る方は、公式のチュートリアルを見ながらやることをお勧めします。
チュートリアルを終えるとtasty bytesのデータがSnowflakeに取り込めているはずです。
運用を見据えたMLパイプラインの構築を目指す、と言っておきながら大変恐縮ですが、今回はACCOUNTADMINを使用して全てのコードを実行しています。
dbtの構成
全体のデータフローは以下のようになっています。
プロジェクトのディレクトリ構成はこちらです。
tasty_bytes_dbt_ml_demo/
├─ macros
│ ├─ generate_schema_name.sql
│ └─ materialization_snowflake__feature_view.sql
├─ models
│ ├─ features
│ ├─ ml_models
│ └─ staging
├─ setup
│ └─ setup.sql
├─ dbt_project.yml
├─ packages.yml
└─ profiles.yml
簡単にmodels配下を説明しておきます。
stagingは特徴量作成のために、初期の加工を行うスキーマです。
内部でraw_モデルとcurated_モデルに分かれており、curated_モデルでデータが実体化されます。
featuresはSnowflakeの特徴量ストアとなるスキーマです。
Snowflakeの特徴量ストアはDynamic TableかViewかどちらかを選択できますが、本記事ではViewを選択しています。
また詳細は後述しますが、特徴量を作成するためのエンコーダーやスケーラーはfeaturesのモデルレジストリ内に保存します。
そして最後はml_modelsです。ここはMLモデルや予測結果を保存するスキーマです。
MLモデル学習時はtrainタグを使います
dbt build --select tag:train --traget dev
推論を実行したい時はpredictです。
dbt build --select tag:predict --traget dev
それでは詳細の解説に入ります。
RAW
このスキーマはtasty bytesの生のデータが入っています。
いわゆるBronzeレイヤーです。
特筆することは無いのですが、Snowflake公式のtasty bytesデモでは、このレイヤーのテストでこけるようになっているようです。
私のリポジトリでは、こけないように微修正しています。
STAGING
ここは特徴量開発に向けてデータを整形するスキーマです。
いわゆるSilverレイヤーです。
多くの場合でスタースキーマやDataVault2.0など、何かしらの体系立ったデータ加工が実施されます。
今回はMLの機能にスポットを当てたいので、大したことはしていません。
このスキーマ内はraw_とcurated_という接頭語がついたモデルがありますが、raw_はビューで実装しており、curated_はテーブルで実体化しています。実運用ではincrementalとなる想定です。
raw_は列の名寄せや日付のフォーマット統一、データ型の変換などを実施しておきます。
curated_では特徴量作成以外のビジネスロジックを実装します。例えば欠損値埋め、外れ値除去などです。その他、学習データから外さなければならないデータの除去、集計、他データとの結合までを済ませておきます。
curated_モデルから特徴量を作ることになります。
ここを大福帳のような巨大なテーブルにしてしまうと、特徴量の定義が変わる度に再計算が必要になる可能性があるため、効率が悪いです。一方で正規化に拘りすぎても、特徴量ストア内での結合が複雑となり可読性が下がります。
扱うデータに合わせた程良いモデリングが求められるレイヤーです。
今回はキッチンカー別のデータを日付ごと、地域ごとに集計する作業があります。
FEATURES
さて、ここからが本題です。
Snowflakeには特徴量ストアという機能があり、学習データ・推論データで特徴量の定義を一元管理することができます。
特徴量ストアはエンティティという仕組みで特徴量をグルーピングしており、その中にFeature Viewという形でデータを保存します。これらはSnowflakeの一般的なオブジェクトと対応しています。
| 特徴量ストア | スキーマ |
| エンティティ | タグ |
| Feature View | ビュー・動的テーブル |
特徴量ストアについては、こちらの記事で詳しく解説しています。
dbtを使ってFeature Viewを作成するには、事前にスキーマを特徴量ストアとして初期化することと、エンティティを作成すること、が必要です。
今回はこれらをまとめて作成できるcustom fearurizationを作りました。
特徴量ストアの初期化からFeature Viewの作成までをストアドプロシージャで実行します。使ったプロシージャは同じスキーマに保存されるようにしています。
マクロは長いのでアコーディオンにしています
dbtモデルはSnowpark dataframe又はsqlテキストで実装できます。
このモデルは販売数量のLAG特徴量を作成しています。
一つのモデルで必要なエンティティの作成まで済みます。
ここでさらにひと工夫します。
特徴量の中にはStandardScalerやLabelEncoderを使って変換をしないといけないものがあります。
これらのツールを使う時は、MLモデル学習時にfitして、推論時には同じpythonオブジェクトでtransform
します。推論時に別のオブジェクトでfit→transformしてしまうと、学習時と異なる特徴量が生まれてしまう為、やってはいけません。
したがってFeature Viewは学習時のみ更新し、推論時には動かないようにする必要があります。
dbtは各モデルにタグを付加することができ、実行時に対象のタグがついたモデルのみ実行することが可能です。特徴量ストアのモデルは基本的にこの機能を使って、学習時のみ動くように制御します。
また、スケーラーなどのオブジェクトの保管にはSnowflakeのモデルレジストリを活用することができます。
モデルレジストリに保管したモデルはSQLで推論(transform)が実行できます。
そのお陰で、簡単にビューの定義に組み込むことができます。
Feature ViewはSnowsightで確認できます。
スケーラーの作成、保存そして特徴量定義の統一、これらが1つのdbtモデルで完結するのはとてもありがたいことです。
Snowflake MLの強みの1つと言えるでしょう。
ML_MODELS
特徴量の定義が終わったので、MLモデルを作ります。
ここで朗報ですが、株式会社estieさんのリポジトリでMLモデルをモデルレジストリに登録するcustom materializationが公開されています。
こちらを拝借し、dbt projects on Snowflakeで動くように微修正しました。
このmaterializationを使うと、MLモデルをレジストリに保存できます。
学習時のみ動けば良いので、trainタグを付けておきます。
特徴量ストアから学習データを作りますが、ここではSnowflake MLのAPIを使うためdbt上のリネージュが繋がらなくなります。少々不格好ではありますが、dbt.ref()を実行して参照するFeature Viewとのリネージュを繋いでおきます。
これはあくまでdbtのリネージュの話であり、Snowflake内部のリネージュはこんなことをしなくても繋がっているのでご安心ください。
学習データはFEATURESスキーマ内のFeature Viewから取得しますが、generate_training_set()を使うとエンティティに指定したjoin_keysに基いて適切にJoinしてくれます。時系列データの場合はASOF JOINが実行されます。
あとはオーソドックスにxgboostのモデルを書いていきます。
作成したモデルはSnowsightで確認できます。
続いては推論時に稼働するモデルです。
推論時も特徴量ストアからデータを取る為、Pythonモデルを使います。
せっかく予測値を特徴量とセットで保存したのでSHAP値も計算し、説明可能性の向上を図ります。
モデルレジストリに保存したモデルのうち、対象のモデルはexplainメソッドが付与されます。
これを使ってSHAP値を計算し、テーブルに保管しておきましょう。
Snowflake MLのAPIで描画することもできます。
感想
dbtを使ってSnowflake MLをフル活用したパイプラインを構築してみました。
Snowflake MLはAPIがとてもシンプルに作られており、SQLとも親和性が良いためdbtを使ったパイプライン構築が可能です。
すでにSnowflake x dbtでデータパイプラインを作っている場合には、他のDWHにデータを持ち出すことなく、同じアーキテクチャの延長線上でML環境を作ることができます。
Snowflake上で全てが完結するというシンプルさを持ちつつ、申し分ない機能を発揮するSnowflake ML、皆さんもぜひ使ってみてください✨
Discussion