dbtの深淵『Custom Materialization』へようこそ

この記事はdbt Advent Calendar 2022の15日目です。
こんにちは、SREチーム所属の@mashiikeです。

皆さま、dbt活用していますか？便利ですよね。
今日はdbtの深淵の一つCustom Materializationについて話したいと思います。

dbtのmaterialization とは

dbtの便利な点、その本質の一つは.sqlや.pyにデータの取得方法さえ書いてしまえば、configに指定したmaterializedの文字列だけで面倒な更新処理をdbt側で全部やってくれる点だと思います。
すごく簡単なdbtのモデルを一つ用意してみましょう。

one.sql

{{
    config(
        materialized='table',
    )
}}
select 1 as num

このように書くと、1行1列で1という数字が入った物理テーブルが作成されます。

{{
    config(
        materialized='view',
    )
}}
select 1 as num

materialized='view' という記述に変えれば、今度はDB Viewとして作成されます。

{{
    config(
        materialized='incremental',
    )
}}
select 1 as num

materialized='incremental' という記述に変えれば、実行のたびに1を追記していってくれます。
非常に便利ですよね。 configの指定方法を変えるだけで、どのように実体化するのかが変えられてしまう。更新処理の方法でさえ、簡単に変えられる。
これは、dbtが持っている materializationという概念が為せる技なのです。

dbtの標準パッケージでは4つの materializationが提供されています。 table, view, incremental, ephemeral これらの４つがあれば基本的には困りません。多くの方は、４つを使いこなすことにのみ集中するのでよいのですが、極稀にこれら４つのmaterialization以外の特殊なmaterializationが欲しくなるケースがあります。

例えば、

DWHに存在しているMaterializedViewと呼ばれる機能を使って実体化したい。
forを駆使しすぎてBigQueryから HTTP Status Code 413 Request Entity Too Largeが帰ってくるのでTableを多段生成したい
Lambda Viewと呼ばれる概念を実現するいい感じの実体化がほしい
機械学習をSQLから実行するCREATE MODEL構文を駆使したいい感じ予測tableを作りたい
etc...

という感じです。
そんな特殊な状況に遭遇したときに覗き込むdbtの深淵が『Custom Matrialization』、、、新しい独自のmaterializationを作るという道です。

docs.getdbt.com

さて、ここで一つ注意を申し上げるならば、
多くの場合、えっ！？それって○○でできるじゃん！ と言われたり、なにか回避策あるんじゃないの？　とか、pre-hookとpost-hookを使えばなんとかなるんじゃないの？ などなどとなるため、
この深淵は知識として知ってるのはいいが実用的かと言われるとなかなか難しいでしょう。

利用に際しては、用法と用量を考えて控えめに使いましょう。

Custom Materialization 事始め

先程のドキュメント Creating new materializationsでは冒頭でこう書かれています。

The model materializations you're familiar with, table, view, and incremental are implemented as macros in a package that's distributed along with dbt. You can check out the source for these materializations here. If you need to create your own materializations, reading these files is a good place to start. Continue reading below for a deep-dive into dbt materializations.

要するに、dbtのデフォルトで提供されている、table, view, incremental, ephemeral のコードを読め！　ということですね。

github.com

おそらく、読んで分かる人にはこの記事は必要ないでしょう。ですので、簡単な例と応用的な例を踏まえて解説していきたいと思います。
まず手始めに、materializationのやっていることを列挙しましょう。

新しいモデルを構築するためにデータベースの準備する
pre-hookを実行する
実際の更新処理に必要なSQLを実行する。
post-hookを実行する
必要に応じて、バックアップや中間テーブルなどをクリーンアップする。
dbtが保持しているRelationキャッシュを更新する

さて、実用も何も考えないのであるならば、ぶっちゃけると1,3,6を行うだけでmaterializationとしては最小限動きます。 1行1列で1という数字が入っているテーブルを作るだけのmaterialization one を作ってみましょう。

macors/materializations/one/one.sql

