fix: omit cross-shard aggregate helpers on single-shard queries

integration/rust/tests/integration/stddev.rs

@@ -1,6 +1,25 @@
 use crate::setup::{admin_sqlx, connections_sqlx};
 use rust_decimal::prelude::*;
-use sqlx::{Executor, Pool, Postgres, Row};
+use sqlx::{Column, Executor, Pool, Postgres, Row};
+
+#[tokio::test]
+async fn test_direct_stddev_hides_cross_shard_helpers() {
+    let conns = connections_sqlx().await;
+
+    setup_schema(&conns, "test_direct_stddev", "int8").await;
+    setup_data(&conns, "test_direct_stddev", TEST_DATA).await;
+
+    let row = conns[1]
+        .fetch_one(
+            "SELECT stddev_pop(value) AS deviation \
+             FROM test_direct_stddev WHERE customer_id = 1",
+        )
+        .await
+        .unwrap();
+
+    assert_eq!(row.columns().len(), 1);
+    assert_eq!(row.columns()[0].name(), "deviation");
+}
 
 #[tokio::test]
 async fn test_variance_numeric() {

pgdog/src/frontend/client/query_engine/test/mod.rs

@@ -28,6 +28,7 @@ mod pipeline_execution;
 pub(crate) mod prelude;
 mod prepared_syntax_error;
 mod replicas;
+mod rewrite_aggregate;
 mod rewrite_extended;
 mod rewrite_insert_split;
 mod rewrite_offset;

pgdog/src/frontend/client/query_engine/test/rewrite_aggregate.rs

@@ -0,0 +1,89 @@
+use crate::frontend::router::parser::Shard;
+
+use super::prelude::*;
+use super::test_sharded_client;
+
+async fn route_and_rewrite_request(messages: Vec<ProtocolMessage>) -> (Shard, String) {
+    let mut client = test_sharded_client();
+    client.client_request = ClientRequest::from(messages);
+
+    let mut engine = QueryEngine::from_client(&client).unwrap();
+    let mut context = QueryEngineContext::new(&mut client);
+
+    let rewrite = engine.parse_and_rewrite(&mut context).unwrap();
+    engine
+        .route_query(&mut context, rewrite.as_ref())
+        .await
+        .unwrap();
+
+    let route = context.client_request.route().shard().clone();
+    let sql = context
+        .client_request
+        .iter()
+        .find_map(|message| match message {
+            ProtocolMessage::Query(query) => Some(query.query().to_owned()),
+            ProtocolMessage::Parse(parse) => Some(parse.query().to_owned()),
+            _ => None,
+        })
+        .unwrap();
+
+    (route, sql)
+}
+
+async fn route_and_rewrite(sql: &str) -> (Shard, String) {
+    route_and_rewrite_request(vec![Query::new(sql).into()]).await
+}
+
+#[tokio::test]
+async fn test_direct_aggregates_do_not_include_cross_shard_helpers() {
+    for function in ["avg", "stddev", "variance"] {
+        let (route, sql) = route_and_rewrite(&format!(
+            "SELECT {function}(region_id) FROM sharded WHERE id = 1"
+        ))
+        .await;
+
+        assert!(matches!(route, Shard::Direct(_)));
+        assert!(
+            !sql.contains("__pgdog_"),
+            "direct {function} query included helper columns: {sql}"
+        );
+    }
+}
+
+#[tokio::test]
+async fn test_cross_shard_aggregate_keeps_helpers() {
+    let (route, sql) = route_and_rewrite("SELECT stddev(region_id) FROM sharded").await;
+
+    assert!(route.is_all());
+    assert!(sql.contains("__pgdog_count_col0"));
+    assert!(sql.contains("__pgdog_sum_col0"));
+    assert!(sql.contains("__pgdog_sumsq_col0"));
+}
+
+#[tokio::test]
+async fn test_direct_extended_aggregate_does_not_include_helpers() {
+    let (route, sql) = route_and_rewrite_request(vec![
+        Parse::new_anonymous("SELECT avg(region_id) FROM sharded WHERE id = $1").into(),
+        Bind::new_params("", &[Parameter::new(b"1")]).into(),
+        Execute::new().into(),
+        Sync.into(),
+    ])
+    .await;
+
+    assert!(matches!(route, Shard::Direct(_)));
+    assert!(!sql.contains("__pgdog_"));
+}
+
+#[tokio::test]
+async fn test_cross_shard_extended_aggregate_keeps_helpers() {
+    let (route, sql) = route_and_rewrite_request(vec![
+        Parse::new_anonymous("SELECT avg(region_id) FROM sharded").into(),
+        Bind::new_params("", &[]).into(),
+        Execute::new().into(),
+        Sync.into(),
+    ])
+    .await;
+
+    assert!(route.is_all());
+    assert!(sql.contains("__pgdog_count_col0"));
+}

pgdog/src/frontend/client/query_engine/test/rewrite_offset.rs

@@ -18,7 +18,7 @@ async fn run_test(messages: Vec<ProtocolMessage>) -> Option<OffsetPlan> {
     let rewrite_result = engine.parse_and_rewrite(&mut context).unwrap();
 
     match rewrite_result {
-        Some(RewriteResult::InPlace { offset }) => offset,
+        Some(RewriteResult::InPlace { offset, .. }) => offset,
         other => panic!("expected InPlace, got {:?}", other),
     }
 }

pgdog/src/frontend/router/parser/rewrite/statement/aggregate/engine.rs

@@ -64,6 +64,27 @@ impl AggregatesRewrite {
         }
     }
 
