Pgrust: Gravity Query Engine

How we built Gravity, the query engine in pgrust v0.4-beta, by starting from the speed limit of the hardware and working backward.

With the release of pgrust v0.4-beta, our new query engine, Gravity, is now as much as 3x faster than ClickHouse on ClickBench, the benchmark that ClickHouse themselves designed to show off how fast ClickHouse is. Even better, Gravity achieves that speed while supporting the full Postgres SQL dialect and type system.

In a couple of previous posts, I wrote about how we used batching, operator fusion, and SIMD to make pgrust hundreds of times faster than Postgres for analytical queries, and how we built a copy-and-patch JIT compiler that compiles stencils in 5 microseconds. All of that work, built up over several releases, is what became Gravity.

Copy-and-patch was a great starting point, but gluing pre-compiled templates together only gets us so far. Because we did not have a real intermediate representation or register allocation across the whole loop, we ended up leaving a lot of performance on the table compared to an optimizing compiler like rustc or Clang. For pgrust v0.4, we wanted to go all the way and build a query compiler that generates optimal machine code for any SQL query. That is why we built Gravity.

Normally, when engineers try to make software faster, they work from the bottom up. They run a benchmark, look at a flamegraph, find a hotspot, and figure out how to speed up their code by 15% or 20%. The problem with working from the bottom up is that you don't know what the ceiling is. You might be 20% away from the fastest possible code the hardware can run, or you might be 100x away.

When we started building Gravity, we flipped this question around. We asked ourselves: how fast can a database actually be? Instead of incrementally speeding up our existing query engine, we wanted to start from the physical speed limit of our hardware and file format and work backward. I like to call this “programming at light speed”.

To find out what "light speed" looked like, we had Claude write standalone, heavily optimized Rust programs directly against our columnar storage crates for a corpus of analytical queries. These programs bypassed the entire SQL executor and just ran hand-tuned, multithreaded Rust loops with SIMD directly over our on-disk format. When we ran those handwritten Rust programs, they executed queries roughly 4x faster than existing analytical databases like ClickHouse across the board. We then mined those hand-tuned Rust loops for optimizations we could incorporate into Gravity.

Our goal for Gravity became simple: for any SQL query, JIT compile machine code at runtime that gets within 20% of the handwritten Rust code Claude wrote directly against our file format.

Our main goal when designing Gravity's architecture was to give ourselves a ton of flexibility in the generated output. We knew that if we made it easy to change Gravity to emit whatever low-level optimizations a query needed, it would become much easier for us to close the gap and approach light speed.

Instead of hardcoding execution patterns or stitching fixed stencils together, we structured Gravity as a multi-stage compiler pipeline that progressively lowers a high-level Postgres plan down to raw assembly:

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

To see what a query looks like at each stage of the compilation process, let’s take the following query:

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;

Stage 1: The Postgres Plan

Gravity does not replace Postgres' parser or query planner. When the query arrives, the pgrust planner analyzes the table statistics and produces a standard query plan.

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

Stage 2: The Pipeline Graph

Next, Gravity takes the query plan and converts it into a “pipeline graph”. This is an intermediate representation that expresses how data will flow throughout the query execution. Each pipeline represents how a row is transformed and operated on throughout the query.

Here is a sketch of what the Pipeline graph looks like for our query:

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

In this case, the query is split into three pipelines:

Next, Gravity compiles each pipeline into a single function in our Lisp-based, assembly-like intermediate representation. Because the IR is general, we can fuse the scan, filter, projection, and aggregation inside p0 into whatever loop structure we want. Here’s a sample of the code for 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

Stage 4: Assembly

Finally, we have backends for each target that compile the IR into assembly:

;; ── 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

Starting from the Bottom: Designing and Benchmarking the Lisp IR First

The first part of Gravity we built was the backends. We wanted to prove that our IR and code generator were actually capable of matching rustc -O3 before we built a SQL compiler on top of them. Once we had the Lisp IR and the backends working, we tested them by having Claude write optimized IR programs that mirrored the previous handwritten Rust code.

Benchmarking those hand-written IR programs gave us a huge piece of news: our Lisp IR compiled in under a millisecond and ran within 10% to 20% of the handwritten Rust code compiled by rustc. Even more importantly, across almost every query we tested, we didn't need to invent any fancy new IR instructions. The IR was expressive enough that the entire gap between our naive SQL engine and "light speed" came down to compiling SQL queries to the IR.

When we compared Gravity against Claude's handwritten Rust across our query corpus, we found queries where the handwritten code was still orders of magnitude faster, not because of how it compiled the loop, but because it was running a fundamentally different algorithm. Even though we kept Postgres’ planner, standard Postgres plan nodes could not express the light-speed algorithms. We had to make sure our plans could encode them.

Across the queries we optimized, we found three algorithms we needed to add to the planner. All three algorithms are specific to a columnar format, so it makes sense that the Postgres planner can't support them. Keep in mind that for a columnar table, you usually specify a sort key, so the data is usually in an order that places similar values close to each other. The three algorithms:

This works because the most expensive part of the query is the first pass. You’re able to save a lot of memory bandwidth by fetching only one column on the first pass and then fetching the rest on a second pass

All three of these plans were added as costed plans to the pgrust planner. At plan time, the planner estimates