跳到正文
Hacker News (official RSS)(pgrust.com)·· 4 小時前精選品質評分79

pgrust v0.4-beta 推出 Gravity 查詢引擎,速度比 ClickHouse 快 3 倍

Programming at Light Speed: Building the Gravity Query Engine

重點導讀

Gravity 是 pgrust v0.4-beta 內建的查詢引擎,透過 JIT 編譯與 SIMD 最佳化,於 ClickBench 上比 ClickHouse 快 3 倍,同時完全支援 PostgreSQL SQL 方言與型別系統。

推薦理由

Gravity 透過 JIT 編譯與低階最佳化,將 PostgreSQL 分析查詢速度提升 3 倍以上,為資料庫效能提供可落地的實務參考。

正文 · 中文翻譯

隨著 pgrust v0.4-beta 的發布,我們的新查詢引擎 Gravity 現已在 ClickBench 上 比 ClickHouse 快多達 3 倍,這是 ClickHouse 自己設計的基準測試,用於展示 ClickHouse 的速度。更好的是,Gravity 在支援完整的 Postgres SQL 語法與型別系統的同時,也達成了這一速度。

在之前的幾篇文章中,我提到我們如何使用批次處理、運算子融合與 SIMD,讓 pgrust 在分析查詢方面比 Postgres 快數百倍,以及我們如何構建一個 copy‑and‑patch JIT 編譯器,能在 5 微秒內編譯 stencil。這些工作在多次釋出中累積,最終形成了 Gravity。

copy‑and‑patch 是一個不錯的起點,但僅將預編譯的模板拼接在一起只能做到這樣。由於我們沒有真正的中間表示或跨整個迴圈的暫存器分配,與 rustc 或 Clang 等最佳化編譯器相比,我們最終留下了大量潛在效能。對於 pgrust v0.4,我們想徹底做到位,打造一個能為任何 SQL 查詢產生最佳機器碼的查詢編譯器。這就是我們構建 Gravity 的原因。

通常,工程師嘗試提升軟體效能時,會自下而上。先跑基準測試、檢視熱圖,找出熱點,然後想辦法把程式碼提升 15% 或 20%。自下而上的做法問題在於,你不知道上限是多少。你可能距離硬體能執行的最快程式碼只差 20%,也可能差距達 100 倍。

當我們開始構建 Gravity 時,我們把這個問題反轉。問自己:資料庫究竟能有多快?我們不再逐步提升現有查詢引擎,而是從硬體與檔案格式的實體速度極限出發,倒推。這我稱之為「以光速程式設計」。

為了弄清「光速」是什麼樣子,我們讓 Claude 為一批分析查詢直接對我們的列式儲存 crate 撰寫獨立、重度最佳化的 Rust 程式。這些程式完全跳過 SQL 執行器,只在磁碟格式上直接執行手工調校的多執行緒 SIMD Rust 迴圈。當我們執行這些手寫 Rust 程式時,查詢速度大約比現有的分析資料庫(如 ClickHouse)快 4 倍。接著我們從這些手工調校的 Rust 迴圈中挖掘可納入 Gravity 的最佳化。

Gravity 的目標變得簡單:對任何 SQL 查詢,在執行時 JIT 編譯機器碼,達到與 Claude 直接針對我們檔案格式手寫的 Rust 程式相差不超過 20%。

Gravity 的整體架構

在設計 Gravity 架構時,我們的主要目標是為生成的輸出提供極大的彈性。我們知道,如果能輕鬆讓 Gravity 發出任何查詢所需的低階最佳化,便能更容易縮小差距,逼近光速。

我們不再硬編執行模式或拼接固定 stencil,而是將 Gravity 結構為多階段編譯器管線,逐步將高階 Postgres 計畫降至原始組譯碼:

SQL Query
  -> Postgres Plan
  -> Pipeline Graph
  -> Lisp IR
  -> Assembly

為了觀察查詢在編譯過程每個階段的樣子,我們以以下查詢為例:

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