+    /// Remove helper aggregate columns previously appended by [`Self::rewrite_select`].
+    pub(crate) fn rollback_select<'a>(
+        select: &mut nodes::SelectStmtMut<'a, '_>,
+        mem: make::MemoryToken<'a>,
+        plan: &AggregateRewritePlan,
+    ) {
+        let original_len = select
+            .target_list()
+            .len()
+            .checked_sub(plan.helpers().len())
+            .expect("aggregate helper plan cannot exceed the SELECT target list");
+        let targets = select
+            .target_list()
+            .iter()
+            .take(original_len)
+            .map(|target| mem.make_unique(target))
+            .collect::<Vec<_>>();
+
+        select.target_list_mut().replace(mem.make_list(&targets));
+    }
+
     fn build_sum_of_squares_func<'a>(
         original: &nodes::FuncCall,
         mem: make::MemoryToken<'a>,
@@ -288,4 +309,27 @@ mod tests {
         // Expect original STDDEV plus three helpers.
         assert_eq!(ast.target_list().len(), 4);
     }
+
+    #[test]
+    fn rewrite_engine_rolls_back_helpers() {
+        let sql = "SELECT AVG(price), STDDEV(discount) FROM menu";
+        let ast = pg_raw_parse::parse(sql).unwrap();
+
+        let select = make::owned(|mem| {
+            let Node::SelectStmt(stmt) = ast.stmts().next().unwrap() else {
+                unreachable!("not a select");
+            };
+            let mut stmt = mem.make_unique(stmt);
+            let aggregate = Aggregate::parse(&stmt, &Default::default());
+            let output = AggregatesRewrite::rewrite_select(&mut stmt.as_mut(), mem, &aggregate);
+
+            assert_eq!(stmt.target_list().len(), 6);
+            AggregatesRewrite::rollback_select(&mut stmt.as_mut(), mem, &output.plan);
+            assert_eq!(stmt.target_list().len(), 2);
+            stmt
+        });
+
+        let statement = pg_raw_parse::deparse(&*select).unwrap();
+        assert!(!statement.as_str().contains("__pgdog_"));
+    }
 }

pgdog/src/frontend/router/parser/rewrite/statement/mod.rs

@@ -166,8 +166,27 @@ impl<'a> StatementRewrite<'a> {
             return Err(err);
         }
 
