Posted inAI

AI主導のSQL最適化：LLMを活用して実行計画を解読する

AI tutorial - IT technology blog
AI tutorial - IT technology blog

背景と理由：午前2時のPagerDutyの悪夢

先週の火曜日、午前2時14分にアラートが私のスマートフォンを直撃しました。本番環境のデータベースCPU使用率は99%に張り付き、APIのレスポンスタイムは12秒まで急上昇していました。プロセスリストを素早く確認すると、犯人が見つかりました。4500万行のordersテーブルに対する、巨大でネストされたクエリです。ある開発者（おそらく3ヶ月前の私でしょう）が、適切なインデックスを貼らずにデプロイしてしまったのです。

通常であれば、EXPLAIN ANALYZEを実行し、500行に及ぶJSON形式の実行計画を凝視することになります。Sequential Scan（シーケンシャルスキャン）がどこでボトルネックを引き起こしているのかを正確に可視化するために20分は費やすでしょう。しかし、睡眠不足の状態では、ヒューマンエラーは避けられません。この精神的な重労働を肩代わりさせる方法が必要でした。ここで大規模言語モデル（LLM）の出番です。退屈な手作業を、合理化されたワークフローへと変えてくれます。

PostgreSQLで直接機械学習を実行する場合と同様に、PostgreSQLやMySQLの標準的なオプティマイザは、既存のインデックスの中から最適なパスを選択することには長けています。しかし、「そもそも作成すべきだった」インデックスを提案してくれることは滅多にありません。構造化された実行計画をLLMに読み込ませることで、不足しているインデックスに対する文脈に沿った即座の推奨事項を得ることができます。私はこのアプローチを本番環境に導入しました。その結果、数時間ではなくわずか数分で、クエリコストを一貫して85%削減できています。

インストール：AIデータベースアシスタントの設定

使い始めるのに、肥大化したエンタープライズ向けツールセットは必要ありません。データベースと、GPT-4oやClaude API Extended ThinkingをサポートするClaude 3.5 SonnetのようなLLMを橋渡しする軽量なPythonスクリプトを使用します。これらのモデルは、InstructorとPydanticでLLMから構造化データを抽出する際に見られるように、JSON実行計画の構造化された階層を解析することに驚くほど長けています。

1. 環境設定

まずは仮想環境を作成することから始めましょう。これにより、ローカル環境をクリーンに保ち、OpenAIとPostgresドライバの適切なバージョンを確保できます。

# 仮想環境を作成して有効化
python3 -m venv sql-ai-env
source sql-ai-env/bin/activate

# 依存関係のインストール
pip install openai psycopg2-binary python-dotenv

2. データベースへのアクセス

データベースユーザーにはEXPLAINを実行する権限が必要です。このガイドではPostgreSQLを使用しますが、ロジックはMySQLやSQL Serverにも完全に応用可能です。また、選択したLLMプロバイダーの標準的なAPIキーも必要になります。

設定：AIに実行計画を読み込ませる

本当のブレイクスルーは、生のクエリをAIに送るのをやめた時に起こります。単に「これをどう最適化すればいい？」と聞くだけでは、一般的なアドバイスしか返ってきません。しかし、実行計画（Execution Plan）を提供すれば、AIはエンジンの内部で何が起きているかを把握できます。Hash Joinがメモリを消費しすぎているタイミングや、Parallel Seq Scanが巨大なテーブルを叩いている箇所を正確に特定できるのです。

1. 最適化スクリプト

optimize.pyという名前のスクリプトを作成します。これは実行計画を抽出し、モデル用にフォーマットする重労働を処理します。

import os
import psycopg2
import json
from openai import OpenAI
from dotenv import load_dotenv

load_dotenv()
client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

def get_execution_plan(query):
    conn = psycopg2.connect(os.getenv("DATABASE_URL"))
    cur = conn.cursor()
    # AIに構造化データを提供するためにFORMAT JSONを使用する
    cur.execute(f"EXPLAIN (FORMAT JSON, ANALYZE) {query}")
    plan = cur.fetchone()[0]
    cur.close()
    conn.close()
    return plan

