Digital Reactor
機械学習

BigQueryオプティマイザのコストモデルに外部APIの料金は考慮されない：BigQuery MLの注意点

BigQueryオプティマイザのコストモデルに外部APIの料金は考慮されない：BigQuery MLの注意点

はじめに

SQLは、書いた順に処理されるとは限りません。実行計画を組むのはオプティマイザで、フィルタと関数呼び出しのどちらを先に評価するかもエンジン側に委ねられています。ふだんは意識しなくてよい性質です。ところが1行ごとに外部のモデルを呼ぶ関数が混ざると、話が変わります。何行に対して呼ぶかがそのまま外部サービスへの課金になるので、評価の順序が結果は変えずとも、請求額を左右するようになります。そして、オプティマイザが計画を選ぶときのコストモデルには、その外部APIの料金は入っていません。

今回、会議トランスクリプトを文章のベクトル（embedding）に変換する毎時のバッチで、未処理の行だけを処理する差分更新のはずが、毎回ソース全件の行数ぶん呼ばれていました。新しく処理した行が0件の時間帯にも、1時間あたり2,000万トークン（数値はクライアント情報のため丸めています）が計上され続けていました。フィルタは書いてあり、テーブルの中身も正しく、ジョブも成功しています。

対象読者:

  • BigQuery MLやDataformで、embeddingや生成AIの呼び出しをバッチに組み込んでいる方
  • Vertex AIの請求の内訳が説明できず、切り分けたいデータ基盤の担当者
  • LLMを含むパイプラインで、コストの出どころを把握しておきたい方

記事のポイント:

  • 絞り込みをML.GENERATE_EMBEDDINGの入力サブクエリに置いてDML（Data Manipulation Language、INSERTやMERGEなどテーブルを書き換える文）に包むと、先にembeddingが評価されることがあります
  • BigQueryのコストモデルに外部APIの料金は入っていないため、呼び出し回数を減らす向きに計画が寄る保証はありません
  • 対処は、実行計画に任せず入力の集合を前のステートメントで確定させる、明示的な書き方に変えることです

危ない書き方と安全な書き方

まず、今回の危ない書き方と安全な書き方をお見せします。問題になるのは、未処理だけを残すanti-join（既存側に一致がない行だけを取る結合）を、embeddingの入力サブクエリの内側に置いた形です。これをINSERTやMERGEに包むと、絞り込みより先にembeddingが評価されることがあります。Dataformのtype: "incremental"でuniqueKeyを指定すると、書いたクエリがそのままMERGEに包まれるので、この形に当たります。

-- 危ない形：絞り込みが embedding の入力サブクエリの中にある
INSERT INTO mart.transcript_embeddings
SELECT chunk_id, ml_generate_embedding_result AS embedding
FROM ML.GENERATE_EMBEDDING(
  MODEL mart.gemini_embedding_001,
  ( SELECT s.chunk_id, s.text AS content
    FROM stg.transcript_chunks s
    LEFT JOIN mart.transcript_embeddings d ON d.chunk_id = s.chunk_id
    WHERE d.chunk_id IS NULL ),        -- ここで0件に絞れているつもり
  STRUCT(TRUE AS flatten_json_output));

安全なのは、絞り込みを前のステートメントで終わらせてから渡す形です。

-- 安全な形：候補を確定させてから embedding に渡す
CREATE TEMP TABLE pending AS
SELECT s.chunk_id, s.text AS content
FROM stg.transcript_chunks s
LEFT JOIN mart.transcript_embeddings d ON d.chunk_id = s.chunk_id
WHERE d.chunk_id IS NULL;

INSERT INTO mart.transcript_embeddings
SELECT chunk_id, ml_generate_embedding_result AS embedding
FROM ML.GENERATE_EMBEDDING(
  MODEL mart.gemini_embedding_001,
  (SELECT chunk_id, content FROM pending),
  STRUCT(TRUE AS flatten_json_output));

違いは、絞り込みがembeddingと同じステートメントにあるか、前のステートメントで完了しているかだけです。書き込まれる結果はどちらも同じで、変わるのはモデルに渡る行数です。DataformならCREATE TEMP TABLEをpre_operationsブロックに移します（同じアクションのステートメントは順に実行されるので、本クエリから参照できます）。実行環境の都合でジョブが分かれる場合は、TEMPではなく_pendingのような通常のテーブルにします。

同じ本体でも、SELECTでは再現せずDMLで課金される

ML.GENERATE_EMBEDDINGはテーブル関数で、入力のサブクエリを受け取ってリモートのモデルを呼びます。書き方が入れ子になっているので実行順もそのとおりだと読めてしまいますが、見た目が入れ子になっていることは、必ずしも絞り込みが先に評価されることを保証しません。

読んで返すだけのSELECTと、テーブルの中身を書き換えるDMLでは、同じサブクエリを抱えていても実行計画の組み立てが変わります。DMLでは書き込み先テーブルの読み取りと更新が1つの文に同居し、MERGEならさらに照合条件（ON句）の評価が加わるため、ソース側をどこまでそろえてから照合に渡すかという段取りが計画に入ってきます。今回、絞り込みの位置が入れ替わったのはこのDML側だけでした。

