[jkaczman]

Building off of #1922, this PR now executes and re-writes plain CTEs (specifically, WITH queries that aren't recursive and don't modify rows) as part of the on-going support for cross-shard subqueries (#381) and cross-shard CTEs (#380).

If we receive a CTE that references a sharded table, we now evaluate the CTE separately, and then re-write the CTE within the outer statement to use the rows it returned from the separate evaluation.

E.g., if we have a statement such as

WITH a AS (SELECT id, status FROM subusers) SELECT x.id, y.id FROM a x JOIN a y ON x.status = y.status

If we sent the statement to all of our shards as-is, Shard A would only join the rows of subusers on Shard A, and Shard B would only join the rows of subusers on Shard B.

Now, we first execute the query of the CTE on all shards. Then we re-write the CTE as explicit VALUES with the rows that the shards return:

WITH a (id, status) AS (VALUES (1::bigint, 'VIP'::text), (2, 'regular')) SELECT x.id, y.id FROM a x JOIN a y ON x.status = y.status

Since the statement now doesn't read any tables, it'll be sent round robin and return the correct answer now (as the CTE values are inlined).

Very similar handling to what's done for derived tables in #1922.

Some other cases are handled:

• Statements can have several CTEs, which we execute in order.
• If we get a CTE that only reads earlier CTEs (e.g., b AS (SELECT id FROM a WHERE status = 'VIP')), b is not executed separately; Postgres will run b against the VALUES of a when we send the outer statement (nice perf handling)
• If we get a CTE that reads a table AND an earlier CTE, the CTE is executed with the earlier CTEs VALUES referenced in a WITH clause.

[github-actions]

Metal bench (arm64) passed. All 475 measures within bounds.

suite	throughput	cpu/txn	user cycles/txn	user instr/txn	kernel cycles/txn	kernel instr/txn	ctx/txn
pooler		19.06 (+1.2% ± 1.0%)	18,324 (+0.3% ± 0.5%)	31,998 (-0.0% ± 0.03%)	30,472 (+1.8% ± 1.7%)	51,004 (-0.3% ± 0.5%)
lb		19.78 (+0.8% ± 0.9%)	19,038 (+0.4% ± 0.5%)	32,869 (-0.1% ± 0.04%)	31,207 (+0.4% ± 1.2%)	52,152 (-0.0% ± 0.3%)
sharding-w2		31.85 (-0.1% ± 2.1%)	43,212 (-0.3% ± 1.1%)	97,319 (+0.0% ± 0.01%)	39,253 (-0.7% ± 3.5%)	53,726 (+0.0% ± 0.2%)	1.55 (+0.2% ± 0.1%)
sharding-w8	250,133 (+0.6% ± 1.8%)	28.55 (-0.6% ± 1.2%)	43,959 (+1.3% ± 3.0%)	52,838 (-0.0% ± 0.2%)	34,882 (-2.9% ± 2.4%)	42,396 (+0.0% ± 0.2%)	0.1716 (+0.8% ± 2.0%)

Counters per txn

suite	frontend stall	l1i refill	backend stall
pooler	7,736 (+0.9% ± 1.1%)	2,612 (+0.3% ± 0.6%)
lb	8,174 (+0.5% ± 1.6%)	2,685 (+0.5% ± 0.5%)
sharding-w2	16,951 (-2.1% ± 3.1%)	4,021 (+0.0% ± 1.9%)
sharding-w8	19,259 (+2.7% ± 7.3%)	4,006 (+2.8% ± 4.2%)

Microbench: all 462 measures within bounds.

Hot code per txn

path	base	head	change
pooler	109.7 KB	109.4 KB	-0.2 KB
sharding-w2	166.2 KB	166.9 KB	+0.7 KB

pooler:

crate	base	head	change
pgdog::frontend::client::query_engine	20.0 KB	19.9 KB	-0.1 KB
tokio	18.4 KB	18.4 KB	-0.1 KB
pgdog::backend::pool	17.1 KB	17.4 KB	+0.2 KB
aws-lc (TLS crypto)	10.2 KB	9.4 KB	-0.8 KB
rustls (TLS)	8.4 KB	8.4 KB	+0.1 KB
pgdog::frontend (other)	7.6 KB	7.7 KB	+0.1 KB
pgdog::backend::server	7.2 KB	7.2 KB	+0.0 KB
other	6.0 KB	6.3 KB	+0.3 KB
everything else	14.7 KB	14.7 KB	+0.0 KB

