ドメイン知識をうまく利用してLLMにSQLを生成させる

みなさんこんにちは、グループ研究開発本部 AI研究開発室のK.Fです。

普段はモバイルアプリのデータ分析業務を中心に行っているのですが、業務の6割くらいはBigQueryでSQLを書いていることが多いです。個人的に、LLMとSQL生成はかなり相性のよいタスクだと感じており、他社さん（mercariさんやCyberAgentさんなど）でもすでに試されている方はいるようです。私も、過去の記事で、LLMにSQLを生成させる方法を検証しましたが、今回はその延長で、少しでも精度改善できそうなアイデアをいろいろ試してみようの回をしてみようと思います。

ユースケースを考える

何においてもユースケース/ペルソナを考えないと、誰にも刺さらないプロダクト（ツール）になってしまったということが起こってしまいますよね。正直、迷いなくSQLをすらすら書けるような熟練のデータアナリストにとっては、SQL生成のありがたみは大きくなく、GitHub Copilotの補完機能で十分かもしれません。一方で、オンボーディング中の新人データアナリストやSQLの基本は理解しているが、すらすら書くほどの自信がないビジネスサイドの人間は、SQL生成の恩恵を大いにうけることができるでしょう。

当然、ユーザの入力（問い合わせ）は自然言語で記述されることが想定されますし、必ずしもデータ構造やテーブルの命名を理解しているとは言えないでしょう。そのため、自然言語から適切なテーブル情報を引き出して、SQLを生成する必要があります。

改善のアイデア

SQL生成で精度を上げるには、データ構造やカラムの使われ方などアナリストやエンジニアしか知り得ないドメイン知識を利用することも多いです。そのため、いかにドメイン知識を利用するかが鍵になります。

1. ドメイン知識をLLMに与える

テーブルの概要やカラムの情報などユーザの入力した情報に付け加えてプロンプトでLLMに入力します。Retrieval Augmented Generation (RAG) といわれるテクニックで、プロンプトに入れるべきメタデータをvector storeなどのDBから動的に取得します。

2. ドメイン知識をLLMに生成させる

集計対象のテーブルが20~30個あるといった状況もよくあるでしょう。そういった場合、すべてのテーブルに対してテーブル概要やカラムの情報などを用意するのはなかなか困難です。BigQueryでは、カラムのdescriptionを入れることができますが、私が関わっているプロジェクトでもきちんと埋まっているテーブルはごく一部です…足りないものは補ってもらおうということで、LLMに生成させます。

BigQueryには監査ログをクエリすることができる機能があり、これを利用すれば対象のテーブルがいつどのようなクエリで使われたのか？を取得することができます。これを利用してクエリ例をプロンプトに入れてメタデータを生成させます。

システムに落とし込む

システムの流れに落とし込むと以下のようになります。

1. 問い合わせに対して、関連度の高いテーブルをリストアップする
2. SQL例を与えて、テーブルの概要を生成させる
3. テーブルの概要からSQLを生成する

検証

パーツパーツでの検証になりますが、実際にシステムの流れを実現できそうなのか？を検証していきます。

1. 問い合わせに対して関連度の高いテーブルをリストアップする

まず、BigQueryの監査ログから使われている頻度の高いテーブル一覧をリストアップします。

with
audit_logs as (
  select
    protopayload_auditlog.servicedata_v1_bigquery.jobGetQueryResultsResponse.job.jobStatistics as job_statistics,
  from
    `your-project.cloud_audit_logs.cloudaudit_googleapis_com_data_access_20*`
  where
    parse_date("%y%m%d", _table_suffix) between current_date() - 30 and current_date() -- 仮に過去30日間のデータを対象とする
)

select
  tables.projectId || "." || tables.datasetId || "." || tables.tableId as table_id,
  count(1) as count_uses,
from
  audit_logs, unnest(job_statistics.referencedTables) as tables
group by
  table_id