在此情況下,查詢被拆分為三條管線:

  • p0 從 hits 表中取出資料列,過濾 url 包含「google」的列,然後將列插入雜湊表以計算 COUNT(*) 與 MIN(url),按 user_id 與事件時間的分鐘進行分組。
  • 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 ...))                                ; -> b214: row index + 1 -> back to b210

階段 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 的第一部分是後端。我們想證明我們的 IR 與程式碼產生器實際上能與 rustc -O3 匹敵,才在其上構建 SQL 編譯器。當 Lisp IR 與後端運作正常後,我們讓 Claude 撰寫最佳化的 IR 程式,模擬先前手寫的 Rust 程式。

對這些手寫 IR 程式進行基準測試,我們得到一個重大訊息:我們的 Lisp IR 編譯時間不到一毫秒,執行速度與 rustc 編譯的手寫 Rust 程式相差 10% 至 20%。更重要的是,在我們測試的幾乎所有查詢中,我們不需要發明任何新奇的 IR 指令。IR 的表達力足以,將我們簡陋的 SQL 引擎與「光速」之間的差距縮小到僅僅是將 SQL 查詢編譯為 IR。

教導規劃器新演算法

當我們將 Gravity 與 Claude 的手寫 Rust 在我們的查詢語料庫中比較時,發現某些查詢手寫程式仍然快數個量級,不是因為它編譯迴圈的方式,而是因為執行了根本不同的演算法。儘管我們保留了 Postgres 的規劃器,標準 Postgres 計畫節點無法表達光速演算法。我們必須確保我們的計畫能編碼它們。

在我們最佳化的查詢中,我們發現需要加入三種演算法到規劃器。這三種演算法皆特定於列式格式,因此 Postgres 規劃器無法支援它們是合理的。請記住,對於列式表,通常會指定排序鍵,資料通常以將相似值靠近的順序排列。這三種演算法:

  • ZoneWalk - 對於類似 ORDER BY … LIMIT n 的查詢,Postgres 會使用索引或掃描整張表來執行。pgrcolumnar 提供了另一種執行方式。pgrcolumnar 將表切分成 8192 行一組的區塊,稱為 granule。每個 granule 擁有追蹤該批次最小值與最大值的元資料。與掃描整張表並排序以尋找前 n 行不同,我們利用 min/max 限界按順序讀取 granule,並在 n 行確定後即停止。
  • ZoneTopK - 對於形如 GROUP BY ... LIMIT n 的查詢,我們可以使用與 ZoneWalk 非常相似的演算法。與其掃描整個表並對每一列做雜湊以僅返回 n 個群組,不如利用元資料按 GROUP BY 鍵順序遍歷 granule。當我們已經看到 n 個完整群組,其最大可能鍵低於所有剩餘批次的最小鍵時,就可以立即停止掃描。ClickHouse 26.10 也有類似的最佳化,enable_group_by_top_k_dynamic_filtering,同樣需要完整群組。
  • TwoPass - For a query of the shape GROUP BY ... ORDER BY count(*) DESC LIMIT 10, if one of the columns in the group by has high cardinality, it can be faster to execute the query in two passes:
    1. 透過計算每個鍵的行數來估算每個群組的上限:SELECT unique_key, COUNT(*) ... GROUP BY unique_key
    2. 接著回去只抓取上限足以進入前 10 的群組

這之所以可行是因為查詢中最昂貴的部分是第一次掃描。你可以藉由在第一次掃描只抓取一個欄位,然後在第二次掃描抓取其餘欄位,節省大量記憶體頻寬

這三個計畫都被加入為成本計算的計畫到 pgrust 計畫器。於規劃階段,計畫器估算一次遍歷實際需要讀取多少 granule 才能停止,只有在此重寫比標準掃描更便宜時才會選擇它

現在我們的查詢已使用與高速實作相同的演算法,且知道我們的 IR 足夠快速以編譯查詢計畫,最後一步就是將 SQL 編譯成我們的 IR

最佳化 SQL→IR 編譯

對於 SQL→IR 編譯而言,唯一重要的是我們能把 IR 最佳化到什麼程度。由於我們使用與高速 Rust 程式相同的演算法,我們可以將自己的程式碼與 Rust 程式碼做差異比較,找出差異所在。這導致超過 100 個最佳化。以下列出三個,讓你瞭解 Gravity 為何如此快速。

