fix cross-shard ORDER BY on non-projected columns

integration/rust/tests/integration/limit.rs

@@ -27,6 +27,37 @@ async fn limit_across_shards() -> Result<(), Box<dyn std::error::Error>> {
     )
     .await?;
 
+    let rows = sharded
+        .fetch_all("SELECT id FROM limit_test ORDER BY value")
+        .await?;
+    assert_eq!(
+        rows.iter()
+            .map(|row| {
+                assert_eq!(row.len(), 1);
+                row.get::<i32, _>("id")
+            })
+            .collect::<Vec<_>>(),
+        vec![1, 2, 3, 4, 5, 1, 2, 3, 4, 5]
+    );
+
+    let rows = sharded
+        .fetch_all("/* pgdog_shard: 0 */ SELECT id FROM limit_test ORDER BY value")
+        .await?;
+    assert_eq!(
+        rows.iter()
+            .map(|row| {
+                assert_eq!(row.len(), 1);
+                row.get::<i32, _>("id")
+            })
+            .collect::<Vec<_>>(),
+        vec![1, 2, 3, 4, 5]
+    );
+
+    let row = sharded
+        .fetch_one("/* pgdog_shard: 0 */ SELECT stddev(value) FROM limit_test")
+        .await?;
+    assert_eq!(row.len(), 1);
+
     // LIMIT 5
     let rows = sharded
         .fetch_all("SELECT value FROM limit_test ORDER BY value LIMIT 5")

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,
@@ -237,37 +237,37 @@ pub(super) struct Aggregates<'a> {
     mappings: HashMap<Grouping, GroupState<'a>>,
     decoder: &'a Decoder,
     aggregate: &'a Aggregate,
