refactor: apply route-dependent rewrites after routing by murex971 · Pull Request #1545 · pgdogdev/pgdog

murex971 · 2026-09-14T17:59:40Z

Aggregate helpers, ORDER BY helpers, and LIMIT/OFFSET pushdown now run only after the route is known, so direct shard execution keeps original SQL and does not leak helper columns.

closes #1425 closes #1135

codecov · 2026-09-14T18:06:28Z

Codecov Report

❌ Patch coverage is 96.96970% with 35 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...end/client/query_engine/test/rewrite_projection.rs	97.03%	12 Missing ⚠️
pgdog/src/backend/prepared_statements.rs	92.38%	8 Missing ⚠️
...ontend/router/parser/rewrite/statement/order_by.rs	95.97%	6 Missing ⚠️
...tend/router/parser/rewrite/statement/projection.rs	97.29%	4 Missing ⚠️
pgdog/src/backend/server.rs	62.50%	3 Missing ⚠️
...og/src/frontend/client/query_engine/route_query.rs	83.33%	1 Missing ⚠️
...g/src/frontend/prepared_statements/global_cache.rs	99.11%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

levkk · 2026-09-14T21:07:16Z

👀

levkk · 2026-09-15T14:24:52Z

+                                    describe.anonymize();
+                                }
+                                if let Some(describe) = describe {
+                                    return Ok(HandleResult::PrependProtocolMessageRewrite {


Q: why do we need this? I'm not too familiar with this piece of code, would be good to have a small example explaining.

levkk

This looks good to me I think. It would be good to get feedback from @jkaczman after we merge #1541 (both touch the rewriter heavily). Should be easy to resolve conflicts, but just want to double check one doesn't break the other and vice versa.

levkk · 2026-09-15T14:27:47Z

        }))
    }

+    fn internal_describe(&mut self, name: &str) -> Option<ProtocolMessage> {


Oh I see. The cross-shard Describe has different columns that direct-to-shard, so we need to rewrite it. I'm remembering now.

jkaczman

LGTM!

It would be great if you could document certain decisions e.g. why you chose to move from an explicit A_Const node with the limit + offset to an A_Expr executing on Postgres. I think I understand the tradeoffs, but it would be great to see the author's POV when reading the code.

I would also be curious to know if the benchmark from before vs after changed at all.

murex971 · 2026-09-17T17:41:58Z

Benchmarked release builds of main and this PR using pgbench prepared protocol against the same three shards. Each case used a 3s warmup followed by three 15s runs at 10 and 50 clients.

Results were effectively neutral: direct and cross-shard AVG ranged from −0.6% to +1.0%; comparable ORDER BY was ~1% slower (expected from post-route bookkeeping). CPU was unchanged (+0.13%), with ~476 KiB higher peak RSS.

jkaczman · 2026-09-18T16:18:13Z

    pub(super) fn select(
        &mut self,
-        cached_ast: &Ast,
+        _cached_ast: &Ast,


I think we're fine to remove this as it's no longer needed in this func

jkaczman · 2026-09-18T16:19:12Z

-        }
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub(crate) enum HelperKind {
+    Count,


Why remove the comments that were originally in aggregate/plan.rs?

jkaczman · 2026-09-18T16:22:22Z

@@ -0,0 +1,199 @@
+use pg_raw_parse::{ConstValue, Node, make, nodes};
+


It would be great to have some comments explaining some of the logic in this file.

jkaczman · 2026-09-18T16:35:00Z

+}
+
+#[derive(Debug, Default, Clone)]
+pub(crate) struct RewriteOutput {


It looks like this struct is just a basic wrapper around ProjectionRewritePlan. Why not just use ProjectionRewritePlan by itself?

jkaczman · 2026-09-18T16:36:39Z

+            )
+    }
+
+    pub(crate) fn aggregate_helpers(&self) -> &[AggregateHelper] {


These fields can be exposed pub(crate) and used directly if we're providing basic getters/setters without any special invariants enforced or anything.

jkaczman · 2026-09-18T16:46:28Z

+                        && projected.fields().into_iter().next_back().and_then(Node::as_str)
+                            == name))
+    })
+}