1. 內聯函式呼叫與聚合

在標準 Postgres 中,評估像 extract(minute FROM EventTime) 或更新聚合函式 AVG、MIN 這類 SQL 函式時,會在每一行透過 Postgres 的函式管理器 (fmgr) 呼叫一次函式。每行(甚至每個小批次)呼叫一次函式會完全毀掉效能。

為瞭解決這個問題,我們將 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 的 SQL 型別為 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 不會逐行比較字串。壓縮字典已排序,因此對於兩個來自同一字典的字串,編碼較小的字串實際上是較小的字串。由此,在計算使用同一字典壓縮的行範圍內的最小字串時,實際上大多是整數比較,而非字串比較。

3. 在暫存器中執行 Collapse

即使我們的過濾器和表示式已被內聯,對數以百萬行的多鍵雜湊表進行探測仍然昂貴。

當我們將 Gravity 的輸出與 Claude 手寫的 Rust 針對 GROUP BY … 查詢做對比時,我們注意到一個有趣的現象。在實際的資料表中,像 UserID 或事件時間戳這類欄位的值往往在磁碟上部分聚簇,這意味著我們經常會看到一連串連續的行屬於完全相同的群組。Claude 的 Rust 程式碼利用了這一點,將目前群組的鍵和值計數儲存在本地變數(rustc 將其保留在 CPU 暫存器中)。只要下一行的鍵與前一行相同,它就只在暫存器中遞增計數器,並且只有在連續行結束時才探測雜湊表。在一張 1 億行的資料表上,這將 1 億次雜湊表探測壓縮到 2200 萬。

我們修改了 Gravity,追蹤前一個群組鍵,並僅在看到新的群組鍵時才探測雜湊表。

確保其通用性

在構建 Gravity 時,我們面臨的一個挑戰是 Claude 容易將最佳化過度貼合我們的語料查詢。Claude 經常僅為我們正在基準測試的那個查詢編寫最佳化。若使用者寫成 strpos(URL, 'google') > 0 而非 URL LIKE '%google%',或按包裝在 Postgres DOMAIN 型別中的欄位分組,一個脆弱的模式匹配器會錯過快速路徑,並靜默地回落到慢速程式碼。

為了對抗這一點,我們製作了每一項最佳化的目錄。接著我們寫了一個 fuzz 測試工具,將單一最佳化應用於多種不同查詢形狀進行測試。它會嘗試各種型別(int2、int4、int8、numeric、text、domain)分組,還有更多,以確認最佳化在多個查詢中都能啟動,而不只是單一查詢。Dan Luu 寫了關於 pgrust 需要多少測試,這也發現了 Postgres 本身的實際錯誤。

Gravity 如今的定位

先以手寫 Rust 為基準設計我們的 Lisp IR,然後教導 Gravity 發出內聯、延遲實體化、暫存器最佳化的 IR,讓我們離「光速」目標驚人地接近。在廣泛的分析查詢範圍內,Gravity 的執行時 JIT 在不到一毫秒內編譯查詢,執行速度則在手寫 Rust 程式碼(Claude 直接針對我們的檔案格式編寫)之 10% 至 20% 之內。

當我們將查詢引擎隔離,並在相同的 pgrust 二進位檔上執行相同的列式表格,將 Gravity 啟用後與我們的 Volcano 行執行器比較,Gravity 在冷啟動時快 102 倍,在熱啟動時快 499 倍,且速度可達 ClickHouse 的 3 倍,同時支援 PostgreSQL 的完整 SQL 方言與型別系統。對於從 PostgreSQL 過來的使用者:在 ClickBench 的 c6a.4xlarge 機器上,pgrust v0.4-beta 的綜合得分比 PostgreSQL 17 高 1,000 倍以上,冷啟動時更高達 1,365 倍。

若您想深入瞭解 Lisp IR、TPDE 後端或 lowerer,您可以檢視程式碼並追蹤我們的常規提交,網址為 github.com/malisper/pgrust。

關於我們早期 v0.4-preview ClickBench 條目的說明:我在 Hacker News 上寫了相關經驗。