光速のプログラミング：Gravityクエリエンジンの構築

光速のプログラミング：Gravityクエリエンジンの構築

pgrust v0.4-betaのリリースにより、当社の新しいクエリエンジン「Gravity」は、ClickHouse自身がClickHouseの速さを示すために設計したベンチマークであるClickBenchにおいて、ClickHouseよりも最大3倍高速になりました。さらに良いことに、Gravityはその速度を達成しながら、Postgres SQL方言と型システム全体を完全にサポートしています。

以前の投稿で、バッチ処理、演算子融合、SIMDを使用してpgrustを分析クエリでPostgresよりも数百倍高速にしたこと、そして5マイクロ秒でステンシルをコンパイルするコピー＆パッチJITコンパイラを構築したことについて書きました。これら数回のリリースで積み上げてきたすべての作業が、Gravityとなりました。

コピー＆パッチは素晴らしい出発点でしたが、事前にコンパイルされたテンプレートを貼り合わせるだけでは限界があります。中間表現やループ全体でのレジスタ割り当てがなかったため、rustcやClangのような最適化コンパイラと比較して、多くのパフォーマンスを犠牲にしていました。pgrust v0.4では、あらゆるSQLクエリに対して最適なマシンコードを生成するクエリコンパイラを構築するという、完全な目標を掲げました。それがGravityを構築した理由です。

通常、エンジニアがソフトウェアを高速化しようとする場合、ボトムアップで作業します。ベンチマークを実行し、フレイムグラフを見て、ホットスポットを見つけ、コードを15％または20％高速化する方法を見つけます。ボトムアップで作業する問題は、天井がわからないことです。ハードウェアが実行できる最速のコードから20％離れている場合もあれば、100倍離れている場合もあります。

Gravityの構築を開始したとき、私たちはこの質問を逆転させました。データベースは実際にどれだけ速くなれるのか？既存のクエリエンジンを段階的に高速化するのではなく、ハードウェアとファイル形式の物理的な速度限界から始めて、逆方向に作業したかったのです。私はこれを「光速のプログラミング」と呼びたいと思います。

「光速」がどのようなものかを知るために、Claudeにスタンドアロンで高度に最適化されたRustプログラムを、分析クエリのコーパスに対して直接、当社のカラム型ストレージクレートに対して書かせました。これらのプログラムは、SQLエグゼキュータ全体をバイパスし、ディスク上のフォーマットに対して直接、手作業でチューニングされたマルチスレッドRustループとSIMDを実行しました。これらの手作業でチューニングされたRustプログラムを実行したところ、既存の分析データベース（ClickHouseなど）と比較して、クエリは全体的に約4倍高速に実行されました。その後、これらの手作業でチューニングされたRustループから、Gravityに組み込むことができる最適化を抽出しました。

Gravityの目標はシンプルになりました。あらゆるSQLクエリに対して、実行時にマシンコードをJITコンパイルし、Claudeが直接書いたRustコードの90%〜80%のパフォーマンスに近づけることです。

Gravityの全体アーキテクチャ

Gravityのアーキテクチャを設計する際の主な目標は、生成される出力に多大な柔軟性を持たせることでした。クエリが必要とする低レベルの最適化をGravityで容易に変更できるようにすれば、ギャップを埋めて光速に近づくことがはるかに容易になるとわかっていました。

実行パターンをハードコーディングしたり、固定ステンシルを貼り合わせたりする代わりに、Gravityをマルチステージコンパイラパイプラインとして構造化し、高レベルのPostgresプランを段階的に生のアセンブリに低下させます。

SQLクエリ
  -> Postgresプラン
  -> パイプライングラフ
  -> Lisp IR
  -> アセンブリ

コンパイルプロセスの各段階でクエリがどのように見えるかを確認するために、次のクエリを使用します。

SELECT UserID, extract(minute FROM EventTime) AS m, MIN(URL), COUNT(*)
FROM hits
WHERE URL LIKE '%google%'
GROUP BY UserID, m
ORDER BY COUNT(*) DESC
LIMIT 10;

ステージ1：Postgresプラン