-        if let NodeMut::SelectStmt(mut select) = stmt.stmt_mut() {
+        if matches!(stmt.stmt(), Node::SelectStmt(_)) {
+            let NodeMut::SelectStmt(mut select) = stmt.stmt_mut() else {
+                unreachable!("statement was checked to be SELECT");
+            };
             self.rewrite_aggregates(&mut select, mem, &mut plan, self.db_schema)?;
+            if !plan.aggregates.is_noop() {
+                aggregate::AggregatesRewrite::rollback_select(&mut select, mem, &plan.aggregates);
+            }
+
+            if !plan.aggregates.is_noop() {
+                plan.direct_stmt = Some(pg_raw_parse::deparse(&*stmt)?.as_str().to_owned());
+
+                let NodeMut::SelectStmt(mut select) = stmt.stmt_mut() else {
+                    unreachable!("statement was checked to be SELECT");
+                };
+                self.rewrite_aggregates(&mut select, mem, &mut plan, self.db_schema)?;
+            }
+
+            let NodeMut::SelectStmt(select) = stmt.stmt_mut() else {
+                unreachable!("statement was checked to be SELECT");
+            };
             self.limit_offset(&select, &mut plan);
         }
 

pgdog/src/frontend/router/parser/rewrite/statement/plan.rs

@@ -1,6 +1,6 @@
 use crate::frontend::{ClientRequest, PreparedStatements};
 use crate::net::messages::bind::{Format, Parameter};
-use crate::net::{Bind, Parse, ProtocolMessage, Query};
+use crate::net::{Bind, ProtocolMessage};
 use crate::unique_id::UniqueId;
 
 use super::insert::build_split_requests;
@@ -9,6 +9,21 @@ use super::{
     Error, InsertSplit, PrepareExecute, ShardingKeyUpdate, aggregate::AggregateRewritePlan,
 };
 
+fn apply_statement(request: &mut ClientRequest, stmt: &str) {
+    for message in request.messages.iter_mut() {
+        match message {
+            ProtocolMessage::Query(query) => query.set_query(stmt),
+            ProtocolMessage::Parse(parse) => {
+                parse.set_query(stmt);
+                if !parse.anonymous() {
+                    PreparedStatements::global().write().rewrite(parse);
+                }
+            }
+            _ => {}
+        }
+    }
+}
+
 /// Statement rewrite plan.
 ///
 /// Executed in order of fields in this struct.
@@ -30,6 +45,11 @@ pub(crate) struct RewritePlan {
     /// Rewritten SQL statement.
     pub(crate) stmt: Option<String>,
 
+    /// Rewritten SQL statement without cross-shard aggregate helper columns.
+    ///
+    /// Restored after routing when a query only needs one shard.
+    pub(crate) direct_stmt: Option<String>,
+
     /// Prepared statements to prepend to the client request.
     /// Each tuple contains (name, statement) for ProtocolMessage::Prepare.
     pub(crate) prepare_rewrites: Vec<PrepareExecute>,
@@ -52,7 +72,10 @@ pub(crate) struct RewritePlan {
 
 #[derive(Debug, Clone)]
 pub(crate) enum RewriteResult {
-    InPlace { offset: Option<OffsetPlan> },
+    InPlace {
+        offset: Option<OffsetPlan>,
+        direct_stmt: Option<String>,
+    },
     InsertSplit(Vec<ClientRequest>),
     ShardingKeyUpdate(ShardingKeyUpdate),
 }
@@ -61,8 +84,25 @@ impl RewriteResult {
     pub(crate) fn apply_after_parser(&self, request: &mut ClientRequest) -> Result<(), Error> {
         match self {
             Self::InPlace {
-                offset: Some(offset),
-            } => offset.apply_after_parser(request),
+                offset,
+                direct_stmt,
+            } => {
+                if request.is_executable()
+                    && request
+                        .route
+                        .as_ref()
+                        .is_some_and(|route| !route.is_cross_shard())
+                    && let Some(stmt) = direct_stmt
+                {
+                    apply_statement(request, stmt);
+                }
+
+                if let Some(offset) = offset {
+                    offset.apply_after_parser(request)?;
+                }
+
+                Ok(())
+            }
             _ => Ok(()),
         }
     }
@@ -100,23 +140,6 @@ impl RewritePlan {
         Ok(())
     }
 
-    /// Apply the rewrite plan to a Parse message by updating the SQL.
-    pub(crate) fn apply_parse(&self, parse: &mut Parse) {
-        if let Some(ref stmt) = self.stmt {
-            parse.set_query(stmt);
-            if !parse.anonymous() {
-                PreparedStatements::global().write().rewrite(parse);
-            }
-        }
-    }
-
-    /// Apply the rewrite plan to a Query message by updating the SQL.
-    pub(crate) fn apply_query(&self, query: &mut Query) {
-        if let Some(ref stmt) = self.stmt {
-            query.set_query(stmt);
-        }
-    }
-
     /// Apply the rewrite plan to a ClientRequest.
     pub(crate) fn apply(&self, request: &mut ClientRequest) -> Result<RewriteResult, Error> {
         // Prepend any required Prepare messages for EXECUTE statements.
@@ -136,12 +159,13 @@ impl RewritePlan {
                 });
         }
 
+        if let Some(stmt) = &self.stmt {
+            apply_statement(request, stmt);
+        }
+
         for message in request.messages.iter_mut() {
-            match message {
-                ProtocolMessage::Parse(parse) => self.apply_parse(parse),
-                ProtocolMessage::Query(query) => self.apply_query(query),
-                ProtocolMessage::Bind(bind) => self.apply_bind(bind)?,
-                _ => {}
+            if let ProtocolMessage::Bind(bind) = message {
+                self.apply_bind(bind)?;
             }
         }
 
@@ -163,6 +187,7 @@ impl RewritePlan {
 
         Ok(RewriteResult::InPlace {
             offset: self.offset.clone(),
+            direct_stmt: self.direct_stmt.clone(),
         })
     }
 }