AI・機械学習
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個のクエリでプランを生成できませんでした。このアプローチは、クエリプラン生成という、検証可能な出力を持つタスクに言語モデルを適用する可能性を示しています。
全文翻訳
クエリオプティマイザは、実際にはどれほど優れているのでしょうか? Leisらは2015年にまさにこの問いを投げかけました。そして、10年後に再び問いかけました。彼らの最初の探求から10年が経過し、膨大な研究が行われているにもかかわらず、クエリオプティマイザは依然として多くの改善の余地を残していることが判明しました。
私がこれを初めて知ったとき、驚きました。Postgresデータベースは、テーブル内のデータについてすべて知っているはずですよね? それはどれほど難しいのでしょうか? 実際には、非常に難しいのです。事実、クエリオプティマイザが実行する必要のある特定のタスクである結合順序付けは、NP困難であることが知られています。したがって、クエリオプティマイザは難しいのです。
オプティマイザが選択したクエリプランが良いものであるかどうかを検証することは、それほど難しくありません。簡単に言えば、良いクエリオプティマイザは高速に実行されるプランを生成し、悪いオプティマイザは遅いプランを生成します。言語モデルは、検証可能な出力を持つタスクを学習するのに特に優れています。最適化すべき単一の軸、つまりクエリの実行時間があるため、問題はモデルをより高速なクエリプランを生成するように導く動作を強化することに美しく還元されます。
以下は、小規模なオープンウェイトモデルを、教師ありファインチューニング(SFT)とエージェント型強化学習(RL)を用いて、Postgresのデフォルトプランを上回るクエリプランを生成できるようにするかどうかを探求するために実行した実験の内訳です。私たちの質問に対する答えは、断固たる「はい」です。
ハイライトには以下が含まれます:
当初は99個の結合負荷の高いクエリでプランを生成できなかった4Bモデルで、113個のクエリ全体で44.7%のレイテンシ削減を達成
並行コンテナ間のLinuxページキャッシュ競合ノイズを最小限に抑えるPostgres測定リグの構築
本質的にノイズの多い環境でのRLロールアウトのスコアリングのためのカスタムGRPOバリアントの設計
RLを2台のマシンに分割: vLLMとトレーナーはレンタルした2x H100ノードで実行し、4つのPostgresコンテナは私のデスクで実行
GPT-6 Astraエージェントの軌跡を500個以上横断するオフポリシー蒸留の実行
最初から始めましょう。
クエリオプティマイザの内部
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つの有効な結合ツリーが生成されます。
(cn ⋈ mc) ⋈ t
⋈
title relation
cn company_name relation
mc movie_companies relation
(cn ⋈ mc) ⋈ t
(t ⋈ mc) ⋈ cn
⋈
company_name relation
t title relation
mc movie_companies relation
(t ⋈ mc) ⋈ cn
私たちのクエリの2つの結合ツリー。下側の結合が最初に実行され、その結果がルート結合への入力となります。
テーブルまたはクエリ結果のカーディナリティは、行数です。関連テーブルのカーディナリティが次のようになっていると仮定します:
cn=100k, mc=2m, t=1m
結合を考慮すると、次のカーディナリティが得られます:
(cn ⋈ mc) = 2m, then ⋈ t = 2m
(t ⋈ mc) = 2m, then ⋈ cn = 2m
これらの3つのテーブルをどの順序で結合しても、2番目の結合には常に同じ2m行が渡されます。
ここで、選択的述語を追加しましょう:
cn' = 5k (10万社のうち5%が日本企業であると仮定)
mc = 2m (変化なし)
t' = 200k (100万件のタイトルのうち20%が2000年代に制作されたと仮定)
(cn' ⋈ mc) ≈ 100k, then ⋈ t' ≈ 20k
(t' ⋈ mc) ≈ 400k, then ⋈ cn' ≈ 20k
最初の結合順序は、2mのmovie_companiesエントリを、日本企業である5%のスライスにフィルタリングします。一様分布を仮定すると(なぜこの仮定をするのかは後で議論します)、この結合は約10万行の結果をもたらします。結果をフィルタリングされたタイトルテーブルと結合すると、2000年代のそれらの行の20%のみが保持されます。
2番目の結合順序は、2mのmovie_companiesエントリを、2000年代に制作されたタイトルの20%のスライスにフィルタリングします。同じ一様性の仮定が成り立つとすると、最初の結合は約40万行の結果をもたらし、これは2番目の結合に40万行を渡していることを意味します。2番目の結合順序を選択した場合、4倍の作業を行います。
残念ながら、それだけでは終わりません。
組み合わせ爆発
各結合は、以下のいずれかを使用できます:
Hash join
Merge join
Nested-loop join
可換性を考慮に入れると、各結合には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種類のスキャンのみを考慮すると:
Sequential
Index
Index-only
Bitmap
2つの結合ツリー:どのテーブルのペアが最初に結合されるか。
× 22 の向き:2つの結合のそれぞれが、どちらの入力が外側でどちらが内側かを入れ替えることができます。
× 32 のアルゴリズム:2つの結合のそれぞれが、ハッシュ、マージ、またはネストループを選択します。
× 43 のスキャン:3つのテーブルのそれぞれが、シーケンシャルまたはインデックス、インデックスオンリー、またはビットマップスキャンで読み取られます。
= 4,608
このクエリを実行する4,608通りの方法があります。
これは実際には過小評価です。プランは並列で実行でき、集計はハッシュまたはソートできます。また、Postgresが...