GravityはPostgresのパーサーやクエリプランナーを置き換えません。クエリが到着すると、pgrustプランナーはテーブル統計を分析し、標準的なクエリプランを生成します。

Limit (cost=... rows=10)
  -> Sort (Sort Key: count(*) DESC)
       -> HashAggregate (Group Key: userid, extract(minute FROM eventtime))
            -> Seq Scan on hits (Filter: (url ~~ '%google%'))

ステージ2：パイプライングラフ

次に、Gravityはクエリプランを「パイプライングラフ」に変換します。これは、クエリ実行中にデータがどのように流れるかを表現する中間表現です。各パイプラインは、行がクエリ全体でどのように変換され、操作されるかを表します。

このクエリのパイプライングラフの概要は次のとおりです。

graph pipelines=3 breakers=2 root=p2

p0 eligible end=breaker(slot=0 HashAgg)
  source range rel=hits am=pgrc
    lanes Attr(UserID):byval8 Attr(EventTime):byval8 Attr(URL):dict4
  filter like(URL, param0)
  project (key0=UserID, key1=extract_minute(EventTime), agg0=min(URL), agg1=count_star())

p1 serial end=breaker(slot=1 TopK)
  source breaker slot=0
  topk k=10 order=agg1:DESC

p2 serial end=sink
  source breaker slot=1
  project (UserID, m, min_url, count) -> wire

この場合、クエリは3つのパイプラインに分割されます。

• p0はhitsテーブルから行を取得し、urlが「google」を含む行をフィルタリングし、user_idとevent_timeの分に基づいてCOUNT(*)とMIN(url)を計算するために行をハッシュテーブルに挿入します。
• p1は結果のハッシュテーブルを取得し、サイズ10のヒープに入れて、COUNT(*)が最も高い上位10グループを計算します。
• p2は結果のヒープを取得し、行を取り出してワイヤーに送信します。

ステージ3：Lisp IR

次に、Gravityは各パイプラインを、Lispベースのアセンブリライクな中間表現の単一関数にコンパイルします。IRは汎用的なため、p0内のスキャン、フィルタリング、射影、集計を、望むループ構造に融合させることができます。以下はp0のコードのサンプルです。

;; ── loop head: one trip per row ─────────────────────────────── (Seq Scan)
(block b210 ((v2574 i64) ...)                   ; v2574 = row index; loop state rides along as block params
  (v2585 (icmp slt v2574 v103))                 ; more rows in this batch?
  (condbr v2585 (b211) (b35 ...)))              ;   no -> fetch the next batch

(block b213 ((v2587 i64))                       ; ── read the row's columns ── (Seq Scan)
  (v2591 (ptr_add v105 v2587 8 0))              ; &EventTime[row]
  (v2592 (ptr_add v109 v2587 8 0))              ; &UserID[row]
  (v2593 (ptr_add v113 v2587 4 0))              ; &URL[row]   (dictionary code)
  (v2612 (loadz i32 v2593 0))                   ; ── WHERE URL LIKE '%google%' ── (Filter)
  (v2616 (load i64 (ptr_add v2605 (ushr v2612 6) 8 0)))
  (v2621 (iand (ushr v2616 (iand v2612 63)) 1)) ;   one bit per distinct URL
  ;; …                                          ;   row fails -> straight to the next row
  )

(block b227 ()                                  ; ── extract(minute FROM EventTime) ── (Projection)
  (v2695 (load i64 v2591 0))
  (v2711 (iadd v2695 211813660800000000))
  (v2715 (umulh (ushr v2711 10) 10746007285463371))
  (v2720 (isub v2711 (imul (ushr v2715 11) 3600000000)))
  (v2726 (ushr (imul (ushr v2720 8) 18764999) 42)) ; minute
  (v2732 (load i64 v2592 0))                    ; UserID
  ;; …
  )

(block b229 ()                                  ; ── GROUP BY UserID, m ── (HashAggregate)
  (v2897 (icmp eq v2576 v2732))                 ;   same UserID as the group in registers?
  (v2898 (icmp eq v2577 v2739))                 ;   same minute?
  (v2899 (iand v2897 v2898))
  (condbr v2901 (b249) (b250)))                 ;   no -> write it to the hash table (b250), start a new group