function	head	change
pgdog::frontend::client::query_engine::query::<impl pgdog::frontend::client::query_engine:	5.4 KB	+0.0 KB
pgdog::frontend::client::Client::run::{{closure}}	5.3 KB	+0.0 KB
pgdog::frontend::client::query_engine::QueryEngine::handle::{{closure}}	4.1 KB	-0.1 KB
tokio::runtime::task::raw::poll	3.2 KB	+0.1 KB
pgdog::frontend::client::query_engine::query::<impl pgdog::frontend::client::query_engine:	3.0 KB	+0.0 KB

sharding-w2:

crate	base	head	change
pgdog::frontend::router	26.6 KB	27.2 KB	+0.6 KB
pgdog::frontend::client::query_engine	23.0 KB	23.1 KB	+0.1 KB
tokio	19.2 KB	19.4 KB	+0.1 KB
pgdog::backend::pool	18.1 KB	18.2 KB	+0.1 KB
aws-lc (TLS crypto)	12.6 KB	12.8 KB	+0.2 KB
other	9.9 KB	9.9 KB	-0.1 KB
pgdog::frontend (other)	8.6 KB	8.8 KB	+0.1 KB
rustls (TLS)	8.2 KB	8.3 KB	+0.1 KB
everything else	39.8 KB	39.2 KB	-0.5 KB

function	head	change
pgdog::frontend::client::Client::run::{{closure}}	6.1 KB	+0.0 KB
pgdog::frontend::client::query_engine::query::<impl pgdog::frontend::client::query_engine:	5.4 KB	+0.1 KB
pgdog::frontend::client::query_engine::QueryEngine::handle::{{closure}}	4.2 KB	-0.1 KB
pgdog::frontend::router::parser::query::select::<impl pgdog::frontend::router::parser::que	3.4 KB	+0.1 KB
tokio::runtime::task::raw::poll	3.4 KB	+0.1 KB

lb flamegraph vs 50ea4f9

Open interactive for hover, zoom and search.

Open interactive.

microbench flamegraph vs 50ea4f9

Open interactive for hover, zoom and search.

Open interactive.

pooler flamegraph vs 50ea4f9

Open interactive for hover, zoom and search.

Open interactive.

sharding-w2 flamegraph vs 50ea4f9

Open interactive for hover, zoom and search.

Open interactive.

sharding-w8 flamegraph vs 50ea4f9

Open interactive for hover, zoom and search.

Open interactive.

Logs

lb, microbench, pooler, sharding-w2, sharding-w8, verdict.json

Compared ba27030 with 50ea4f9 on m8g.metal-24xl (spot).

[github-actions]

Metal bench (x86) failed. 1 of 475 measures outside bounds.

suite	throughput	cpu/txn	user cycles/txn	user instr/txn	kernel cycles/txn	kernel instr/txn	ctx/txn
pooler		34.88 (+0.5% ± 0.3%)	23,126 (+1.5% ± 0.8%)	33,854 (-0.0% ± 0.02%)	58,059 (+0.3% ± 0.3%)	38,472 (-0.1% ± 0.1%)
lb		35.8 (+0.2% ± 0.1%)	24,157 (+0.8% ± 0.7%)	34,810 (-0.0% ± 0.03%)	58,867 (-0.0% ± 0.1%)	39,181 (+0.0% ± 0.1%)
sharding-w2		57.09 (+1.1% ± 0.9%)	54,278 (+2.1% ± 1.3%)	101,183 (-0.0% ± 0.01%)	73,440 (+0.3% ± 1.1%)	43,061 (-0.1% ± 0.3%)	1.389 (-0.3% ± 0.3%)
sharding-w8	146,605 (-1.7% ± 2.3%)	48.26 (+0.3% ± 1.6%)	58,089 (+0.1% ± 1.7%)	56,690 (+0.0% ± 0.1%)	64,533 (+0.7% ± 1.8%)	34,199 (+0.0% ± 0.2%)	0.1673 (-1.3% ± 2.3%)

Counters per txn

suite	frontend stall	l1i refill	backend stall
pooler	75,168 (+3.9% ± 0.9%)	3,602 (+0.0% ± 0.10%)
lb	77,756 (+3.6% ± 1.3%)	3,719 (+0.1% ± 0.2%)
sharding-w2	164,234 (+4.3% ± 3.0%)	6,118 (+1.8% ± 0.6%)
sharding-w8	127,285 (+1.4% ± 3.4%)	4,844 (+0.4% ± 0.6%)

benchmark	measure	result	baseline	change	bound
full_simple::order_limit_offset	instructions	205,582	202,822	+1.36% ± 0.00%	+1%

Microbench: 461 of 462 measures within bounds.