-    helper_columns: HashMap<usize, HelperColumns>,
+    projected_columns: HashMap<usize, HelperColumns>,
 }
 
 impl<'a> Aggregates<'a> {
     pub(super) fn new(
         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();
+        let mut projected_columns: HashMap<usize, HelperColumns> = HashMap::new();
 
         for target in aggregate.targets() {
             let key = target.column();
             match target.function() {
                 AggregateFunction::Count => {
-                    helper_columns.entry(key).or_default().count = Some(target.column());
+                    projected_columns.entry(key).or_default().count = Some(target.column());
                 }
                 AggregateFunction::Sum => {
-                    helper_columns.entry(key).or_default().sum = Some(target.column());
+                    projected_columns.entry(key).or_default().sum = Some(target.column());
                 }
                 _ => {}
             }
         }
 
-        for helper in plan.helpers() {
+        for helper in plan.aggregate_helpers() {
             let Some(index) = decoder.row_description().field_index(&helper.alias) else {
                 continue;
             };
 
-            let entry = helper_columns.entry(helper.target_column).or_default();
+            let entry = projected_columns.entry(helper.target_column).or_default();
             match helper.kind {
                 HelperKind::Count => entry.count = Some(index),
                 HelperKind::Sum => entry.sum = Some(index),
@@ -278,14 +278,14 @@ impl<'a> Aggregates<'a> {
         let helpers_present = aggregate.targets().iter().all(|target| {
             let key = target.column();
             match target.function() {
-                AggregateFunction::Avg => helper_columns
+                AggregateFunction::Avg => projected_columns
                     .get(&key)
                     .and_then(|columns| columns.count)
                     .is_some(),
                 AggregateFunction::StddevPop
                 | AggregateFunction::StddevSamp
                 | AggregateFunction::VarPop
-                | AggregateFunction::VarSamp => helper_columns
+                | AggregateFunction::VarSamp => projected_columns
                     .get(&key)
                     .map(|columns| {
                         columns.count.is_some() && columns.sum.is_some() && columns.sumsq.is_some()
@@ -305,7 +305,7 @@ impl<'a> Aggregates<'a> {
                 decoder,
                 mappings: HashMap::new(),
                 aggregate,
-                helper_columns,
+                projected_columns,
             })
         } else {
             None
@@ -333,7 +333,7 @@ impl<'a> Aggregates<'a> {
                 Entry::Occupied(o) => o.into_mut(),
                 Entry::Vacant(v) => {
                     let accumulators =
-                        Accumulator::from_aggregate(self.aggregate, &self.helper_columns)?;
+                        Accumulator::from_aggregate(self.aggregate, &self.projected_columns)?;
 
                     // Gather all col vals corresponding to list of passthrough indices.
                     let mut passthrough = Vec::new();
@@ -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()
@@ -557,7 +557,7 @@ mod test {
             &rows,
             &decoder,
             &aggregate,
-            &AggregateRewritePlan::default(),
+            &ProjectionRewritePlan::default(),
         )
         .unwrap()
         .aggregate()
@@ -602,7 +602,7 @@ mod test {
             &rows,
             &decoder,
             &aggregate,
-            &AggregateRewritePlan::default(),
+            &ProjectionRewritePlan::default(),
         )
         .unwrap()
         .aggregate()
@@ -649,7 +649,7 @@ mod test {
             &rows,
             &decoder,
             &aggregate,
-            &AggregateRewritePlan::default(),
+            &ProjectionRewritePlan::default(),
         )
         .unwrap()
         .aggregate()

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

@@ -7,12 +7,17 @@ use crate::{
             TwoPcPhase,
             two_pc::{TwoPcTransaction, statement::phase_control},
         },
+        router::parser::rewrite::statement::projection::ProjectionRewritePlan,
+    },
+    net::{
+        DataRow, FromBytes, FrontendPid, ProtocolMessage, Query, RowDescription, ToBytes,
+        parameter::Parameters,
     },
-    net::{FrontendPid, ProtocolMessage, Query, parameter::Parameters},
     state::State,
 };
 
 use futures::future::join_all;
+use std::collections::BTreeSet;
 
 use super::*;
 use crate::util::safe_sleep;
@@ -85,15 +90,18 @@ impl Binding {
         }
     }
 
-    pub(super) async fn read(&mut self) -> Result<Message, Error> {
-        match self {
-            Binding::Direct(guard, _) => guard.read().await,
+    pub(super) async fn read(
+        &mut self,
+        projection_rewrite: &ProjectionRewritePlan,
+    ) -> Result<Message, Error> {
+        let mut message = match self {
+            Binding::Direct(guard, _) => guard.read().await?,
 
             Binding::NotConnected => loop {
                 safe_sleep(Duration::MAX).await
             },
 
-            Binding::Admin(backend) => Ok(backend.read().await?),
+            Binding::Admin(backend) => backend.read().await?,
             Binding::MultiShard(shards, state) => {
                 if shards.is_empty() {
                     loop {
@@ -102,37 +110,42 @@ impl Binding {
                 } else {
                     // Loop until we read a message from a shard
                     // or there are no more messages to be read.
-                    loop {
-                        // Return all sorted data rows if any.
-                        if let Some(message) = state.get_server_message() {
-                            return Ok(message);
-                        }
-                        let mut read = false;
-                        for server in shards.iter_mut() {
-                            if !server.has_more_messages() {
-                                continue;
+                    'message: {
+                        loop {
+                            // Return all sorted data rows if any.
+                            if let Some(message) = state.get_server_message() {
+                                break 'message message;
                             }
+                            let mut read = false;
+                            for server in shards.iter_mut() {
+                                if !server.has_more_messages() {
+                                    continue;
+                                }
 
-                            let message = server.read().await?;
+                                let message = server.read().await?;
 
-                            read = true;
-                            if let Some(message) = state.handle_server_message(message)? {
-                                return Ok(message);
+                                read = true;
+                                if let Some(message) = state.handle_server_message(message)? {
+                                    break 'message message;
+                                }
                             }
-                        }
 
-                        if !read {
-                            break;
+                            if !read {
+                                break;
+                            }
                         }
-                    }
 
-                    loop {
-                        state.query_complete();
-                        safe_sleep(Duration::MAX).await;
+                        loop {
+                            state.query_complete();
+                            safe_sleep(Duration::MAX).await;
+                        }
                     }
                 }
             }
-        }
+        };
+
+        drop_projected_columns(&mut message, projection_rewrite)?;
+        Ok(message)
     }
 
     /// Send an entire buffer of messages to the servers(s).
@@ -554,3 +567,60 @@ impl Binding {
         })
     }
 }
+
+fn drop_projected_columns(
+    message: &mut Message,
+    plan: &ProjectionRewritePlan,
+) -> Result<(), Error> {
+    if plan.is_noop() {
+        return Ok(());
+    }
+
+    let drop = plan.drop_columns().collect::<BTreeSet<_>>();
+    let payload = match message.code() {
+        'D' => {
+            let mut row = DataRow::from_bytes(message.to_bytes())?;
+            row.drop_columns(&drop);
+            row.to_bytes()
+        }
+        'T' => RowDescription::from_bytes(message.to_bytes())?
+            .drop_columns(drop)
+            .to_bytes(),
+        _ => return Ok(()),
+    };
+    message.replace_payload(payload);
+    Ok(())
+}
+
+#[cfg(test)]
+mod response_tests {
+    use super::*;
+    use crate::{
+        frontend::router::parser::rewrite::statement::projection::OrderByHelper,
+        net::{Field, Format},
+    };
+
+    #[test]
+    fn hidden_columns_are_removed_at_binding_boundary() {
+        let mut plan = ProjectionRewritePlan::default();
+        plan.add_order_by_helper(OrderByHelper {
+            sort_position: 0,
+            projected_column: 1,
+        });
+
+        let description =
+            RowDescription::new(&[Field::bigint("id"), Field::text("__pgdog_order_by_0")]);
+        let mut description = description.message();
+        drop_projected_columns(&mut description, &plan).unwrap();
+        let description = RowDescription::from_bytes(description.to_bytes()).unwrap();
+        assert_eq!(description.fields.len(), 1);
+
+        let mut row = DataRow::new();
+        row.add(42_i64).add("alice");
+        let mut row = row.message();
+        drop_projected_columns(&mut row, &plan).unwrap();
+        let row = DataRow::from_bytes(row.to_bytes()).unwrap();
+        assert_eq!(row.len(), 1);
+        assert_eq!(row.get::<i64>(0, Format::Text), Some(42));
+    }
+}