(block b249 ()                                  ; ── MIN(URL), COUNT(*) ── (HashAggregate)
  (v2903 (loadz i32 v2593 0))                   ;   this row's URL code
  (v2911 (icmp ule v2903 (iand v2581 4294967295)))
  (v2912 (select v2911 v2908 v2581))            ;   MIN: keep the smaller code
  (v2921 (iadd v2584 1))                        ;   COUNT(*) += 1
  (br b228 v2576 v2577 v2578 v2579 v12 v2912 v2913 v2915 v2921))

(block b228 (...)                               ; ── next row ──
  (br b214 ...))

ステージ4：アセンブリ

最後に、各ターゲット用のバックエンドがあり、IRをアセンブリにコンパイルします。

;; ── b210  loop head ─────────────────────────────────────────── (Seq Scan) --
2b3c  cmp   x10, x24                 ; v2585 = row < rows_in_batch
2b40  b.lt  0x2b70                   ; condbr v2585 (b211) (b35)

;; ── b213  read the row's columns ───────────────────────────── (Seq Scan) --
                                     ; v2591/v2592/v2593 (ptr_add) emit no instructions:
                                     ; they fold into the indexed loads below ([base, row, lsl #n])
2b8c  ldr   w6, [x27, x5, lsl #2]    ; v2612 = URL[row]            (loadz i32)

;; ── WHERE URL LIKE '%google%' ──────────────────────────────── (Filter) --
2b90  lsr   x7, x6, #6               ;        code >> 6
2b94  ldr   x8, [x0, x7, lsl #3]     ; v2616 = words[code >> 6]
2b98  and   x0, x6, #0x3f            ;        code & 63
2b9c  lsr   x7, x8, x0               ;        word >> (code & 63)
2ba0  and   x7, x7, #0x1             ; v2621 = matched bit

;; ── b227  extract(minute FROM EventTime) ───────────────────── (Projection) --
2e04  ldr   x1, [x25, x5, lsl #3]    ; v2695 = EventTime[row]
2e38  mov   x7, #0x2000              ;  ┐
2e3c  movk  x7, #0xfa8f, lsl #16     ;  │ constant 211813660800000000
2e40  movk  x7, #0x8366, lsl #32     ;  │
2e44  movk  x7, #0x2f0, lsl #48      ;  ┘
2e48  add   x1, x1, x7               ; v2711
2e4c  lsr   x7, x1, #10              ;        v2711 >> 10
2e50  mov   x8, #0xf94b              ;  ┐
2e54  movk  x8, #0xcb00, lsl #16     ;  │ constant 10746007285463371
2e58  movk  x8, #0x2d6f, lsl #32     ;  │
2e5c  movk  x8, #0x26, lsl #48       ;  ┘
2e60  umulh x7, x7, x8               ; v2715
2e64  lsr   x7, x7, #11              ;        v2715 >> 11   (hours)
2e68  mov   x8, #0xa400              ;  ┐ constant 3600000000
2e6c  movk  x8, #0xd693, lsl #16     ;  ┘
2e70  mul   x7, x7, x8               ;        hours * 3600000000
2e74  sub   x1, x1, x7               ; v2720 = µs into the hour
2e78  lsr   x1, x1, #8               ;        v2720 >> 8
2e7c  mov   x7, #0x54c7              ;  ┐ constant 18764999
2e80  movk  x7, #0x11e, lsl #16      ;  ┘
2e84  mul   x1, x1, x7
2e88  lsr   x1, x1, #42              ; v2726 = minute
2e9c  ldr   x0, [x26, x5, lsl #3]    ; v2732 = UserID[row]

;; ── b229  same group as the run in registers? ──────────────── (HashAggregate) --
2eb0  cmp   x12, x0                  ; v2897 = run.UserID == UserID
2eb4  ccmp  x13, x1, #0x0, eq        ; v2898, v2899 = && run.minute == minute
2eb8  ccmp  x14, x6, #0x0, eq        ; v2901 = && run.flags == flags   (the third key compare, cut from the IR sample)
2ec8  b.ne  0x2efc                   ; condbr v2901 (b249) (b250)

;; ── b249  MIN(URL), COUNT(*) ───────────────────────────────── (HashAggregate) --
2ecc  ldr   w7, [x27, x5, lsl #2]    ; v2903 = URL[row]   (loaded again)
2ed0  ldr   w8, [x19, #264]          ;  ┐ v2908 = (part << 32) | code
2ed4  orr   x0, x7, x8, lsl #32      ;  ┘   (the select's operand; its definition is cut from the IR sample)
2ed8  and   x8, x30, #0xffffffff     ;        v2581 & 0xFFFFFFFF   (run.min lives in x30)
2edc  cmp   x7, x8                   ;  ┐ v2911 = code <= run.min's code
2ee0  cset  x7, ls                   ;  ┘
2ee4  cmp   x7, #0x0                 ;  ┐ v2912 = select(v2911, v2908, v2581)
2ee8  csel  x30, x0, x30, ne         ;  ┘
2ef0  add   x4, x4, #0x1             ; v2921 = COUNT(*) + 1   (count lives in x4)
2ef8  b     0x337c                   ; br b228 …

;; ── b228 → b214  next row ────────────────────────────────────── --
337c  mov   x18, #0x1                ; block argument v12 ("a run is open")
3380  add   x10, x10, #0x1           ; row + 1
3384  b     0x2b3c                   ; back to b210

Lisp IRの設計とベンチマークをまず行う

Gravityの最初の部分としてバックエンドを構築しました。SQLコンパイラをその上に構築する前に、IRとコードジェネレータが実際にrustc -O3に匹敵できることを証明したかったのです。Lisp IRとバックエンドが動作したら、Claudeに以前の手作業Rustコードを模倣した最適化IRプログラムを書かせ、それらをテストしました。

これらの手作業IRプログラムのベンチマークは、大きなニュースをもたらしました。当社のLisp IRは1ミリ秒未満でコンパイルされ、手作業Rustコード（rustcでコンパイル）の10%〜20%以内のパフォーマンスで実行されました。さらに重要なのは、テストしたほとんどすべてのクエリで、派手な新しいIR命令を考案する必要がなかったことです。IRは十分に表現力豊かであったため、単純なSQLエンジンと「光速」との間のギャップ全体が、SQLクエリをIRにコンパイルすることに集約されました。

プランナーに新しいアルゴリズムを教える

GravityをClaudeの手作業Rustと比較した際、手作業コードが、ループのコンパイル方法ではなく、根本的に異なるアルゴリズムを実行していたために、依然として桁違いに高速なクエリが見つかりました。Postgresのプランナーを維持していても、標準的なPostgresプランノードでは光速アルゴリズムを表現できませんでした。当社のプランがそれらをエンコードできるようにする必要がありました。

最適化したクエリ全体で、プランナーに追加する必要がある3つのアルゴリズムが見つかりました。これら3つのアルゴリズムはカラム型フォーマットに特化しているため、Postgresプランナーがそれらをサポートできないのも理にかなっています。カラム型テーブルの場合、通常はソートキーを指定するため、データは類似の値が互いに近くに配置される順序になっていることを覚えておいてください。3つのアルゴリズムは次のとおりです。

• ZoneWalk - ORDER BY … LIMIT nのようなクエリの場合、Postgresはインデックスを使用するか、テーブル全体をスキャンしてこれを実行します。pgrcolumnarは、クエリを実行するための追加の方法を可能にします。pgrustcolumnarは、テーブルをグラニュールと呼ばれる8192行のバッチに分割します。各グラニュールには、そのバッチの最小値と最大値を追跡するメタデータがあります。テーブル全体をスキャンして上位n行を見つけるためにソートする代わりに、最小/最大境界を使用してグラニュールを順序通りに読み取り、n行が解決され次第停止します。
• ZoneTopK - GROUP BY ... LIMIT nの形状を持つクエリの場合、ZoneWalkと非常に似たアルゴリズムを使用できます。すべての行をスキャンしてハッシュ化するだけでnグループを返すのではなく、メタデータを使用してGROUP BYキーで順序付けられたグラニュールをウォークできます。グループの最大キーが残りのすべてのバッチの最小キーよりも小さい、n個の完全なグループを確認したら、スキャンをすぐに停止できます。ClickHouse 26.10には同様の最適化enable_group_by_top_k_dynamic_filteringがありますが、これも完全なグループを必要とします。
• TwoPass - GROUP BY ... ORDER BY count(*) DESC LIMIT 10の形状のクエリの場合、GROUP BYの列のいずれかが高カーディナリティを持つ場合、クエリを2パスで実行する方が高速になる可能性があります。
  1. 各キーあたりの行数をカウントして、各グループのサイズの推定上限を計算します：SELECT unique_key, COUNT(*) ... GROUP BY unique_key
  2. その後、戻って、上位10に入るのに十分な推定上限を持つグループのみを取得します。

これは、クエリの最もコストのかかる部分が最初のパスであるため機能します。最初のパスでは1つの列のみを取得し、2番目のパスで残りを取得することで、多くのメモリ帯域幅を節約できます。

これら3つのプランすべてが、pgrustプランナーにコスト付きプランとして追加されました。プラン時に、プランナーはウォークが実際に読み取る必要があるグラニュールの数を推定し、標準スキャンよりも安価な場合にのみリライトを選択します。

これで、クエリが光速実装と同じアルゴリズムを使用し、IRがクエリプランをコンパイルするのに十分高速であることがわかったので、最後のステップはSQLからIRへのコンパイルでした。

SQL→IRコンパイルの最適化

SQL→IRコンパイルでは、IRをどれだけ最適化できるかがすべてでした。光速Rustプログラムと同じアルゴリズムを使用したため、コードをRustコードと比較して差分を取り、違いがどこから来るかを確認できました。これにより、100以上の最適化が行われました。Gravityを非常に高速にしているものを理解するための例として、3つ挙げます。

1. 関数呼び出しと集計のインライン化

標準的なPostgresでは、extract(minute FROM EventTime)のようなSQL関数の評価や、AVGやMINのような集計の更新は、行ごとにPostgresの関数マネージャー（fmgr）を介した関数呼び出しを行います。1行あたり（または小さなバッチあたり）に1回関数呼び出しを行うと、パフォーマンスが完全に損なわれます。

これを修正するために、Postgresの組み込み関数をすべてティアに分類し、238コア関数と集計については、Lisp IRに直接インライン化しました。COUNT(*)とSUMを含むクエリは、次のようなものになります。

(block b32 ()
  (v313 (load i64 v307 24))         ; load count from aggregate state
  (v314 (iadd v313 v194))           ; count + 1
  ((store v307 24 v314))
  (v315 (load i64 v307 32))         ; load sum from aggregate state
  (v316 (iadd v315 v309))           ; sum + val
  ((store v307 32 v316))
  (br b22 v197))

AVGのパフォーマンスをベンチマークしたところ、インライン化により、行あたりのループが392命令から120命令に削減され、クエリは2.56倍高速になりました。

2. 遅延マテリアライゼーション

カラム型フォーマットでは、URLのようなテキスト列は辞書エンコードされています。テーブルの各部分はユニークな文字列の辞書を格納し、行自体は辞書への整数コードを格納します。

Gravityの初期バージョンでは、SQL型（TEXT、INT4、NUMERIC）のみに基づいて式のIRを生成する方法を決定していました。URLはTEXT型なので、Gravityは辞書コードを文字列にデコードし、各行にコピーしていました。これらのクエリに対するClaudeの手作業Rustを見ると、行ループ内で文字列に触れることはありませんでした。遅延マテリアライゼーションを使用していました。可能な限り辞書コードを操作し、実際の文字列は最後の最後にのみマテリアライズしました。

Gravityを再設計し、すべての値を可能な限り遅延してマテリアライズするようにしました。各行でurl LIKE '%google%'のようなフィルタを計算する際、コンパイラはユニークな辞書エントリごとに最適化されたサブストリングマッチャーを実行し、フィルタに一致する辞書コードのビットマップを作成するステップを生成します。デフォルトでは、v0.4-betaは、そのビットマップをサーバーメモリに、フィルタとその定数、およびテーブル部分をキーとして保持するため、変更されていないデータに対して同じLIKEを持つ後続のクエリは、マッチャーを再度実行せずに再利用します。pgrust.dict_filter_cache = off、またはpgrust.condition_cacheをgranuleまたはoffに設定すると、すべてのクエリが独自のキャッシュを構築します。

MIN(URL)のような集計を計算する際、Gravityは行ごとに文字列を比較しません。圧縮辞書はソートされているため、同じ辞書からエンコードされた2つの文字列では、コードが小さい方が小さい文字列です。このため、同じ辞書で圧縮された行の範囲内で最小文字列を計算する場合、文字列比較ではなく、ほとんど整数比較になります。

3. レジスタでのラン折り畳み

フィルタや式がインライン化されていても、数百万行のマルチキーハッシュテーブルをプローブするのはコストがかかります。

Gravityの出力をClaudeの手作業Rustと比較した際、興味深いことに気づきました。実世界のテーブルでは、UserIDやイベントタイムスタンプのような列の値はディスク上で部分的にクラスタ化されていることが多く、同じグループに属する連続した行の実行が頻繁に見られます。ClaudeのRustコードは、現在のグループのキーをローカル変数（rustcがCPUレジスタに保持）に保持することでこれを活用しました。次の行が前の行と同じキーを持つ限り、レジスタ内のカウンターをインクリメントするだけで、実行が終了するまでハッシュテーブルをプローブしました。1億行のテーブルでは、1億回のハッシュテーブルプローブが2200万回に折り畳まれました。

Gravityを変更して、前のグループキーを追跡し、新しいグループキーが見つかった場合にのみハッシュテーブルをプローブするようにしました。

一般化の保証

Gravityを構築する上で直面した課題の1つは、Claudeが最適化をコーパスクエリに過剰適合させる傾向があったことです。多くの場合、Claudeはベンチマークしているクエリに正確に合わせて最適化を記述していました。ユーザーがURL LIKE '%google%'の代わりにstrpos(URL, 'google') > 0と記述したり、Postgresのドメイン型でラップされた列でグループ化したりすると、壊れやすいパターンマッチャーは高速パスを見逃し、サイレントに低速コードにフォールバックします。

これに対抗するために、すべての最適化のカタログを作成しました。その後、単一の最適化を取得し、さまざまなクエリ形状でテストするファズィーを作成しました。さまざまな型（int2、int4、int8、numeric、text、domain）、GROUP BY、その他多くのものを試して、最適化が単一のクエリだけでなく、多くのクエリで有効になることを確認しました。Dan Luuは、pgrustにどれだけのテストが費やされているかについて書いており、それによりPostgres自体の実際のバグも発見されています。

Gravityの現状

まず手作業Rustに対してLisp IRを設計し、次にGravityにインライン化され、遅延マテリアライズされ、レジスタ最適化されたIRを生成するように教えることで、「光速」ターゲットに驚くほど近づきました。幅広い分析クエリにおいて、Gravityのランタイムは1ミリ秒未満でクエリをJITコンパイルし、Claudeが直接ファイルフォーマットに対して書いた手作業Rustコードの90%〜80%以内のパフォーマンスで実行します。

クエリエンジンを分離し、Gravityが有効な場合とVolcano行エグゼキュータを使用した場合で、全く同じpgrustバイナリを全く同じカラム型テーブルで実行すると、Gravityはコールドで102倍、ホットで499倍高速になり、ClickHouseよりも最大3倍高速に実行されながら、Postgresの完全なSQL方言と型システムをサポートします。Postgresから来た方へ：ClickBenchのc6a.4xlargeマシンでは、pgrust v0.4-betaの総合スコアはPostgreSQL 17よりも1,000倍以上、コールド実行では1,365倍優れています。

Lisp IR、TPDEバックエンド、またはロワーに深く入りたい場合は、コードを確認し、github.com/malisper/pgrustで定期的なコミットをフォローできます。

以前のv0.4-preview ClickBenchエントリに関する注記：Hacker Newsで何が起こったかについて書きました。