def get_ai_suggestion(query, plan):
    prompt = f"""
    あなたはシニアデータベース管理者です。
    このSQLクエリとPostgreSQLの実行計画を分析してください。
    Seq Scanやコストの高いノードなどのボトルネックを特定してください。
    
    以下を提供してください：
    1. 必要な正確なCREATE INDEXコマンド。
    2. パフォーマンスを向上させるための具体的なクエリの書き換え。
    
    クエリ：
    {query}
    
    実行計画：
    {json.dumps(plan, indent=2)}
    """
    
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content": prompt}]
    )
    return response.choices[0].message.content

# 使用例
slow_query = "SELECT * FROM orders WHERE status = 'pending' AND created_at > '2023-01-01'"
plan = get_execution_plan(slow_query)
suggestion = get_ai_suggestion(slow_query, plan)
print(suggestion)

2. システムプロンプトの微調整

高品質な結果を得るには、具体的な制約を与えることが重要です。AIに対して「カーディナリティの高いカラムを優先する」や「カバリングインデックスを好む」ように指示することで、的外れな提案を防ぐことができると分かりました。モデルには、WHERE、JOIN、またはGROUP BY句に含まれるカラムに基づいてのみインデックスを提案するように伝えてください。

検証とモニタリング：結果の妥当性確認

AIが生成したSQLを盲目的に本番環境で実行するのは、災いの元です。本番環境でのGuardrails AIの重要性を念頭に、私のワークフローには、本番データの代表的なサブセットを含むステージング環境での必須の検証ステップが必ず含まれています。

1. 「ビフォー・アフター」のチェック

変更を適用する前にパフォーマンスを測定します。実行時間と実行計画の「Total Cost」メトリクスを記録しておきましょう。インデックスを適用したら、再度EXPLAIN ANALYZEを実行します。計画がSeq ScanからIndex Scanに切り替わっていることを確認してください。最近のあるテストでは、これによりノードコストが145,000からわずか120に減少しました。

-- AIが提案したインデックス
CREATE INDEX idx_orders_status_created_at ON orders(status, created_at);

-- 改善結果の確認
EXPLAIN ANALYZE SELECT * FROM orders WHERE status = 'pending' AND created_at > '2023-01-01';

2. インデックスの肥大化（Bloat）の管理

インデックスは無料ではないことを忘れないでください。新しいインデックスを追加するたびに、INSERTやUPDATE操作にオーバーヘッドが加わります。AIが差し迫った問題を解決した後は、1週間ほどインデックスの使用状況を監視してください。以下のクエリを使用して、新しいインデックスが実際に役立っているかを確認します。

SELECT 
    relname AS table_name, 
    indexrelname AS index_name, 
    idx_scan AS times_used
FROM pg_stat_user_indexes 
WHERE indexrelname = 'idx_orders_status_created_at';

数日間のトラフィックの後でもidx_scanがゼロのままであれば、そのインデックスを削除しましょう。AIが積極的すぎたか、アプリケーションのクエリパターンが変化した可能性があります。

3. パイプラインの自動化

現在、私たちのセットアップでは、LLMとGitHub Actionsを活用して、これをCI/CDパイプラインに直接統合しています。開発者が新しい複雑なクエリを含むプルリクエスト（PR）を送信すると、GitHub Actionがサニタイズされた開発用データベースでEXPLAINを実行します。その後、LLMが最適化案をPRにコメントします。これにより、午前2時のトラブルが本番環境に到達するのを未然に防いでいます。

手動の実行計画分析からAI支援の診断に移行することで、修正にかかる時間は数時間から数秒に短縮されます。これにより、エンジニアリングチームはネストされたループのツリー図を凝視するのではなく、ハイレベルなアーキテクチャの設計に集中できるようになります。

Share:
Tags: