Postgresよりも81%高速なクエリプランを生成する4Bモデルのトレーニング
Training a 4B model to produce 81% faster query plans than Postgres (rohanbansal.com)
要約
本記事では、小規模なオープンウェイトモデルを、教師ありファインチューニング(SFT)とエージェント型強化学習(RL)を用いて、Postgresのデフォルトプランを上回るクエリプランを生成できるようにする実験について詳述しています。実験の結果、4Bモデルは113個の結合負荷の高いクエリにおいて44.7%のレイテンシ削減を達成し、当初は99個のクエリでプランを生成できませんでした。このアプローチは、クエリプラン生成という、検証可能な出力を持つタスクに言語モデルを適用する可能性を示しています。
全文翻訳
4BモデルでPostgresより81%高速なクエリプランを生成するトレーニング
クエリプランナーは、本当にどれほど優れているのか?
Leisらは2015年にまさにこの問いを投げかけました。そして、10年後に再び問いかけました。
彼らの最初の調査から10年間にわたる膨大な研究にもかかわらず、クエリプランナーは依然として改善の余地が大きいことが判明しました。
私がこれを初めて知ったとき、驚きました。Postgresデータベースは、テーブル内のデータについてすべてを知っているはずですよね? それがどうしてそんなに難しいのでしょうか?
実際には、非常に難しいのです。クエリプランナーが実行する必要のある特定のタスク、つまり結合順序付けは、NP困難であることが知られています。
したがって、クエリプランナーは難しいです。しかし、プランナーが選択したクエリプランが良いものであるかどうかを検証することは、それほど難しくありません。簡単に言えば、良いクエリプランナーは高速に実行されるプランを生成し、悪いプランナーは遅いプランを生成します。言語モデルは、検証可能な出力を持つタスクを学習することに特に長けています。最適化すべき単一の軸(クエリの実行時間)があるため、この問題は、モデルをより高速なクエリプランを生成するように導く行動を強化することに美しく還元されます。
以下は、小規模なオープンウェイトモデルを、教師ありファインチューニング(SFT)とエージェント型強化学習(RL)でポストトレーニングすることにより、Postgresのデフォルトプランを上回るPostgresクエリプランを生成できるか、という問いを探求するために実行した実験の概要です。
私たちの問いに対する答えは、紛れもない「はい」です。ハイライトは以下の通りです。
- 99件のクエリでプランを生成できなかった初期状態の4Bモデルが、113件の結合が多いクエリ全体でレイテンシを44.7%削減しました。
- 並行コンテナ間のLinuxページキャッシュの競合ノイズを最小限に抑えるPostgres測定リグを構築しました。
- 本質的にノイズの多い環境でRLのロールアウトをスコアリングするためのカスタムGRPOバリアントを設計しました。
- RLを2台のマシンに分割しました。vLLMとトレーナーはレンタルした2x H100ノードで実行し、4つのPostgresコンテナは私のデスクで実行しました。
- 半千個のGPT-6 Astraエージェント軌跡にわたるオフポリシー蒸留を実行しました。
最初から始めましょう。
クエリプランナーの内部
IMDbデータセットの以下のスライスを考えてみましょう。
-- IMDbのタイトル(映画、シリーズ、エピソードなど)[約100万行]
title (
id integer PRIMARY KEY,
title text,
production_year integer,
kind_id integer -- FK -> kind_type
)
-- 映画 <> 会社 結合テーブル [約200万行]
movie_companies (
id integer PRIMARY KEY,
movie_id integer, -- FK -> title.id
company_id integer, -- FK -> company_name.id
company_type_id integer, -- FK -> company_type.id
note text
)
-- 会社の名前、出身地など [約10万行]
company_name (
id integer PRIMARY KEY,
name text,
country_code text -- '[us]', '[jp]', ...
)
-- タイトルの会社の種類ルックアップテーブル [4行]
company_type (
id integer PRIMARY KEY,
kind text -- 'production companies', 'distributors', ...
)
-- タイトルが何であるかのルックアップテーブル [7行]
kind_type (
id integer PRIMARY KEY,
kind text -- 'movie', 'tv series', 'episode', ...
)
「2000年代に最も多くのタイトルを公開した日本企業はどこか?」という質問に答えようとしているとします。以下のようなクエリを作成するかもしれません。
SELECT cn.name,
COUNT(*) AS titles
FROM title AS t,
movie_companies AS mc,
company_name AS cn
WHERE t.id = mc.movie_id
AND mc.company_id = cn.id
AND cn.country_code = '[jp]'
AND t.production_year BETWEEN 2000 AND 2009
GROUP BY cn.name
ORDER BY titles DESC
LIMIT 10;
このクエリは、2000年から2009年の間に公開されたタイトルの数でソートされた、上位10社の日本企業を出力します。
しかし、Postgresはこれらの結果をどのように得たのでしょうか?
Postgresがこのデータを提供するためにたどったパスは、必然的なものではなく、選択的述語(つまり、WHERE句のフィルタリング条件)とすべて関係があります。
これを説明するために、日本企業のフィルタリングや日付範囲のフィルタリングがない場合の同じクエリを想像してみましょう。
SELECT cn.name,
COUNT(*) AS titles
FROM title AS t,
movie_companies AS mc,
company_name AS cn
WHERE t.id = mc.movie_id
AND mc.company_id = cn.id
GROUP BY cn.name
ORDER BY titles DESC
LIMIT 10;
mcはmc.company_id = cn.idを介してのみcnと結合でき、tはt.id = mc.movie_idを介してのみmcと結合できます。
これらの制約により、2つの結合ツリーが生成されます。交換可能性を考慮すると、技術的には8つの結合ツリーが存在します。この場合、関係のサイズに影響しないため、考慮しません。
⋈ ⋈ t cn mc (cn ⋈ mc) ⋈ t
⋈ ⋈ cn t mc (t ⋈ mc) ⋈ cn
私たちのクエリの2つの結合ツリー。下側の結合が最初に実行され、その結果がルート結合の入力となります。
テーブルまたはクエリ結果のカーディナリティは、それが含む行数です。関連テーブルが以下のカーディナリティを持つと仮定します。
- cn=100k
- mc=2m
- t=1m
結合を考慮すると、以下のカーディナリティが得られます。
(cn⋈mc)=2m、その後⋈t=2m
(t⋈mc)=2m、その後⋈cn=2m
これらの3つのテーブルを結合する順序に関係なく、2番目の結合には常に同じ2m行が渡されます。
ここで、選択的述語を戻しましょう。
- cn'=5k(10万社のうち5%が日本企業であると仮定)
- mc=2m(変更なし)
- t'=200k(100万件のタイトルのうち20%が2000年代に作成されたと仮定)
(cn'⋈mc)≈100k、その後⋈t'≈20k
(t'⋈mc)≈400k、その後⋈cn'≈20k
最初の結合順序は、200万件のmovie_companiesエントリを、日本企業である5%のスライスにフィルタリングします。一様分布を仮定すると(なぜこれを仮定するのかは後で議論します)、この結合は約10万行の結果をもたらします。結果をフィルタリングされたtitleテーブルと結合すると、2000年代のそれらの行の20%のみが保持されます。
2番目の結合順序は、200万件のmovie_companiesエントリを、2000年代に作成されたタイトルの20%のスライスにフィルタリングします。同じ一様性の仮定が当てはまるため、最初の結合は約40万行の結果をもたらし、2番目の結合に40万行が渡されることになります。
2番目の結合順序を選択した場合、4倍の作業を行います。
残念ながら、それだけでは終わりません。
組み合わせ爆発
各結合は、以下のいずれかを使用できます。
- ハッシュ結合
- マージ結合
- ネストループ結合
交換可能性を再度考慮すると、4つの異なる外側/内側結合方向があり、8つの可能な組み合わせが生じます。
(cn⋈mc)⋈t
t⋈(cn⋈mc)
(mc⋈cn)⋈t
t⋈(mc⋈cn)
(t⋈mc)⋈cn
cn⋈(t⋈mc)
(mc⋈t)⋈cn
cn⋈(mc⋈t)
最後に、各テーブルは異なる方法でスキャンできます。4種類のスキャンのみを考慮すると、次のようになります。
- シーケンシャル
- インデックス
- インデックスのみ
- ビットマップ
これにより、4,608通りのクエリ実行方法が存在します。これは実際には過小評価です。プランは並列実行でき、集計はハッシュ化またはソートされる可能性があります。
また、Postgresはこれらすべてのプランを評価するわけではないことに注意してください。動的計画法(および12以上の結合を含むクエリには遺伝的アルゴリズム)を使用して探索空間を刈り込みます。
[画像: 画像: クエリプランの探索空間のサイズを示すグラフ]
さらに悪いことに、各結合は探索空間を組み合わせ的に爆発させます。
SELECT MIN(t.title) AS movie_title
FROM company_name AS cn,
keyword AS k,
movie_companies AS mc,
movie_keyword AS mk,
title AS t
WHERE cn.country_code ='[de]'
AND k.keyword ='character-name-in-title'
AND cn.id = mc.company_id
AND mc.movie_id = t.id
AND t.id = mk.movie_id
AND mk.keyword_id = k.id
AND mc.movie_id = mk.movie_id;
結合されているテーブルの数を選択するにはクリックしてください。4つのテーブル以降、左側のクエリはJOBからのものです。右側は探索空間のサイズの概算です。
推定、カウントではない
Postgresはここで難しい立場にいます。カーディナリティを単純にカウントし、後続の結合に渡される行数を最小限に抑えるプランを選択できると考えるのは合理的でしょう。
しかし、これはPostgresがクエリプランニング中にカーディナリティをカウントできることを意味します。それはできません。これを把握するには、各結合を実行して結果の行数を数える必要があります。これは高速なクエリプランナーの目的を損ないます。クエリプランナーはコスト最小化を正確に行うことを目指すのではなく、多くの種類のクエリで「十分に良い」ことを目指します。
代わりに、Postgresは統計情報を使用してカーディナリティを推定します。プランナーはpg_statisticテーブルに問い合わせ、各列の一般的な値とその頻度、および残りの値のヒストグラムを取得します。結合を追加すると、事態はさらに複雑になります。Postgresは、あるテーブルの行が別のテーブルの行にどのように分布しているかを知りません。これを回避するために、最初のテーブルの特定の値の頻度を2番目のテーブルに単純に適用できると仮定します。これが前述の一様分布の仮定です。
一様分布を仮定することはヒューリスティックとしては問題ありませんが、失敗すると深刻な問題を引き起こします。前の結合順序に戻って見てみましょう(cn'⋈mc)≈100k、その後⋈t'≈20k。200万件のmovie_companiesエントリを、5%が日本企業であるという仮定でフィルタリングしました。しかし、日本企業である5%が、実は映画の50%を占めていたらどうなるでしょうか?最初の結合は100万行を生成するでしょう!コストモデルは最初の結合順序を選択するように言いますが、実際には、2番目の結合の方が、2番目の結合に40万行しか渡さないため、実際にはより良いです。
Postgresの仮定 - 100k行
実際 - 100k行
⋈ ⋈ t′ cn′ mc (cn′ ⋈ mc) ⋈ t′
⋈ ⋈ cn′ t′ mc (t′ ⋈ mc) ⋈ cn′
movie_companies行のうち、日本企業に属する割合:(一様)
スライダーをドラッグして、日本企業がより生産的になるようにしてください。Postgresの推定値は静的である一方、実際の行数は影響を受けることに注意してください。
最初の結合での1つの悪い推定値は、結合ツリー全体に連鎖し、他のすべての推定値を破損させる可能性があります。
象を操る方法
Postgresは常に最もコストの低いプランを選択しますが、ソースコードを変更せずにコストモデルを変更することはできません。では、コストの高い異なるプランを選択するように、どのように誘導できるでしょうか?
そこで登場するのがpg_hint_planです。
pg_hint_planは非常にシンプルなサードパーティ拡張機能です。SQLステートメントの上に構造化された「ヒント」をコメントとして追加するだけで、Postgresをヒントで提供された指示を使用するプランに誘導できます。例えば:
/*+ HashJoin(a b) SeqScan(a)*/EXPLAIN SELECT * FROM pgbench_branches b JOIN pgbench_accounts a ON b.bid = a.bid ORDER BY a.aid;
QUERY PLAN--------------------------------------------------------------------------------- Sort (cost=31465.84..31715.84 rows=100000 width=197) Sort Key: a.aid -> Hash Join (cost=1.02..4016.02 rows=100000 width=197) Hash Cond: (a.bid = b.bid) -> Seq Scan on pgbench_accounts a (cost=0.00..2640.00 rows=100000 width=97) -> Hash (cost=1.01..1.01 rows=1 width=100) -> Seq Scan on pgbench_branches b (cost=0.00..1.01 rows=1 width=100)(7 rows)
pg_hint_planのドキュメントからの例。
このヒントは、pgbench_accountsとpgbench_branchesを結合するためにHashJoinを使用し、pgbench_accountsテーブルのシーケンシャルスキャンを実行することを義務付けています。実際のクエリプランもそれにうまく従っています。
問題の定式化
pg_hint_planヒントを使用してPostgresに異なる、そして潜在的により良いクエリプランを選択させることを考えると、私たちが開始している質問は次のとおりです。
言語モデルは、より良いクエリプランにつながるヒントを生成することを学習できるか?
有用な研究
これが解決する価値のある問題である可能性のあるものは何でしょうか?
私の最初の考えは、モデルにクエリとPostgresプランナーが持つまったく同じ情報を提供することでした。これは、より良いカーディナリティ推定器を構築できるかどうかを確認することになります。私は、これは探求する価値のある道ではないと結論付けました。カーディナリティ推定に関する数十年の研究と戦うことになります。さらに、推論レイテンシだけでも、Postgresの超高速クエリプランナーと比較して、学習された有用性をはるかに上回るでしょう。
2番目のアイデア、そして私が正しい定式化だと信じているものは、特定のデータベース使用パターン、つまり重い分析ワークロードにあります。クエリが最適化されていないデフォルトのPostgresプランを使用して数千回実行される場合、効率の向上がテーブルに残されています。代わりに、モデルは特定のクエリを実行するより良い方法を見つけるようにトレーニングされる可能性があります。トレーニングプロセスでは、事前にそのクエリを数十回から数百回実行する必要があるかもしれませんが、クエリのすべての実行にわたる償却コストは劇的に低くなります。
目標は、単発クエリのパフォーマンス/効率のパレートフロンティアでPostgresに勝つことではなく、繰り返し実行されるクエリでそれを上回ることができるかもしれません。
モデルとそのハーネス
自分でトレーニング/推論を自宅の2x RTX 3090リグ(愛称FLOPper)で行うのが最も簡単だったため、小規模な4Bモデルから始めることにしました。
このプロジェクトを開始した頃、Qwen 3.8ファミリーのモデルがリリースされましたが、残念ながら4Bバリアントはありませんでした。しかし、ドイツの小さな研究所であるEmperoによるQwen 3.8 4Bの蒸留モデルに出会い、興味を引かれました。彼らはQwen 3.8の2.4Tモデルを教師モデルとして使用し、Qwen 3.5 4Bに学習を蒸留してempero-ai/Qwen3.8-4B-Distillを生成しました。この蒸留モデルは、ベースの3.5モデルよりも outright に優れているわけではありません。MMLUタスクではパフォーマンスが向上し、GSM8Kタスクではわずかに低下します。つまり、この蒸留は、一般的な知識の広さで評価するとパフォーマンスが向上し、多段階の数学的推論ではわずかに低下します。どちらが私たちのタスクに適しているかはわかりません。いずれにしても、蒸留モデルを使用することにしました。
モデルを固定した後、ヒント生成をオーケストレーションする軽量エージェントハーネスqo-agentを構築しました。これには以下の6つのツールが与えられました。
inspect_relation— テーブルのカラム(型とNULL許容性)、インデックス定義、推定行数とバイト数をリストします。get_column_stats— リレーションのカラム(1〜8個)のPostgresプランナー統計情報を取得します。get_plan— デフォルトプランの推定値、または提出された候補の保存済みプランを取得します。evaluate_candidate— 提案されたプランアクションを検証し、タイミング/プラン診断のために実行します。keep_default— Postgresのデフォルトプラン自体を候補として返し、検索を終了します。finish— 提出された候補IDまたはデフォルトプランを入力として受け取り、検索を終了します。
構造化された出力を活用するため、エージェントはPlanAction JSONオブジェクトを生成するように指示されました。evaluate_candidateへの呼び出しは、PlanActionオブジェクトをヒントにコンパイルし、元のクエリの前に付けました。
サンプルエージェント軌跡:
Agent → get_plan(