order by
  count_uses desc

このようなクエリで簡単に取得することができます。この結果をプロンプトに埋め込んで、LLMに問い合わせをします。

あなたはデータアナリストです。Google Cloud BigQueryを使いこなすことができます。

# 質問
{{ question }}

上の質問に対して関連度が高いと推測できるテーブルを以下のテーブル一覧から選んでください。（複数選択可、テーブル一覧は利用頻度の高いものから順に並んでいます）
{% for table_id in table_usages %}
・{{ table_id }}
{% endfor %}

questionが利用者が入力する問い合わせ文言で、table_usagesが上のクエリの出力結果です。

実際に、試してみると、テーブルの名前しか与えていないのにも関わらず、関連するテーブルをそこそこの精度で拾ってこれています。（テーブル名はマスクしてしまっていますが、実際に集計しようと思うと利用するテーブルをきちんと拾ってこれていました。）

問い合わせ）
  アクティブユーザの1日あたりの平均売上額を求めてください

回答）
  アクティブユーザの1日あたりの平均売上額を求めるためには、アクティブユーザに関するデータと売上に関するデータが必要です。選ぶべきテーブルはアクティブユーザの定義と売上の定義に依存しますが、一般的には次のようなテーブルが関連する可能性があります。

  - `datamartにあるユーザの行動ログを集約したテーブル`: アクティブユーザの活動を示す日次ログが含まれている可能性があります。
  - `dashboardのために日次で売上を集計しているテーブル`: 日次の利益に関するデータで、売上に直接関連する内容が含まれる可能性が高いです。
  - `広告売上の生ログが入ったテーブル`: 広告収入に関するレポートが含まれており、売上の一部を表す可能性があります。
    ※ 実際のテーブル名がはいっていましたが、マスクしています

  これらのテーブルの中で、具体的にどのテーブルがアクティブユーザと売上のデータを含んでいるかは、テーブルのスキーマや実際のデータ内容を確認する必要があります。また、売上とは具体的に何を指すのか（商品販売、サービス提供、広告収入など）によって、関連するテーブルが異なることも考慮する必要があります

2. SQL例を与えて、テーブルの概要を生成させる

次に、あるテーブルに対して発行されたSQL例をリストアップします。

declare target_date DATE default "2024-01-01";

create temp function query_priority (
  is_destination BOOL,
  type STRING
) as (
  case
    when is_destination = true and type = "CREATE_TABLE_AS_SELECT" then 1
    when is_destination = true and type = "INSERT" then 2
    when is_destination = true and type = "MERGE"  then 2
    when type in ("SELECT", "CREATE_TABLE_AS_SELECT") then 3
    when type in ("UPDATE", "DELETE") then 4
    else 5
  end
);

with
audit_logs as (
  select
    protopayload_auditlog.authenticationInfo.principalEmail as email,
    protopayload_auditlog.servicedata_v1_bigquery.jobGetQueryResultsResponse.job.jobConfiguration as job_config,
    protopayload_auditlog.servicedata_v1_bigquery.jobGetQueryResultsResponse.job.jobStatistics as job_statistics,
  from
    `your-project.cloud_audit_logs.cloudaudit_googleapis_com_data_access_20*`
  where
    parse_date("%y%m%d", _table_suffix) = target_date
)

select distinct
  job_config.query.query as query,
  query_priority(
    job_config.query.destinationTable.tableId like "your_table_name_with_%",
    job_config.query.statementType
  ) as priority,
from
  audit_logs, unnest(job_statistics.referencedTables) as refrenced_table
where 1=1
  and email = "your_email_address or your_service_account_email_address"
  and (
    (
      job_config.query.destinationTable.datasetId = "your_dataset_name"
      and job_config.query.destinationTable.tableId like "your_table_name_with_%"
    ) or (
      refrenced_table.datasetId = "your_dataset_name"
      and refrenced_table.tableId like "your_table_name_with_%"
    )
  )
