Databricksのai_query関数をDeltaLiveTables上で使って手軽かつ強力なデータパイプラインを作る

# 導入

Databricksの```ai_query```関数、Delta Live Tables(以下、DLT)上でも使えるようなことが[ドキュメント](https://docs.databricks.com/en/sql/language-manual/functions/ai_query.html#examples)に書かれていました。
簡単にLLMを使うデータパイプラインをDLTで作れるとかなり便利なので、試してみた記録です。

# 準備

今回は、日本語文書の感情分類(Negative/Positiveの分類)処理を```ai_query```で行い、その結果を保管するというシンプルなパイプラインを構築します。

データとして、以下のようなCSVファイルを適当に作成しました。

```text:sample_001.csv
id,comment,category
1,滋賀レイクスのB.LEAGUE PREMIERへの初年度参入が決定しました！,sports
2,日々の運動は体の健康を維持するために特に重要です。,health
3,日経平均株価が7月以来の４万円台となりました,economy
```

このファイルをUnity Catalog管理下の適当なボリュームに保管します。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/883dd679-18b4-f854-fa29-273b6de9db7c.png)

今回は```/Volumes/training/llm/raws/samples/dlt_query/```という場所に保管しています。

# DLTパイプライン処理の定義

ノートブックを作成し、DLT用のパイプライン処理を定義します。
```ai_query```関数を使う関係上、全てSQLで記述しました。

まず、CSVファイルからデータを取り込む処理。
ファイルが増分で増えていくという想定として、[ストリーミングテーブル](https://docs.databricks.com/ja/tables/streaming.html)にしています。

```sql
CREATE OR REFRESH STREAMING TABLE sample_comment_raw
COMMENT "ai_queryによる加工用サンプルデータ取込用ストリーミングテーブル" AS
SELECT
  *
FROM
  STREAM read_files(
    '/Volumes/training/llm/raws/samples/dlt_query/',
    format => 'csv',
    header => true,
    mode => 'FAILFAST'
  )
```

次に```ai_query```を使った感情分析処理を記述します。
内容はシンプルですね。

```sql
CREATE
OR REFRESH MATERIALIZED VIEW sentiment_added
TBLPROPERTIES('pipelines.channel' = 'PREVIEW')
COMMENT "感情分析結果の追加" AS
SELECT
  id,
  category,
  comment,
  ai_query(
    "<ai_queryで使うエンドポイント名>",
    concat(
      "次のテキストをポジティブかネガティブに分類してください。返答はPositiveかNegativeかだけで返してください。\n\n",
      comment
    )
  ) as sentiment
FROM
  live.sample_comment_raw;
```


# DLTパイプラインの作成

次にDelta Live Tablesのパイプラインを作成します。

メニューからDelta Live Tableを開いて、パイプラインを作成します。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/84e9ef97-519e-cecf-1935-159b876b1c15.png)

パイプライン設定は、名前は適当に設定して先ほど作成したノートブックを指定します。
クラスタはサーバレスを使用しました。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/a55180aa-521a-77d2-4c6b-770b04efbd56.png)

配信先はUnity Catalogの適切な場所を指定。
大事なポイントとして、チャネル設定は現状「プレビュー」にする必要があるようです。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/767b6445-e980-e746-46d6-0eda40ff8393.png)

# 実行

あとは開始ボタンを押して実行。
正常に処理が終わるとDAGで処理結果が表示されます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/320c6370-5431-b067-9f4f-de09c9bc1860.png)

カタログエクスプローラから実行結果を見ると以下の感じ。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/a287b8d7-b354-c692-1264-619d4117a572.png)

ついでに「Genieに聞く」ボタンから、Genieにサマリさせてみます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/32c088a7-a4d5-e3c2-02c6-f6459a2bad77.png)

３件だけかつ全部Positiveだったので見え方はちょっとイマイチですが、簡単に感情分析とその結果確認ができました。うーん、簡単。

ちなみに、同様のフォーマットでCSVファイルを作成してボリュームに置いていけば、DLTのパイプラインを実行するたびにデータが追記・更新されていきます。
例えば商品レビューなどのデータを定期取得する環境が整っていれば、簡単にいろんなLLMによるデータ加工処理をすることができますね。

# まとめ

ai_query&DLTを使ったデータパイプラインを作成しました。
手軽に作れる上、応用範囲も広いのでこれをベースにいろいろやれそうです。

今回やっていてハマったのはai_queryのところで、[ドキュメント](https://docs.databricks.com/en/sql/language-manual/functions/ai_query.html#arguments)IにはfailOnErrorやmodelParametersといったパラメータが紹介されているのですが、これらの指定ができなかったこと。そのあたりのアップデートがまだ自環境に反映されてないのかな。。。

とはいえ、非常に簡単にやれたので様々なユースケースへの適用をトライしてみたいと思います。