{%- materialization one, default %}
    {%- set identifier = model['alias'] -%}
    {%- set target_relation = api.Relation.create(identifier=identifier,
                                                schema=schema,
                                                database=database,
                                                type='table') -%}
    -- 1. 古いテーブルを削除する
    {{ adapter.drop_relation(target_relation) }}

    -- 3. 実際に更新する
    {% call statement('main', auto_begin=True) -%}
        {{ get_create_table_as_sql(False, target_relation, 'SELECT 1 as num') }}
    {%- endcall %}
    {{ adapter.commit() }}

    -- 6. dbtが保持しているRelationキャッシュを更新する
    {{ return({'relations': [target_relation]}) }}
{%- endmaterialization %}

models/one.sql

{{
    config(
        materialized='one',
    )
}}

　ものすごい暴挙をしています！おそらく色々突っ込みどころはあるとは思いますが、dbt 1.3.1ではコレで動いてしまいます。
このmaterialization oneを使用した models/one.sql では、何を書いても CREATE TABLE <table名> AS SELECT 1 as num が実行されます。
上記の materializationを理解するためには、いくつか説明が必要だと思います。

　まずはじめに api.Relation.create ですね。こちらの詳細はドキュメントを参照してください。
dbtにはいくつかのAPIが提供されていて、このapi.Relation.create の API はDWH上に存在しているRelation(テーブル、View、etc...)を表すためのオブジェクトを生成してくれます。
おそらく、皆様がよく使うとしたら、 incremental のmaterializationでよく使うであろう this がこのオブジェクトのインスタンスだったりします。
コレを使うことで、テーブル名等をいい感じにquoteしてくれるようになります。
このRelationのオブジェクトは、dbt側がある程度キャッシュしてくれています。
materializationの最後で {{ return({'relations': [target_relation]}) }} という形で、実際に更新されたRelationが何なのかを、materializationの実行元に返すことでそのキャッシュを更新されます。

　さて、次に気になるのは adapter.drop_relation(target_relation) や adapter.commit() でしょう。
こちらも詳細はドキュメントを参照するとしましょう。
OSS版のdbtを使っている皆様は最初に以下のようなpython モジュールを入れたと思います。

dbt-redshift, dbt-bigquery, dbt-snowflake, dbt-postgres, etc... これらがadpaterと呼ばれるもので、dbtのJinja関数としてadapter.**という共通インタフェースが提供されています。
要するに、実行対象のDWHやデータベースによって具体的なRelationのdrop方法やTransactionの取り扱い方が異なるので、このadapterという関数群がいい感じにそのへんを抽象化してくれているというわけですね。

　最後はcall statement の中にある get_create_table_as_sql macroでしょうか。
コレ自体は、dbtにbuiltinされている便利マクロです。実装はこちらです。
get_create_table_as_sql(temporary, relation, sql) というふうになっており、作成対象が第１引数が一時テーブルであるかどうか？第２引数が作成先のRelation、最後がSELECT構文のSQLとなります。
ここではSELECT 1 as numに決め打ちしていますが、第３引数には基本的に実際に使えるSQL・・・つまりmodelファイルの中の描画されたSQLを使いたくなると思います。 materializationのブロックの中では、 sql という変数がすでに用意されており、この中に描画済みのSQLが格納されています。

dbt 1.3 時点でのtable の materializationの実装では https://github.com/dbt-labs/dbt-core/blob/1.3.latest/core/dbt/include/global_project/macros/materializations/models/table/table.sql#L31

  -- build model
  {% call statement('main') -%}
    {{ get_create_table_as_sql(False, intermediate_relation, sql) }}
  {%- endcall %}

となっております。

以上が、非常にシンプルなmaterializationになります。
このmaterializationを見たあとで上記のbuiltinされているtable materialization の実装をみるとだいぶいろいろなことがされているということを感じるでしょう。

buildが失敗したときのために直接書き換えるのではなく、一旦中間テーブルとして作成する。
中間テーブルが作成できたら、古いRelationはbackupとしてrenameする。 (この際、relationがtableではなくviewである場合も考慮されている)
コメントの付与やインデックスの作成、権限の付与が成功したらtransactionをCommitする。
Commit後にbackupを削除する。
run_hooks macroによる configに設定されたmacroの実行

これらが、 materialization one ではやっていないが table ではやっている内容になります。

