refactor: apply route-dependent rewrites after routing

pgdog/src/backend/pool/connection/aggregate.rs

@@ -6,7 +6,7 @@ use std::mem;
 use crate::{
     frontend::router::parser::{
         Aggregate, AggregateFunction, AggregateTarget,
-        rewrite::statement::aggregate::{AggregateRewritePlan, HelperKind},
+        rewrite::statement::{aggregate::HelperKind, projection::ProjectionRewritePlan},
     },
     net::{
         Decoder,
@@ -245,7 +245,7 @@ impl<'a> Aggregates<'a> {
         rows: &'a VecDeque<DataRow>,
         decoder: &'a Decoder,
         aggregate: &'a Aggregate,
-        plan: &AggregateRewritePlan,
+        plan: &ProjectionRewritePlan,
     ) -> Option<Self> {
         let mut helper_columns: HashMap<usize, HelperColumns> = HashMap::new();
 
@@ -262,7 +262,7 @@ impl<'a> Aggregates<'a> {
             }
         }
 
-        for helper in plan.helpers() {
+        for helper in &plan.aggregate_helpers {
             let Some(index) = decoder.row_description().field_index(&helper.alias) else {
                 continue;
             };
@@ -523,7 +523,7 @@ mod test {
         shard1.add("3");
         rows.push_back(shard1);
 
-        let plan = AggregateRewritePlan::default();
+        let plan = ProjectionRewritePlan::default();
         let mut result = Aggregates::new(&rows, &decoder, &aggregate, &plan)
             .unwrap()
             .aggregate()
@@ -571,7 +571,7 @@ mod test {
                 );
                 rows.push_back(row);
             }
-            let plan = AggregateRewritePlan::default();
+            let plan = ProjectionRewritePlan::default();
             let mut result = Aggregates::new(&rows, &decoder, &aggregate, &plan)
                 .expect("count aggregate")
                 .aggregate()
@@ -609,7 +609,7 @@ mod test {
             &rows,
             &decoder,
             &aggregate,
-            &AggregateRewritePlan::default(),
+            &ProjectionRewritePlan::default(),
         )
         .unwrap()
         .aggregate()
@@ -654,7 +654,7 @@ mod test {
             &rows,
             &decoder,
             &aggregate,
-            &AggregateRewritePlan::default(),
+            &ProjectionRewritePlan::default(),
         )
         .unwrap()
         .aggregate()
@@ -701,7 +701,7 @@ mod test {
             &rows,
             &decoder,
             &aggregate,
-            &AggregateRewritePlan::default(),
+            &ProjectionRewritePlan::default(),
         )
         .unwrap()
         .aggregate()

pgdog/src/backend/pool/connection/buffer.rs

@@ -9,7 +9,7 @@ use std::{
 use crate::{
     frontend::router::parser::{
         Aggregate, DistinctBy, DistinctColumn, Limit, OrderBy,
-        rewrite::statement::aggregate::AggregateRewritePlan,
+        rewrite::statement::projection::ProjectionRewritePlan,
     },
     net::{
         Decoder,
@@ -140,31 +140,30 @@ impl Buffer {
         &mut self,
         aggregate: &Aggregate,
         decoder: &Decoder,
-        plan: &AggregateRewritePlan,
+        plan: &ProjectionRewritePlan,
     ) -> Result<(), super::Error> {
         let buffer: VecDeque<DataRow> = std::mem::take(&mut self.buffer);
-        let mut rows = if aggregate.is_empty() {
+        let rows = if aggregate.is_empty() {
             buffer
         } else if let Some(aggregates) = Aggregates::new(&buffer, decoder, aggregate, plan) {
             aggregates.aggregate()?
         } else {
             buffer
         };
 
-        Self::drop_helper_columns(&mut rows, plan);
         self.buffer = rows;
 
         Ok(())
     }
 
-    fn drop_helper_columns(rows: &mut VecDeque<DataRow>, plan: &AggregateRewritePlan) {
+    pub(super) fn drop_helper_columns(&mut self, plan: &ProjectionRewritePlan, decoder: &Decoder) {
         if plan.is_noop() {
             return;
         }
 
-        let drop = plan.drop_columns().collect();
+        let drop = plan.drop_columns(decoder.row_description());
 
-        for row in rows.iter_mut() {
+        for row in self.buffer.iter_mut() {
             row.drop_columns(&drop);
         }
     }
@@ -285,7 +284,7 @@ mod test {
             buf.add(dr.message()).unwrap();
         }
 
-        buf.aggregate(&agg, &Decoder::from(rd), &AggregateRewritePlan::default())
+        buf.aggregate(&agg, &Decoder::from(rd), &ProjectionRewritePlan::default())
             .unwrap();
         buf.mark_full();
 
@@ -312,7 +311,7 @@ mod test {
             }
         }
 
-        buf.aggregate(&agg, &Decoder::from(rd), &AggregateRewritePlan::default())
+        buf.aggregate(&agg, &Decoder::from(rd), &ProjectionRewritePlan::default())
             .unwrap();
         buf.mark_full();
 

pgdog/src/backend/pool/connection/multi_shard/mod.rs

@@ -259,15 +259,19 @@ impl MultiShard {
             self.buffer.mark_full();
 
             if !self.buffer.is_empty() {
+                // Helpers remain in the internal row through aggregation and
+                // sorting, then are removed before client-visible operations.
                 self.buffer
                     .aggregate(
                         self.route.aggregate(),
                         &self.decoder,
-                        self.route.aggregate_rewrite_plan(),
+                        &self.route.projection_rewrite_plan,
                     )
                     .map_err(Error::from)?;
 
                 self.buffer.sort(self.route.order_by(), &self.decoder);
+                self.buffer
+                    .drop_helper_columns(&self.route.projection_rewrite_plan, &self.decoder);
                 self.buffer.distinct(self.route.distinct(), &self.decoder);
                 self.buffer.limit(self.route.limit());
             }
@@ -310,11 +314,11 @@ impl MultiShard {
         {
             // Only send it to the client once all shards sent it,
             // so we don't get early requests from clients.
-            let plan = self.route.aggregate_rewrite_plan();
+            let plan = &self.route.projection_rewrite_plan;
             if plan.is_noop() {
                 forward = Some(message);
             } else {
-                let client_rd = rd.drop_columns(plan.drop_columns());
+                let client_rd = rd.drop_columns(plan.drop_columns(&rd));
                 forward = Some(client_rd.message());
             }
 

pgdog/src/backend/pool/connection/multi_shard/test.rs

@@ -1,5 +1,8 @@
 use crate::{
-    frontend::router::parser::{DistinctBy, Shard, ShardWithPriority},
+    frontend::router::parser::{
+        DistinctBy, OrderBy, Shard, ShardWithPriority,
+        rewrite::statement::projection::{OrderByHelper, OrderBySource, ProjectionRewritePlan},
+    },
     net::{BindComplete, DataRow, Field, Format},
 };
 
@@ -60,6 +63,80 @@ fn test_inconsistent_data_rows() {
     }
 }
 
+#[test]
+fn test_order_by_helper_after_star_expansion_is_dropped_after_sorting() {
+    let mut plan = ProjectionRewritePlan::default();
+    plan.order_by_helpers.push(OrderByHelper {
+        sort_position: 0,
+        source: OrderBySource::Column("price".into()),
+        alias: "__pgdog_order_col0".into(),
+        injected: true,
+    });
+    let mut route = Route::select(
+        ShardWithPriority::new_default_unset(Shard::All),
+        vec![OrderBy::AscColumn("__pgdog_order_col0".into())],
+        Default::default(),
+        Default::default(),
+        None,
+    );
+    route.projection_rewrite_plan = plan;
+    let mut multi_shard = MultiShard::new(vec![0, 1], &route);
+
+    let row_description = RowDescription::new(&[
+        Field::bigint("id"),
+        Field::text("value"),
+        Field::timestamp("created_at"),
+        Field::bigint("__pgdog_order_col0"),
+    ]);
+    assert!(
+        multi_shard
+            .handle_server_message(row_description.message())
+            .unwrap()
+            .is_none()
+    );
+    let client_description = multi_shard
+        .handle_server_message(row_description.message())
+        .unwrap()
+        .unwrap();
+    let client_description = RowDescription::from_bytes(client_description.to_bytes()).unwrap();
+    assert_eq!(
+        client_description
+            .fields
+            .iter()
+            .map(|field| field.name.as_str())
+            .collect::<Vec<_>>(),
+        ["id", "value", "created_at"]
+    );
+
+    let mut first = DataRow::new();
+    first
+        .add(1_i64)
+        .add("first")
+        .add("2026-01-01 00:00:00")
+        .add(20_i64);
+    let mut second = DataRow::new();
+    second
+        .add(2_i64)
+        .add("second")
+        .add("2026-01-02 00:00:00")
+        .add(10_i64);
+    multi_shard.handle_server_message(first.message()).unwrap();
+    multi_shard.handle_server_message(second.message()).unwrap();
+
+    for _ in 0..2 {
+        multi_shard
+            .handle_server_message(CommandComplete::from_str("SELECT 1").message())
+            .unwrap();
+    }
+
+    for expected in [2_i64, 1_i64] {
+        let message = multi_shard.get_server_message().unwrap();
+        let row = DataRow::from_bytes(message.to_bytes()).unwrap();
+        assert_eq!(row.len(), 3);
+        assert_eq!(row.get::<i64>(0, Format::Text).unwrap(), expected);
+    }
+}
+
 #[test]
 fn test_rd_before_dr() {
     let mut multi_shard = MultiShard::new(