Microbench latency: 1 benches changed more than 3%

benchmark	latency	baseline	change	instructions change
extended_cached::delete_by_key	2,211 ns	2,294 ns	-3.62% ± 1.07%	-0.04% ± 0.00%

Hot code per txn

path	base	head	change
pooler	100.6 KB	102.4 KB	+1.9 KB
sharding-w2	155.4 KB	157.5 KB	+2.1 KB

pooler:

crate	base	head	change
pgdog::frontend::client::query_engine	19.1 KB	19.9 KB	+0.8 KB
tokio	17.1 KB	17.0 KB	-0.1 KB
pgdog::backend::pool	16.2 KB	16.9 KB	+0.8 KB
pgdog::frontend (other)	8.8 KB	8.9 KB	+0.2 KB
rustls (TLS)	6.9 KB	6.9 KB	+0.1 KB
aws-lc (TLS crypto)	7.2 KB	6.5 KB	-0.7 KB
pgdog::backend::server	6.1 KB	6.5 KB	+0.4 KB
other	5.7 KB	5.7 KB	+0.0 KB
everything else	13.6 KB	14.0 KB	+0.4 KB

function	head	change
pgdog::frontend::client::Client::run::{{closure}}	6.8 KB	+0.2 KB
pgdog::frontend::client::query_engine::query::<impl pgdog::frontend::client::query_engine:	5.8 KB	+0.2 KB
pgdog::frontend::client::query_engine::QueryEngine::handle::{{closure}}	4.8 KB	+0.4 KB
tokio::runtime::task::raw::poll	3.4 KB	-0.2 KB
pgdog::frontend::client::query_engine::query::<impl pgdog::frontend::client::query_engine:	2.9 KB	-0.1 KB

sharding-w2:

crate	base	head	change
pgdog::frontend::router	24.2 KB	25.9 KB	+1.7 KB
pgdog::frontend::client::query_engine	23.0 KB	23.4 KB	+0.4 KB
tokio	17.8 KB	17.9 KB	+0.2 KB
pgdog::backend::pool	18.0 KB	17.4 KB	-0.6 KB
other	9.9 KB	10.0 KB	+0.1 KB
pgdog::frontend (other)	10.1 KB	9.7 KB	-0.4 KB
aws-lc (TLS crypto)	10.1 KB	9.6 KB	-0.5 KB
pgdog::backend::server	7.1 KB	7.6 KB	+0.4 KB
everything else	35.2 KB	35.9 KB	+0.8 KB

function	head	change
pgdog::frontend::client::Client::run::{{closure}}	7.2 KB	-0.3 KB
pgdog::frontend::client::query_engine::query::<impl pgdog::frontend::client::query_engine:	5.8 KB	-0.1 KB
pgdog::frontend::client::query_engine::QueryEngine::handle::{{closure}}	4.6 KB	-0.1 KB
pgdog::frontend::router::parser::query::select::<impl pgdog::frontend::router::parser::que	4.1 KB	+0.2 KB
tokio::runtime::task::raw::poll	3.8 KB	-0.1 KB

lb flamegraph vs 50ea4f9

Open interactive for hover, zoom and search.

Open interactive.

microbench flamegraph vs 50ea4f9

Open interactive for hover, zoom and search.

Open interactive.

pooler flamegraph vs 50ea4f9

Open interactive for hover, zoom and search.

Open interactive.

sharding-w2 flamegraph vs 50ea4f9

Open interactive for hover, zoom and search.

Open interactive.

sharding-w8 flamegraph vs 50ea4f9

Open interactive for hover, zoom and search.

Open interactive.

Logs

lb, microbench, pooler, sharding-w2, sharding-w8, verdict.json

Compared ba27030 with 50ea4f9 on c7a.metal-48xl (spot).

[codecov]

Codecov Report

❌ Patch coverage is 85.61020% with 79 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...end/client/query_engine/multi_step/subquery/mod.rs	77.50%	45 Missing ⚠️
...nt/query_engine/multi_step/subquery/with_clause.rs	84.44%	14 Missing ⚠️
...nt/query_engine/multi_step/subquery/from_clause.rs	85.18%	8 Missing ⚠️
...r/parser/rewrite/statement/subquery/with_clause.rs	91.54%	6 Missing ⚠️
...t/query_engine/multi_step/subquery/where_clause.rs	90.00%	3 Missing ⚠️
...parser/rewrite/statement/subquery/insert_select.rs	77.77%	2 Missing ⚠️
.../query_engine/multi_step/subquery/insert_select.rs	95.83%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!