もう少し実用的にした、応用例 `multi_stage_table`

こちらは、ベタに書いてしまうと長いので予めカスタムパッケージにしておきました。

github.com

実装はこちらとなっております。

この multi_stage_tableのmaterializationは forを駆使しすぎてBigQueryから HTTP Status Code 413 Request Entity Too Largeが帰ってくるのでTableを多段生成したい というケースを想定したPoCな物となっております。

{{
    config(
        materialized='multi_stage_table',
        stages=[
            'webapp',
            'raw',
        ],
    )
}}

select
    '{{ current_stage }}' as stage
    ,{{ dbt.safe_cast(dbt.date_trunc('day', '"timestamp"'), api.Column.translate_type("date")) }} as ymd
    {%- if current_stage == 'webapp' %}
    ,case when platform <> 'Web' then 'App' else 'Web' end as platform
    {%- else %}
    ,platform
    {%- endif %}
    ,count(distinct user_id) as uu
from {{ ref('action_log') }}
group by 1,2,3

configに書かれたstages の要素の数だけ上記のmodelのSQLを描画します。各stageの名前は current_stage という変数に格納されます。
current_stage に従ってSQLの描画を変更でき、それぞれ一時的なテーブルとして保存したあとに UNION ALLで全部くっけて一つのテーブルにするという奇特な materializationに仕上がっております。
実際に使用するには、もっと色々と考える必要がありますが Custom Materializationの奥深さを伝える例としては良いものではないかと思っています。

この materialization のキモは、こちらのmulti_stage_table.sql#L23-L44 となっております。

    {%- set stages = config.get('stages') %}
    {%- if stages is not sequence %}
        {{ exceptions.raise_compiler_error("Invalid `stages`. must be sequence. Got:" ~ number) }}
    {%- endif %}
    {%- if (stages | length) == 0 %}
        {{ exceptions.raise_compiler_error("Invalid `stages`. stage must be greater than or equal to 1.") }}
    {%- endif %}

    {%- set stage_relations = {} %}
    {%- set stage_sqls = {} %}
    {%- for stage in stages %}
        {%- set stage_unrendered_sql = "{%- set current_stage = '"~stage~"' %}\n"~ model['raw_code'] %}
        {%- set stage_rendered_sql = render(stage_unrendered_sql) | trim %}
        {%- if (stage_rendered_sql | length) == 0 %}
      。i    {{ exceptions.raise_compiler_error("state `"~stage~"` sql is empty") }}
        {%- endif %}
        {%- do stage_sqls.update({stage: stage_renderd_sql}) %}

        {%- set stage_relation = make_temp_relation(target_relation) -%}
        {%- do stage_relations.update({stage: stage_relation}) %}
    {%- endfor %}

materialization中ではmodel というコンテキスト変数があり、この中には様々な情報が入っています。
configに指定したstagesの情報は {%- set stages = config.get('stages') %} の一文で取得し、このstagesをforで回しつつ各stageのSQLを実行しています。
model['raw_code']には描画前の生のテンプレートSQLが保存されていて、文字列結合で {%- set current_stage = '<stage名>' %} という文字列を先頭にくっつけて、buildinで提供されているrender関数によって再描画しています。
そして、後にこの再描画された各stageのSQLを一時テーブルとして保存したあとにUNION ALLで結合するSQLをこのmaterializationのmain処理として実行しています。 multi_stage_table.sql#L62-L71

このように、dbtのmaterializationという仕組みを理解すると、dbtをもっとカスタマイズして使い倒すことができるようになります。

おわりに

dbtの深淵の一つ『Custom Materialization』いかがでしょうか？おそらく、多くの人は『コレ実用的じゃないじゃん！』と思ったことでしょう。
私の場合、実運用では『RedshiftのMaterializedViewを使用する』『Redshift MLのCREATE MODEL構文を使った予測モデルを作る』くらいにしか有効的に使えてないと思っています。
気軽に使うものでは無いですが、いざというときの手段の一つとして覚えておくと面白いかと思います。
皆様のよきdbt活用の一助となれば幸いです。

カヤックでは、様々な深淵に興味があるエンジニアも募集しています。

hubspot.kayac.com