切り分けのために、入力が0件になる3つの形を同じモデル・同じソースに対して順に流し、そのたびに呼び出しトークン数を確認しました。判定の要になるのは、危ない形の中身をそのままSELECTで実行した2番目です。

-- anti-join の結果が 0 件になる SELECT（本番とまったく同じ絞り込み）
SELECT COUNT(*)
FROM ML.GENERATE_EMBEDDING(
  MODEL mart.gemini_embedding_001,
  ( SELECT s.chunk_id, s.text AS content
    FROM stg.transcript_chunks s
    LEFT JOIN mart.transcript_embeddings d ON d.chunk_id = s.chunk_id
    WHERE d.chunk_id IS NULL ),        -- 未処理は 0 件
  STRUCT(TRUE AS flatten_json_output));
実行した形絞り込みの結果書き込み行数呼び出しトークン
空の入力を渡すSELECT0件書き込みなし0
anti-joinで0件になるSELECT0件書き込みなし0
同じ本体をMERGEに包む0件0行約2,000万

SELECTの2つは、入力が0行なら呼び出しも0回で、想定どおりでした。ところが同じ本体をINSERTやMERGEに包むと、書き込みは0行のままトークンだけが1回あたり約2,000万出ます。1チャンクあたりおよそ1,000トークンなので2万行ぶん、つまりソース全件の量です。

デバッグが難しいのはこの一致しなさのせいです。手元で怪しい箇所を切り出してSELECTで実行すれば、行数もトークンも想定どおりに出ます。そこで問題なしと判断してしまうと、実際に課金しているDMLの側は最後まで疑いから外れたままになります。

コストモデルが見ているのは、バイト数とシャッフルとスロットだけ

絞り込みが後ろに回りうる理由は、オプティマイザが計画を選ぶときのコストモデル（どの計画が安いかを判断する計算）に、外部の課金が入っていないことです。コストモデルが見ているのは、読むバイト数、シャッフルする量、消費するスロットです。外部モデルの呼び出しが1行あたりいくらかかるかという情報はそこにないため、呼び出し回数を減らす向きに計画を寄せる動機がありません。どの段階で入れ替わったのかまで実行計画から追い切れたわけではありませんが、絞り込みの位置に課金を任せられないという結論は変わりません。

仕様上の誤りではありません。書き込まれる結果は正しく、SQLの意味論としても筋が通っています。ただ、評価の順序が保証されていないなら、そこに呼び出し回数がぶら下がっている状態のままでは運用できません。順序が入れ替わっても結果は同じなので、将来のバージョンで計画が変わっても誰も気づけないからです。同じ注意はAI.GENERATE系の関数や、外部APIを呼ぶリモート関数にも当てはまります。

書き込み0行なのにトークンが出ている実行を探す

自分の環境で同じことが起きていないかは、2つの数字を突き合わせれば判定できます。厄介なのは、トークン数がBigQueryのジョブ統計に出てこないことです。呼び出し量はCloud Monitoringのaiplatform.googleapis.com/publisher/online_serving/token_countをモデルのラベルで絞って見ます。書き込み行数のほうはジョブ側から取ります。

-- 書き込み行数を時刻つきで取り、トークン数の指標と突き合わせる
SELECT creation_time, job_id, statement_type,
       dml_statistics.inserted_row_count AS inserted_rows, total_slot_ms
FROM `region-asia-northeast1`.INFORMATION_SCHEMA.JOBS_BY_PROJECT
WHERE creation_time > TIMESTAMP_SUB(CURRENT_TIMESTAMP(), INTERVAL 6 HOUR)
  AND query LIKE '%ML.GENERATE_EMBEDDING%'
ORDER BY creation_time DESC;

探すのは、inserted_rowsが0なのに同じ時刻にトークンが立っている実行です。差分更新のバッチであれば、追加が0件の時間帯にトークンも0へ落ちているかどうかが、そのまま健全性の判定になります。毎時実行なら1日で24回ぶんの点が並ぶので、1日ぶんさかのぼれば判断できます。

安全な形に書き換えたあとは、候補が0件の実行でトークンもほぼ0になり、毎時のスケジュールのままスパイクも出ていません。この突き合わせを毎回の実行で自動化するなら、Cloud Run jobsだけで組んだバッチ運用で使っている監視の枠にそのまま載せられます。

まとめ

BigQueryのオプティマイザは、読むバイト数とスロットでコストを見積もります。外部モデルの呼び出しが1行いくらで課金されることはその見積もりに入っておらず、SQLに書いた絞り込みが先に評価される保証もありません。結果は変わらないまま、呼び出し回数だけが変わります。今回は差分更新のはずが毎回ソース全件ぶん呼ばれ、入力を前のステートメントで確定しただけでほぼ0になりました。

自分のパイプラインで1行ごとに外部を呼ぶ関数を探し、その入力がDMLの内側で絞られていないかを見てください。内側にあるなら、絞り込みを前のステートメントへ出す形に書き換えます。そのうえで、書き込みが行数に対してトークン消費が過剰になっていないかを、ジョブ統計とトークン数の指標で一度突き合わせるとよいでしょう。

関連記事

← 技術ブログ一覧へ