order by
  priority

このようなSQLで抽出することができます。BigQueryでは、table_name_20240101のような日付で分割されたシャーディングテーブルを利用することも多いので、likeと%をうまく利用してより多くのクエリが引っかかるようにしています。また、emailで絞り込むことができるので、シニアデータアナリストのメールアドレスやデータマートの作成のためのバッチ処理で利用しているサービスアカウントのメールアドレスで絞り込むと、抽出されるクエリの品質をある程度担保することができそうです。もうひとつ、ポイントがあって、すべてのクエリをそのままプロンプトに入れるのはトークンサイズの制約的に難しいので、クエリを何かしら優先順位付けをしてあげる必要があります。この例では、該当のテーブルを作成するクエリがもっとの情報量が多いだろうと思われるので、テーブルを作成するために利用されたクエリとそのデータを更新するために利用されてクエリの優先度が高くなるようにしています。その次点で、そのテーブルを単純に利用しているだけのクエリがくるでしょう。上記SQLの上部にあるquery_priorityというtemp関数がその処理にあたるものです。

なるべく同じようなクエリをサンプルとして与えないように、クエリ間の類似度をとって、遠いものを選抜するなどできれば、もっと精度があがりそうですね。

ユーザの行動ログ(firebase)を日次で集約したデータマートのテーブルに対しての出力例）

  このテーブル `datamartにあるユーザの行動ログを集約したテーブル`(名前はマスクしています) は、日々のユーザーのアプリケーション利用に関するログデータを格納するためのものであると推測されます。以下は各カラムの概要です:

  - `event_date`: イベントが発生した日付。このカラムによりテーブルがパーティションされています。
  - `user_id`: ユーザーの識別子。
  - `os_name`: ユーザーが使用しているオペレーティングシステムの名前。
  - `app_version`: アプリケーションのバージョン。
  - `event_name`: 発生したイベントの名前。
  - `param_name`: イベントに関連するパラメータの名前。
  ... 略
  - `count_logs`: 特定のイベントが記録された回数。
  - `created_at`: レコードが作成された時刻。
  - `updated_at`: レコードが更新された時刻。
  
  提供されたSQL例を見ると、このテーブルは他のデータソースからのデータを集約して、イベントログデータの集計レポートを作成するために使用されています。例えば、`別の中間テーブル` というテーブルから特定の日付範囲のデータを取得し、それを `このテーブル` テーブルに集約して挿入しています。

3. テーブルの概要からSQLを生成する

最後に、テーブルの候補と質問をあわせてプロンプトを作成し、LLMに問い合わせます。

あなたはデータアナリストです。Google Cloud BigQueryを使いこなすことができます。

# 質問
{{ question }}

以下のテーブル候補を利用して、質問に対するSQLを生成してください。

# テーブル候補
{% for table_id, table_description in table_candidates %}
- テーブルID: {{ id }}
{{ table_description }}
{% endfor %}

実際の出力されたSQLはお見せすることができませんが、アクティブの定義など質問がそもそも抽象的になっている部分はうまく扱えていないですが、すこし手直しすれば正しそうなクエリを出力できていました。

まとめ

今回は、ドメイン知識をうまく利用してLLMでSQLを生成させてみました。LLMには当然、ハルシネーションの問題もあり、完全にデータアナリストの仕事を置き換えることはできないですが、新入りのデータアナリストがLLMに質問しながら、クエリのアウトプットの質をあげることが十分にできそうな状態にあると感じました。かなり可能性を感じたので、今後も検証を続けていきます。

最後に

グループ研究開発本部 AI研究開発室では、データサイエンティスト/機械学習エンジニアを募集しています。ビッグデータの解析業務などAI研究開発室にご興味を持って頂ける方がいらっしゃいましたら、ぜひ募集要項一覧からご応募をお願いします。 一緒に勉強しながら楽しく働きたい方のご応募をお待ちしております。