I think this will fail with something like SELECT p.* FROM products p ORDER BY price. Looks like it would work fine if it was SELECT * FROM products ORDER BY price or SELECT p.* FROM products p ORDER BY p.price, but not if we have a mixed case of one being qualified and the other not.

jkaczman · 2026-09-18T16:52:57Z

+    pub(crate) fn projection_rewrite_plan(&self) -> &ProjectionRewritePlan {
+        &self.projection_rewrite_plan
    }

-    pub(crate) fn set_rewrite_plan(&mut self, plan: AggregateRewritePlan) {
-        self.rewrite_plan = plan;
+    pub(crate) fn set_projection_rewrite_plan(&mut self, plan: ProjectionRewritePlan) {
+        self.projection_rewrite_plan = plan;


This is another case we could probably get rid of the getter/setter

jkaczman · 2026-09-18T17:03:39Z

+fn cross_shard_variant_name(name: &str) -> String {
+    format!("{name}_cross_shard")
+}


I wonder if we could re-use the previous allocation somehow, here's a super simple example (probably not optimal):

buf = "x__pgdog_1"
cross_shard_variant = buf
base = buf[1..]

jkaczman · 2026-09-18T17:22:59Z

+            continue;
+        }
+
+        let projected_column = select.target_list().len() + helpers.len();


I think we shouldn't rely on this column index; it gives us the parse tree entries, but not the actual output columns (considering when * is used). I wonder if it would be better to use what you do in aggregate.rs again (decoder.row_description().field_index(&helper.alias) else {). Right now, I think there are edge cases where this can accidentally delete user columns / sort on the wrong column.

jkaczman · 2026-09-18T17:23:15Z

+
+    #[test]
+    fn projects_column_not_covered_by_qualified_star() {
+        let (sql, plan) = rewrite("SELECT a.* FROM a JOIN b ON a.id = b.a_id ORDER BY b.score");


This should also test the index (going off of my comment above)

murex971 added 2 commits September 14, 2026 22:41

refactor: unify temporary projection rewrites

a01eb6e

refactor: finalize shard rewrites after routing

a6e6532

murex971 mentioned this pull request Sep 14, 2026

fix cross-shard ORDER BY on non-projected columns #1517

Closed

murex971 added 5 commits September 14, 2026 23:45

separate order_by

06b3a38

fix

4507ac5

fmt

ca39461

more fixes

d6075dd

describe croos shard prepared variants

8757c75

murex971 marked this pull request as ready for review September 14, 2026 20:01

murex971 added 2 commits September 15, 2026 15:52

edge cases in order_by

0b80320

fix post route rewrites

8a87626

levkk reviewed Sep 15, 2026

View reviewed changes

Comment thread pgdog/src/frontend/prepared_statements/global_cache.rs Outdated

levkk reviewed Sep 15, 2026

View reviewed changes

Comment thread pgdog/src/frontend/prepared_statements/global_cache.rs Outdated

murex971 requested a review from jkaczman September 16, 2026 13:33

jkaczman reviewed Sep 17, 2026

View reviewed changes

murex971 added 3 commits September 17, 2026 21:37

Merge remote-tracking branch 'forkie/main' into refactor-parse-route

7e9f3b7

add helper

017fc1f

comments

f6550d4

murex971 requested review from jkaczman and levkk September 17, 2026 18:36

jkaczman reviewed Sep 18, 2026

View reviewed changes

		@@ -0,0 +1,199 @@
		use pg_raw_parse::{ConstValue, Node, make, nodes};

Conversation

murex971 commented Sep 14, 2026 • edited by jkaczman Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

codecov Bot commented Sep 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

levkk commented Sep 14, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

levkk left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

jkaczman left a comment

Choose a reason for hiding this comment

Uh oh!

murex971 commented Sep 17, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

murex971 commented Sep 14, 2026 •

edited by jkaczman

Loading

codecov Bot commented Sep 14, 2026 •

edited

Loading