From 5363d1e8c4f90bdfc7bb927550c7abb1e1421cc8 Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Fri, 16 Oct 2026 21:53:58 +0530 Subject: [PATCH 01/24] init --- integration/rust/tests/integration/having.rs | 109 ++++++ integration/rust/tests/integration/mod.rs | 1 + pgdog/src/backend/error.rs | 11 +- .../pool/connection/aggregate/error.rs | 6 + pgdog/src/backend/pool/connection/buffer.rs | 72 ++++ .../backend/pool/connection/buffer/having.rs | 273 +++++++++++++++ .../pool/connection/multi_shard/mod.rs | 8 + .../query_engine/test/rewrite_projection.rs | 72 ++++ pgdog/src/frontend/router/parser/having.rs | 73 ++++ pgdog/src/frontend/router/parser/mod.rs | 2 + .../router/parser/rewrite/statement/error.rs | 3 + .../router/parser/rewrite/statement/having.rs | 311 ++++++++++++++++++ .../router/parser/rewrite/statement/mod.rs | 1 + .../parser/rewrite/statement/projection.rs | 162 ++++++++- pgdog/src/frontend/router/parser/route.rs | 5 + 15 files changed, 1097 insertions(+), 12 deletions(-) create mode 100644 integration/rust/tests/integration/having.rs create mode 100644 pgdog/src/backend/pool/connection/buffer/having.rs create mode 100644 pgdog/src/frontend/router/parser/having.rs create mode 100644 pgdog/src/frontend/router/parser/rewrite/statement/having.rs diff --git a/integration/rust/tests/integration/having.rs b/integration/rust/tests/integration/having.rs new file mode 100644 index 000000000..a43d95338 --- /dev/null +++ b/integration/rust/tests/integration/having.rs @@ -0,0 +1,109 @@ +use crate::setup::{admin_sqlx, connections_sqlx}; +use sqlx::{Executor, Row, postgres::PgPool}; + +#[tokio::test] +async fn having_filters_merged_groups() -> Result<(), Box> { + let sharded = connections_sqlx().await.get(1).cloned().unwrap(); + reset_table(&sharded).await?; + + seed( + &sharded, + 0, + "INSERT INTO having_merge_test VALUES (1, 10), (1, 11), (2, 12)", + ) + .await?; + seed( + &sharded, + 1, + "INSERT INTO having_merge_test VALUES (1, 20), (2, 21), (2, 22)", + ) + .await?; + + let rows = sharded + .fetch_all( + "SELECT customer_id, COUNT(*) AS order_count \ + FROM having_merge_test \ + GROUP BY customer_id \ + HAVING COUNT(*) > 2 \ + ORDER BY customer_id", + ) + .await?; + + assert_eq!(rows.len(), 2); + assert_eq!(rows[0].get::("customer_id"), 1); + assert_eq!(rows[0].get::("order_count"), 3); + assert_eq!(rows[1].get::("customer_id"), 2); + assert_eq!(rows[1].get::("order_count"), 3); + + cleanup(&sharded).await; + Ok(()) +} + +#[tokio::test] +async fn prepared_having_parameter_is_typed_by_postgres() -> Result<(), Box> +{ + let sharded = connections_sqlx().await.get(1).cloned().unwrap(); + reset_table(&sharded).await?; + + seed( + &sharded, + 0, + "INSERT INTO having_merge_test VALUES (1, 10), (1, 11)", + ) + .await?; + seed( + &sharded, + 1, + "INSERT INTO having_merge_test VALUES (1, 20), (2, 21)", + ) + .await?; + + let rows = sqlx::query( + "SELECT customer_id, COUNT(*) AS order_count \ + FROM having_merge_test \ + GROUP BY customer_id \ + HAVING COUNT(*) > $1", + ) + .bind(2_i64) + .fetch_all(&sharded) + .await?; + + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].get::("customer_id"), 1); + assert_eq!(rows[0].get::("order_count"), 3); + + cleanup(&sharded).await; + Ok(()) +} + +async fn reset_table(pool: &PgPool) -> Result<(), sqlx::Error> { + cleanup(pool).await; + for shard in [0, 1] { + pool.execute( + format!( + "/* pgdog_shard: {shard} */ \ + CREATE TABLE having_merge_test(customer_id BIGINT, order_id BIGINT)" + ) + .as_str(), + ) + .await?; + } + admin_sqlx().await.execute("RELOAD").await?; + Ok(()) +} + +async fn seed(pool: &PgPool, shard: usize, statement: &str) -> Result<(), sqlx::Error> { + pool.execute(format!("/* pgdog_shard: {shard} */ {statement}").as_str()) + .await?; + Ok(()) +} + +async fn cleanup(pool: &PgPool) { + for shard in [0, 1] { + pool.execute( + format!("/* pgdog_shard: {shard} */ DROP TABLE IF EXISTS having_merge_test").as_str(), + ) + .await + .ok(); + } +} diff --git a/integration/rust/tests/integration/mod.rs b/integration/rust/tests/integration/mod.rs index 0b3e50b7a..eda72387c 100644 --- a/integration/rust/tests/integration/mod.rs +++ b/integration/rust/tests/integration/mod.rs @@ -18,6 +18,7 @@ mod cross_shard_omni_check; pub mod distinct; pub mod explain; pub mod fake_transactions; +pub mod having; pub mod idle_in_transaction; pub mod limit; pub mod maintenance_mode; diff --git a/pgdog/src/backend/error.rs b/pgdog/src/backend/error.rs index 5696f3051..c3106a720 100644 --- a/pgdog/src/backend/error.rs +++ b/pgdog/src/backend/error.rs @@ -146,10 +146,13 @@ impl Error { pub(crate) fn is_unsupported(&self) -> bool { matches!( self, - Self::Aggregate(AggregateError::UnsupportedAggregation { .. }) - | Self::MultiShard(MultiShardError::Aggregate( - AggregateError::UnsupportedAggregation { .. } - )) + Self::Aggregate( + AggregateError::UnsupportedAggregation { .. } + | AggregateError::UnsupportedHaving(_) + ) | Self::MultiShard(MultiShardError::Aggregate( + AggregateError::UnsupportedAggregation { .. } + | AggregateError::UnsupportedHaving(_) + )) ) } diff --git a/pgdog/src/backend/pool/connection/aggregate/error.rs b/pgdog/src/backend/pool/connection/aggregate/error.rs index 4015d890e..b45d8578b 100644 --- a/pgdog/src/backend/pool/connection/aggregate/error.rs +++ b/pgdog/src/backend/pool/connection/aggregate/error.rs @@ -14,6 +14,12 @@ pub(crate) enum Error { #[error("unsupported aggregation {function}: {reason}")] UnsupportedAggregation { function: String, reason: String }, + #[error("unsupported cross-shard HAVING: {0}")] + UnsupportedHaving(String), + + #[error("invalid cross-shard HAVING state: {0}")] + InvalidHaving(String), + #[error("unexpected aggregation error: {0}")] Unexpected(String), } diff --git a/pgdog/src/backend/pool/connection/buffer.rs b/pgdog/src/backend/pool/connection/buffer.rs index 786063fbf..278f54507 100644 --- a/pgdog/src/backend/pool/connection/buffer.rs +++ b/pgdog/src/backend/pool/connection/buffer.rs @@ -21,6 +21,8 @@ use pgdog_postgres_types::Datum; use super::Aggregates; +mod having; + /// Sort and aggregate rows received from multiple shards. #[derive(Default, Debug, Clone)] pub(super) struct Buffer { @@ -156,6 +158,26 @@ impl Buffer { Ok(()) } + /// Apply `HAVING` after all partial aggregate rows have been merged. + pub(super) fn having( + &mut self, + plan: &ProjectionRewritePlan, + decoder: &Decoder, + ) -> Result<(), super::Error> { + let Some(expression) = plan.having.expression() else { + return Ok(()); + }; + + let mut filtered = VecDeque::with_capacity(self.buffer.len()); + for row in self.buffer.drain(..) { + if having::evaluate(expression, &row, decoder)? { + filtered.push_back(row); + } + } + self.buffer = filtered; + Ok(()) + } + pub(super) fn drop_helper_columns(&mut self, plan: &ProjectionRewritePlan, decoder: &Decoder) { if plan.is_noop() { return; @@ -247,8 +269,10 @@ impl Buffer { #[cfg(test)] mod test { use super::*; + use crate::frontend::router::parser::{CompareOperator, HavingExpr, HavingValue}; use crate::net::{Datum, Field, Format, RowDescription}; use bytes::Bytes; + use pg_raw_parse::Node; #[test] fn test_sort_buffer() { @@ -332,6 +356,54 @@ mod test { } } + #[test] + fn having_filters_after_grouped_aggregate_merge() { + let ast = pg_raw_parse::parse( + "SELECT customer_id, COUNT(*), \ + COUNT(*) AS __pgdog_having_0, 3::bigint AS __pgdog_having_1 \ + FROM purchases GROUP BY customer_id", + ) + .unwrap(); + let Node::SelectStmt(select) = ast.stmts().next().unwrap() else { + panic!("not a select") + }; + let aggregate = Aggregate::parse(select, &Default::default()); + let decoder = Decoder::from(RowDescription::new(&[ + Field::bigint("customer_id"), + Field::bigint("count"), + Field::bigint("__pgdog_having_0"), + Field::bigint("__pgdog_having_1"), + ])); + let plan = ProjectionRewritePlan { + having: crate::frontend::router::parser::HavingRewritePlan::new( + HavingExpr::Compare { + left: HavingValue::new("__pgdog_having_0".into()), + operator: CompareOperator::Greater, + right: HavingValue::new("__pgdog_having_1".into()), + }, + vec!["__pgdog_having_0".into(), "__pgdog_having_1".into()], + ), + ..Default::default() + }; + + let mut buffer = Buffer::default(); + for (customer, partial_count) in [(1_i64, 2_i64), (1, 2), (2, 1), (2, 1)] { + let mut row = DataRow::new(); + row.add(customer) + .add(partial_count) + .add(partial_count) + .add(3_i64); + buffer.add(row.message()).unwrap(); + } + + buffer.aggregate(&aggregate, &decoder, &plan).unwrap(); + buffer.having(&plan, &decoder).unwrap(); + + assert_eq!(buffer.len(), 1); + assert_eq!(buffer.buffer[0].get::(0, Format::Text).unwrap(), 1); + assert_eq!(buffer.buffer[0].get::(1, Format::Text).unwrap(), 4); + } + #[test] fn test_sort_buffer_with_timestamps() { let mut buf = Buffer::default(); diff --git a/pgdog/src/backend/pool/connection/buffer/having.rs b/pgdog/src/backend/pool/connection/buffer/having.rs new file mode 100644 index 000000000..dbd6d0d81 --- /dev/null +++ b/pgdog/src/backend/pool/connection/buffer/having.rs @@ -0,0 +1,273 @@ +use std::cmp::Ordering; + +use pgdog_postgres_types::{Datum, Numeric}; + +use crate::{ + frontend::router::parser::{CompareOperator, HavingExpr, HavingValue}, + net::{Decoder, messages::DataRow}, +}; + +use super::super::aggregate::Error; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum TruthValue { + True, + False, + Unknown, +} + +impl TruthValue { + fn not(self) -> Self { + match self { + Self::True => Self::False, + Self::False => Self::True, + Self::Unknown => Self::Unknown, + } + } +} + +pub(super) fn evaluate( + expression: &HavingExpr, + row: &DataRow, + decoder: &Decoder, +) -> Result { + Ok(matches!( + evaluate_expression(expression, row, decoder)?, + TruthValue::True + )) +} + +fn evaluate_expression( + expression: &HavingExpr, + row: &DataRow, + decoder: &Decoder, +) -> Result { + match expression { + HavingExpr::Compare { + left, + operator, + right, + } => { + let left = value(left, row, decoder)?; + let right = value(right, row, decoder)?; + compare(&left, *operator, &right) + } + HavingExpr::And(expressions) => { + let mut result = TruthValue::True; + for expression in expressions { + match evaluate_expression(expression, row, decoder)? { + TruthValue::False => return Ok(TruthValue::False), + TruthValue::Unknown => result = TruthValue::Unknown, + TruthValue::True => {} + } + } + Ok(result) + } + HavingExpr::Or(expressions) => { + let mut result = TruthValue::False; + for expression in expressions { + match evaluate_expression(expression, row, decoder)? { + TruthValue::True => return Ok(TruthValue::True), + TruthValue::Unknown => result = TruthValue::Unknown, + TruthValue::False => {} + } + } + Ok(result) + } + HavingExpr::Not(expression) => Ok(evaluate_expression(expression, row, decoder)?.not()), + HavingExpr::IsNull { + value: operand, + negated, + } => { + let is_null = value(operand, row, decoder)?.is_null(); + Ok(if is_null != *negated { + TruthValue::True + } else { + TruthValue::False + }) + } + } +} + +fn value(operand: &HavingValue, row: &DataRow, decoder: &Decoder) -> Result { + let index = decoder + .row_description() + .field_index(operand.alias()) + .ok_or_else(|| Error::InvalidHaving(format!("missing operand {}", operand.alias())))?; + Ok(row.get_column_checked(index, decoder)?.value) +} + +fn compare(left: &Datum, operator: CompareOperator, right: &Datum) -> Result { + if left.is_null() || right.is_null() { + return Ok(TruthValue::Unknown); + } + + let ordering = numeric_ordering(left, right) + .or_else(|| comparable_ordering(left, right)) + .ok_or_else(|| { + Error::UnsupportedHaving(format!( + "cannot compare {} with {}", + left.data_type(), + right.data_type() + )) + })?; + let result = match operator { + CompareOperator::Equal => ordering == Ordering::Equal, + CompareOperator::NotEqual => ordering != Ordering::Equal, + CompareOperator::Less => ordering == Ordering::Less, + CompareOperator::LessOrEqual => ordering != Ordering::Greater, + CompareOperator::Greater => ordering == Ordering::Greater, + CompareOperator::GreaterOrEqual => ordering != Ordering::Less, + }; + Ok(if result { + TruthValue::True + } else { + TruthValue::False + }) +} + +fn comparable_ordering(left: &Datum, right: &Datum) -> Option { + match (left, right) { + (Datum::Text(_), Datum::Text(_)) + | (Datum::Interval(_), Datum::Interval(_)) + | (Datum::Timestamp(_), Datum::Timestamp(_)) + | (Datum::TimestampTz(_), Datum::TimestampTz(_)) + | (Datum::Uuid(_), Datum::Uuid(_)) + | (Datum::Boolean(_), Datum::Boolean(_)) + | (Datum::Oid(_), Datum::Oid(_)) => left.partial_cmp(right), + _ => None, + } +} + +fn numeric_ordering(left: &Datum, right: &Datum) -> Option { + if let (Some(left), Some(right)) = (integer(left), integer(right)) { + return Some(left.cmp(&right)); + } + + if let (Some(left), Some(right)) = (exact_numeric(left), exact_numeric(right)) { + return left.partial_cmp(&right); + } + + if let (Some(left), Some(right)) = (float(left), float(right)) { + return Some(postgres_float_ordering(left, right)); + } + + None +} + +fn postgres_float_ordering(left: f64, right: f64) -> Ordering { + match (left.is_nan(), right.is_nan()) { + (true, true) => Ordering::Equal, + (true, false) => Ordering::Greater, + (false, true) => Ordering::Less, + (false, false) => left + .partial_cmp(&right) + .expect("non-NaN floats always have an ordering"), + } +} + +fn integer(value: &Datum) -> Option { + match value { + Datum::Bigint(value) => Some((*value).into()), + Datum::Integer(value) => Some((*value).into()), + Datum::SmallInt(value) => Some((*value).into()), + _ => None, + } +} + +fn exact_numeric(value: &Datum) -> Option { + match value { + Datum::Numeric(value) => Some(*value), + Datum::Bigint(value) => Some(Numeric::from(*value)), + Datum::Integer(value) => Some(Numeric::from(*value)), + Datum::SmallInt(value) => Some(Numeric::from(i32::from(*value))), + _ => None, + } +} + +fn float(value: &Datum) -> Option { + match value { + Datum::Double(value) => Some(value.0), + Datum::Float(value) => Some(value.0.into()), + Datum::Bigint(value) => Some(*value as f64), + Datum::Integer(value) => Some((*value).into()), + Datum::SmallInt(value) => Some((*value).into()), + _ => None, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ + frontend::router::parser::HavingValue, + net::{ + Field, RowDescription, + messages::{DataRow, Datum}, + }, + }; + + fn compare_expression(operator: CompareOperator) -> HavingExpr { + HavingExpr::Compare { + left: HavingValue::new("__pgdog_having_0".into()), + operator, + right: HavingValue::new("__pgdog_having_1".into()), + } + } + + fn integer_field(name: &str) -> Field { + Field { + name: name.into(), + table_oid: 0, + column: 0, + type_oid: 23, + type_size: 4, + type_modifier: -1, + format: 0, + } + } + + #[test] + fn compares_merged_bigint_with_integer_without_losing_precision() { + let decoder = Decoder::from(RowDescription::new(&[ + Field::bigint("__pgdog_having_0"), + integer_field("__pgdog_having_1"), + ])); + let mut row = DataRow::new(); + row.add(9_007_199_254_740_993_i64).add(2_i32); + + assert!( + evaluate( + &compare_expression(CompareOperator::Greater), + &row, + &decoder + ) + .unwrap() + ); + } + + #[test] + fn implements_sql_three_valued_boolean_logic() { + let decoder = Decoder::from(RowDescription::new(&[ + Field::bigint("__pgdog_having_0"), + integer_field("__pgdog_having_1"), + ])); + let mut row = DataRow::new(); + row.add(Datum::Null).add(2_i32); + + let unknown = compare_expression(CompareOperator::Greater); + let is_null = HavingExpr::IsNull { + value: HavingValue::new("__pgdog_having_0".into()), + negated: false, + }; + assert!( + evaluate( + &HavingExpr::Or(vec![unknown.clone(), is_null]), + &row, + &decoder + ) + .unwrap() + ); + assert!(!evaluate(&HavingExpr::Not(Box::new(unknown)), &row, &decoder).unwrap()); + } +} diff --git a/pgdog/src/backend/pool/connection/multi_shard/mod.rs b/pgdog/src/backend/pool/connection/multi_shard/mod.rs index 81de53db3..adcdac68b 100644 --- a/pgdog/src/backend/pool/connection/multi_shard/mod.rs +++ b/pgdog/src/backend/pool/connection/multi_shard/mod.rs @@ -256,6 +256,14 @@ impl MultiShard { return Err(err.into()); } + if let Err(err) = self + .buffer + .having(&self.route.projection_rewrite_plan, &self.decoder) + { + self.buffer.clear(); + return Err(err.into()); + } + self.buffer.sort(self.route.order_by(), &self.decoder); self.buffer .drop_helper_columns(&self.route.projection_rewrite_plan, &self.decoder); diff --git a/pgdog/src/frontend/client/query_engine/test/rewrite_projection.rs b/pgdog/src/frontend/client/query_engine/test/rewrite_projection.rs index d08c39dd0..88fa18741 100644 --- a/pgdog/src/frontend/client/query_engine/test/rewrite_projection.rs +++ b/pgdog/src/frontend/client/query_engine/test/rewrite_projection.rs @@ -95,6 +95,78 @@ async fn cross_shard_aggregate_adds_and_tracks_helpers() { ); } +#[tokio::test] +async fn cross_shard_having_rewrites_and_replaces_route_aggregate() { + let sql = "SELECT customer_id, COUNT(*) AS orders \ + FROM purchases GROUP BY customer_id HAVING COUNT(*) > 2"; + let mut client = test_sharded_client(); + client.client_request = ClientRequest::from(vec![ProtocolMessage::Query(Query::new(sql))]); + + let engine = QueryEngine::from_client(&client).unwrap(); + let (mut context, client_request) = QueryEngineContext::new(&mut client); + let result = engine + .parse_and_rewrite(&mut context, client_request) + .await + .unwrap(); + client_request.route = Some(route(Shard::All)); + + projection::finalize_after_route( + client_request, + &Schema::default(), + result.as_ref().and_then(RewriteResult::offset_plan), + ) + .unwrap(); + + let query = match &client_request.messages[0] { + ProtocolMessage::Query(query) => query, + _ => panic!("expected Query"), + }; + assert!(query.query().contains("__pgdog_having_0")); + assert_eq!( + client_request.route().aggregate().targets().len(), + 2, + "both the selected and hidden COUNT must be merged" + ); + assert!( + client_request + .route() + .projection_rewrite_plan + .having + .expression() + .is_some() + ); +} + +#[tokio::test] +async fn direct_having_keeps_native_postgres_semantics() { + let sql = "SELECT customer_id, COUNT(*) AS orders \ + FROM purchases GROUP BY customer_id HAVING COUNT(*) > 2"; + let mut client = test_sharded_client(); + client.client_request = ClientRequest::from(vec![ProtocolMessage::Query(Query::new(sql))]); + + let engine = QueryEngine::from_client(&client).unwrap(); + let (mut context, client_request) = QueryEngineContext::new(&mut client); + let result = engine + .parse_and_rewrite(&mut context, client_request) + .await + .unwrap(); + client_request.route = Some(route(Shard::Direct(0))); + + projection::finalize_after_route( + client_request, + &Schema::default(), + result.as_ref().and_then(RewriteResult::offset_plan), + ) + .unwrap(); + + let query = match &client_request.messages[0] { + ProtocolMessage::Query(query) => query, + _ => panic!("expected Query"), + }; + assert_eq!(query.query(), sql); + assert!(client_request.route().projection_rewrite_plan.is_noop()); +} + #[tokio::test] async fn named_prepared_aggregate_uses_cross_shard_variant() { let mut client = test_sharded_client(); diff --git a/pgdog/src/frontend/router/parser/having.rs b/pgdog/src/frontend/router/parser/having.rs new file mode 100644 index 000000000..ca83d0b59 --- /dev/null +++ b/pgdog/src/frontend/router/parser/having.rs @@ -0,0 +1,73 @@ +//! Cross-shard `HAVING` execution plan. + +/// Comparison operators supported by the cross-shard `HAVING` evaluator. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum CompareOperator { + Equal, + NotEqual, + Less, + LessOrEqual, + Greater, + GreaterOrEqual, +} + +/// A value projected by the cross-shard query under an internal alias. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct HavingValue { + alias: String, +} + +impl HavingValue { + pub(crate) fn new(alias: String) -> Self { + Self { alias } + } + + pub(crate) fn alias(&self) -> &str { + &self.alias + } +} + +/// Predicate evaluated after partial aggregate results have been merged. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum HavingExpr { + Compare { + left: HavingValue, + operator: CompareOperator, + right: HavingValue, + }, + And(Vec), + Or(Vec), + Not(Box), + IsNull { + value: HavingValue, + negated: bool, + }, +} + +/// Metadata produced while rewriting a cross-shard `HAVING` clause. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub(crate) struct HavingRewritePlan { + expression: Option, + hidden_aliases: Vec, +} + +impl HavingRewritePlan { + pub(crate) fn new(expression: HavingExpr, hidden_aliases: Vec) -> Self { + Self { + expression: Some(expression), + hidden_aliases, + } + } + + pub(crate) fn expression(&self) -> Option<&HavingExpr> { + self.expression.as_ref() + } + + pub(crate) fn hidden_aliases(&self) -> &[String] { + &self.hidden_aliases + } + + pub(crate) fn is_noop(&self) -> bool { + self.expression.is_none() + } +} diff --git a/pgdog/src/frontend/router/parser/mod.rs b/pgdog/src/frontend/router/parser/mod.rs index a74d7a01f..4eb21227a 100644 --- a/pgdog/src/frontend/router/parser/mod.rs +++ b/pgdog/src/frontend/router/parser/mod.rs @@ -16,6 +16,7 @@ pub(crate) mod error; pub(crate) mod explain_trace; mod from_clause; pub(crate) mod function; +mod having; pub(crate) mod key; mod limit; pub(crate) mod multi_tenant; @@ -45,6 +46,7 @@ pub(crate) use distinct::{Distinct, DistinctBy, DistinctColumn}; pub(crate) use error::Error; pub(crate) use from_clause::FromClause; use function::Function; +pub(crate) use having::{CompareOperator, HavingExpr, HavingRewritePlan, HavingValue}; pub(crate) use key::Key; pub(crate) use limit::{Limit, LimitClause}; pub(crate) use order_by::OrderBy; diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/error.rs b/pgdog/src/frontend/router/parser/rewrite/statement/error.rs index 6f45d07cb..f13f3808a 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/error.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/error.rs @@ -14,6 +14,9 @@ pub(crate) enum Error { #[error("sharding key assignment unsupported: {0}")] UnsupportedShardingKeyUpdate(String), + #[error("cross-shard HAVING is unsupported: {0}")] + UnsupportedHaving(String), + #[error("net: {0}")] Net(#[from] crate::net::Error), diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/having.rs b/pgdog/src/frontend/router/parser/rewrite/statement/having.rs new file mode 100644 index 000000000..4973d9b5e --- /dev/null +++ b/pgdog/src/frontend/router/parser/rewrite/statement/having.rs @@ -0,0 +1,311 @@ +//! Route-dependent rewrite for cross-shard `HAVING`. + +use std::collections::HashSet; + +use itertools::Itertools; +use pg_raw_parse::{ConstValue, Node, make, nodes}; + +use crate::frontend::router::parser::{ + CompareOperator, Function, HavingExpr, HavingRewritePlan, HavingValue, +}; + +use super::Error; + +const ALIAS_PREFIX: &str = "__pgdog_having_"; + +struct ParsedHaving<'a> { + expression: HavingTemplate, + values: Vec<(String, Node<'a>)>, +} + +enum HavingTemplate { + Compare { + left: ValueTemplate, + operator: CompareOperator, + right: ValueTemplate, + }, + And(Vec), + Or(Vec), + Not(Box), + IsNull { + value: ValueTemplate, + negated: bool, + }, +} + +struct ValueTemplate { + alias: String, +} + +/// Project every supported predicate operand and disable shard-local filtering. +/// +/// The original expression remains under `OR TRUE` so PostgreSQL still performs +/// parse analysis and parameter type inference. +pub(super) fn rewrite_select<'a>( + select: &mut nodes::SelectStmtMut<'a, '_>, + mem: make::MemoryToken<'a>, +) -> Result { + let having = select.having_clause(); + if matches!(having, Node::None) { + return Ok(HavingRewritePlan::default()); + } + let having_source = mem.make_unique(having); + let having = having_source.as_ref(); + + let aliases = select + .target_list() + .iter() + .filter_map(|target| { + target + .name() + .or_else(|| target_output_name(target.val())) + .map(str::to_owned) + }) + .collect(); + let mut parser = Parser { + aliases, + values: Vec::new(), + next_alias: 0, + }; + let expression = parser.parse_expression(having)?; + let group_sources = select + .group_clause() + .iter() + .map(|group| mem.make_unique(group)) + .collect::>(); + for group_source in &group_sources { + let group = group_source.as_ref(); + match group { + Node::ColumnRef(_) => { + let alias = parser.next_alias(); + parser.values.push((alias, group)); + } + Node::A_Const(_) => {} + _ => { + return Err(Error::UnsupportedHaving(String::from( + "GROUP BY expressions other than columns and ordinals are unsupported", + ))); + } + } + } + let parsed = ParsedHaving { + expression, + values: parser.values, + }; + + let hidden_aliases = parsed + .values + .iter() + .map(|(alias, _)| alias.clone()) + .collect::>(); + let targets = parsed + .values + .iter() + .map(|(alias, node)| mem.make_res_target(Some(alias), mem.empty(), mem.make_unique(*node))) + .collect::>(); + select + .target_list_mut() + .extend(mem, mem.make_list(&targets)); + + let original = mem.make_unique(having); + let true_value = mem.make_a_const(ConstValue::Boolean(true)).uncast(); + let mut neutralized = mem.make_node::(); + neutralized + .as_mut() + .set_boolop(nodes::BoolExprType::OR_EXPR); + neutralized + .as_mut() + .set_args(mem.make_list(&[original, true_value])); + select.set_having_clause(neutralized.uncast()); + + Ok(HavingRewritePlan::new( + parsed.expression.materialize(), + hidden_aliases, + )) +} + +fn target_output_name(node: Node<'_>) -> Option<&str> { + match node { + Node::ColumnRef(column) => column.fields().iter().next_back()?.as_str(), + Node::FuncCall(function) => function.funcname().iter().next_back()?.as_str(), + Node::TypeCast(cast) => target_output_name(cast.arg()), + _ => None, + } +} + +impl HavingTemplate { + fn materialize(self) -> HavingExpr { + match self { + Self::Compare { + left, + operator, + right, + } => HavingExpr::Compare { + left: left.materialize(), + operator, + right: right.materialize(), + }, + Self::And(expressions) => { + HavingExpr::And(expressions.into_iter().map(Self::materialize).collect()) + } + Self::Or(expressions) => { + HavingExpr::Or(expressions.into_iter().map(Self::materialize).collect()) + } + Self::Not(expression) => HavingExpr::Not(Box::new(expression.materialize())), + Self::IsNull { value, negated } => HavingExpr::IsNull { + value: value.materialize(), + negated, + }, + } + } +} + +impl ValueTemplate { + fn materialize(self) -> HavingValue { + HavingValue::new(self.alias) + } +} + +struct Parser<'a> { + aliases: HashSet, + values: Vec<(String, Node<'a>)>, + next_alias: usize, +} + +impl<'a> Parser<'a> { + fn parse_expression(&mut self, node: Node<'a>) -> Result { + match node { + Node::A_Expr(expression) if expression.kind == nodes::A_Expr_Kind::AEXPR_OP => { + let operator = expression + .name() + .iter() + .exactly_one() + .ok() + .and_then(Node::as_str) + .and_then(Self::comparison_operator) + .ok_or_else(|| { + Error::UnsupportedHaving(String::from( + "only =, <>, !=, <, <=, > and >= comparisons are supported", + )) + })?; + Ok(HavingTemplate::Compare { + left: self.parse_value(expression.lexpr())?, + operator, + right: self.parse_value(expression.rexpr())?, + }) + } + Node::BoolExpr(expression) => { + let expressions = expression + .args() + .iter() + .map(|node| self.parse_expression(node)) + .collect::, _>>()?; + match expression.boolop { + nodes::BoolExprType::AND_EXPR => Ok(HavingTemplate::And(expressions)), + nodes::BoolExprType::OR_EXPR => Ok(HavingTemplate::Or(expressions)), + nodes::BoolExprType::NOT_EXPR => { + let mut expressions = expressions.into_iter(); + let expression = expressions.next().ok_or_else(|| { + Error::UnsupportedHaving(String::from("NOT requires one expression")) + })?; + if expressions.next().is_some() { + return Err(Error::UnsupportedHaving(String::from( + "NOT requires one expression", + ))); + } + Ok(HavingTemplate::Not(Box::new(expression))) + } + _ => Err(Error::UnsupportedHaving(String::from( + "unsupported boolean expression", + ))), + } + } + Node::NullTest(test) => { + let negated = match test.nulltesttype { + nodes::NullTestType::IS_NULL => false, + nodes::NullTestType::IS_NOT_NULL => true, + _ => { + return Err(Error::UnsupportedHaving(String::from( + "only IS NULL and IS NOT NULL are supported", + ))); + } + }; + Ok(HavingTemplate::IsNull { + value: self.parse_value(test.arg())?, + negated, + }) + } + _ => Err(Error::UnsupportedHaving(String::from( + "supported predicates are comparisons, AND, OR, NOT and NULL tests", + ))), + } + } + + fn parse_value(&mut self, node: Node<'a>) -> Result { + let supported = match node { + Node::A_Const(_) | Node::ParamRef(_) | Node::ColumnRef(_) => true, + Node::TypeCast(cast) => Self::supported_value(cast.arg()), + Node::FuncCall(_) => Self::supported_aggregate(node), + _ => false, + }; + if !supported { + return Err(Error::UnsupportedHaving(String::from( + "operands must be aggregate calls, grouped columns, constants or parameters", + ))); + } + + let alias = self.next_alias(); + self.values.push((alias.clone(), node)); + Ok(ValueTemplate { alias }) + } + + fn supported_value(node: Node<'_>) -> bool { + match node { + Node::A_Const(_) | Node::ParamRef(_) | Node::ColumnRef(_) => true, + Node::TypeCast(cast) => Self::supported_value(cast.arg()), + Node::FuncCall(_) => Self::supported_aggregate(node), + _ => false, + } + } + + fn supported_aggregate(node: Node<'_>) -> bool { + Function::try_from(node).is_ok_and(|function| { + matches!( + function.name, + "count" + | "max" + | "min" + | "sum" + | "avg" + | "stddev" + | "stddev_samp" + | "stddev_pop" + | "variance" + | "var_samp" + | "var_pop" + ) + }) + } + + fn comparison_operator(operator: &str) -> Option { + Some(match operator { + "=" => CompareOperator::Equal, + "<>" | "!=" => CompareOperator::NotEqual, + "<" => CompareOperator::Less, + "<=" => CompareOperator::LessOrEqual, + ">" => CompareOperator::Greater, + ">=" => CompareOperator::GreaterOrEqual, + _ => return None, + }) + } + + fn next_alias(&mut self) -> String { + loop { + let alias = format!("{ALIAS_PREFIX}{}", self.next_alias); + self.next_alias += 1; + if self.aliases.insert(alias.clone()) { + return alias; + } + } + } +} diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs b/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs index 3c30d15cf..2ff773dd3 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs @@ -11,6 +11,7 @@ use pg_raw_parse::{Node, NodeMut, make, nodes, transform, walk}; pub(crate) mod aggregate; pub(crate) mod auto_id; pub(crate) mod error; +mod having; pub(crate) mod insert; pub(crate) mod insert_split_plan; pub(crate) mod nextval; diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs b/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs index 772246586..0ae20da87 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs @@ -6,10 +6,11 @@ use super::Error; use super::aggregate::{AggregatesRewrite, HelperKind}; +use super::having; use super::offset::{self, OffsetPlan}; use super::order_by; use crate::backend::schema::Schema; -use crate::frontend::router::parser::{Aggregate, OrderBy}; +use crate::frontend::router::parser::{Aggregate, HavingRewritePlan, OrderBy}; use crate::frontend::{ClientRequest, PreparedStatements}; use crate::net::{ProtocolMessage, RowDescription}; use pg_raw_parse::{Node, StmtList, make}; @@ -57,11 +58,14 @@ impl OrderByHelper { pub(crate) struct ProjectionRewritePlan { pub(crate) aggregate_helpers: Vec, pub(crate) order_by_helpers: Vec, + pub(crate) having: HavingRewritePlan, } impl ProjectionRewritePlan { pub(crate) fn is_noop(&self) -> bool { - self.aggregate_helpers.is_empty() && self.order_by_helpers.is_empty() + self.aggregate_helpers.is_empty() + && self.order_by_helpers.is_empty() + && self.having.is_noop() } pub(crate) fn drop_columns(&self, row_description: &RowDescription) -> BTreeSet { @@ -74,6 +78,7 @@ impl ProjectionRewritePlan { .filter(|helper| helper.injected) .map(|helper| helper.alias.as_str()), ) + .chain(self.having.hidden_aliases().iter().map(String::as_str)) .filter_map(|alias| row_description.field_index(alias)) .collect() } @@ -83,6 +88,7 @@ impl ProjectionRewritePlan { pub(crate) struct PostRouteRewrite { sql: Arc, plan: ProjectionRewritePlan, + aggregate: Option, } pub(crate) fn finalize_after_route( @@ -148,6 +154,9 @@ pub(crate) fn finalize_after_route( && let Some(route) = request.route.as_mut() { route.projection_rewrite_plan = rewrite.plan.clone(); + if let Some(aggregate) = &rewrite.aggregate { + route.set_aggregate(aggregate.clone()); + } let mut order_by = route.order_by().to_vec(); for helper in &rewrite.plan.order_by_helpers { // Prefer the structural position. The source fallback handles a @@ -181,29 +190,53 @@ fn build( return Ok(None); }; + let has_having = !matches!(select.having_clause(), Node::None); let aggregate = Aggregate::parse(select, schema); - if aggregate.is_empty() && select.sort_clause().is_empty() && !rewrite_offset { + if aggregate.is_empty() && !has_having && select.sort_clause().is_empty() && !rewrite_offset { return Ok(None); } let mut plan = ProjectionRewritePlan::default(); - let rewritten = make::owned(|mem| { + let mut rewritten_aggregate = None; + let rewritten = make::try_owned(|mem| { let mut select = mem.make_unique(select); + if has_having { + plan.having = having::rewrite_select(&mut select.as_mut(), mem)?; + let aggregate = Aggregate::parse(&select, schema); + if aggregate.is_empty() { + return Err(Error::UnsupportedHaving(String::from( + "at least one aggregate function is required", + ))); + } + if aggregate.group_by().len() != select.group_clause().len() { + return Err(Error::UnsupportedHaving(String::from( + "every GROUP BY expression must be a selected column or ordinal", + ))); + } + rewritten_aggregate = Some(aggregate); + } + let aggregate = rewritten_aggregate.as_ref().unwrap_or(&aggregate); if !aggregate.is_empty() { - plan = AggregatesRewrite::rewrite_select(&mut select.as_mut(), mem, &aggregate); + let aggregate_plan = + AggregatesRewrite::rewrite_select(&mut select.as_mut(), mem, aggregate); + plan.aggregate_helpers = aggregate_plan.aggregate_helpers; } order_by::rewrite_select(&mut select.as_mut(), mem, &mut plan); if rewrite_offset { offset::rewrite_select(&mut select.as_mut(), mem); } - select - }); + Ok::<_, Error>(select) + })?; if plan.is_noop() && !rewrite_offset { return Ok(None); } let sql: Arc = pg_raw_parse::deparse(&*rewritten)?.as_str().into(); - Ok(Some(PostRouteRewrite { sql, plan })) + Ok(Some(PostRouteRewrite { + sql, + plan, + aggregate: rewritten_aggregate, + })) } #[cfg(test)] @@ -253,4 +286,117 @@ mod tests { ]); assert!(plan.drop_columns(&row_description).is_empty()); } + + #[test] + fn having_projects_operands_and_updates_aggregate_plan() { + let ast = pg_raw_parse::parse( + "SELECT customer_id, COUNT(*) AS orders \ + FROM purchases GROUP BY customer_id HAVING COUNT(*) > 2", + ) + .unwrap(); + let rewrite = build(&ast, &Schema::default(), false) + .unwrap() + .expect("HAVING requires a rewrite"); + + assert!(rewrite.sql.contains("__pgdog_having_0")); + assert!(rewrite.sql.contains("__pgdog_having_1")); + assert_eq!(rewrite.plan.having.hidden_aliases().len(), 3); + assert_eq!( + rewrite + .aggregate + .as_ref() + .expect("rewritten aggregate") + .targets() + .len(), + 2, + "the selected and hidden COUNT must both be merged" + ); + } + + #[test] + fn having_internal_aliases_do_not_collide_with_client_aliases() { + let ast = pg_raw_parse::parse( + "SELECT __pgdog_having_0, COUNT(*) \ + FROM purchases GROUP BY __pgdog_having_0 HAVING COUNT(*) > 2", + ) + .unwrap(); + let rewrite = build(&ast, &Schema::default(), false) + .unwrap() + .expect("HAVING requires a rewrite"); + + assert_eq!( + rewrite.plan.having.hidden_aliases(), + &["__pgdog_having_1", "__pgdog_having_2", "__pgdog_having_3"] + ); + } + + #[test] + fn having_projects_unselected_grouping_column() { + let ast = pg_raw_parse::parse( + "SELECT COUNT(*) FROM purchases \ + GROUP BY customer_id HAVING COUNT(*) > 2", + ) + .unwrap(); + let rewrite = build(&ast, &Schema::default(), false) + .unwrap() + .expect("HAVING requires a rewrite"); + let aggregate = rewrite.aggregate.expect("rewritten aggregate"); + + assert_eq!(aggregate.group_by().len(), 1); + assert_eq!(aggregate.targets().len(), 2); + } + + #[test] + fn having_avg_receives_independent_count_helper() { + let ast = pg_raw_parse::parse( + "SELECT customer_id, AVG(price) FROM purchases \ + GROUP BY customer_id HAVING AVG(price) > 10", + ) + .unwrap(); + let rewrite = build(&ast, &Schema::default(), false) + .unwrap() + .expect("HAVING requires a rewrite"); + + assert_eq!( + rewrite + .aggregate + .as_ref() + .expect("rewritten aggregate") + .targets() + .len(), + 2 + ); + assert_eq!( + rewrite.plan.aggregate_helpers.len(), + 2, + "selected and hidden AVG each require their own COUNT" + ); + } + + #[test] + fn aggregate_used_only_by_having_is_merged_and_hidden() { + let ast = pg_raw_parse::parse( + "SELECT customer_id FROM purchases \ + GROUP BY customer_id HAVING COUNT(*) > 2", + ) + .unwrap(); + let rewrite = build(&ast, &Schema::default(), false) + .unwrap() + .expect("HAVING requires a rewrite"); + let aggregate = rewrite.aggregate.expect("rewritten aggregate"); + + assert_eq!(aggregate.targets().len(), 1); + assert_eq!(aggregate.group_by().len(), 1); + } + + #[test] + fn unsupported_having_expression_fails_closed() { + let ast = pg_raw_parse::parse( + "SELECT customer_id, SUM(price) FROM purchases \ + GROUP BY customer_id HAVING SUM(price) / COUNT(*) > 2", + ) + .unwrap(); + let error = build(&ast, &Schema::default(), false).unwrap_err(); + assert!(matches!(error, Error::UnsupportedHaving(_))); + } } diff --git a/pgdog/src/frontend/router/parser/route.rs b/pgdog/src/frontend/router/parser/route.rs index afc53609a..88699d3a2 100644 --- a/pgdog/src/frontend/router/parser/route.rs +++ b/pgdog/src/frontend/router/parser/route.rs @@ -273,6 +273,10 @@ impl Route { &self.aggregate } + pub(crate) fn set_aggregate(&mut self, aggregate: Aggregate) { + self.aggregate = aggregate; + } + /// Set shard on this route, along with reasoning /// for that shard selection. pub(crate) fn set_shard(&mut self, shard: ShardWithPriority) { @@ -343,6 +347,7 @@ impl Route { pub(crate) fn requires_post_processing(&self) -> bool { !self.order_by().is_empty() || !self.aggregate().is_empty() + || !self.projection_rewrite_plan.having.is_noop() || self.distinct().is_some() || self.limit().limit.is_some() || self.limit().offset.is_some() From 98291d08f199aa54a6cd13fd41a21237afa5d556 Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Fri, 16 Oct 2026 22:17:18 +0530 Subject: [PATCH 02/24] improvements --- pgdog/src/backend/pool/connection/buffer.rs | 5 + .../backend/pool/connection/buffer/having.rs | 125 ++++++++++++++++++ .../router/parser/rewrite/statement/having.rs | 80 +++++++++++ .../parser/rewrite/statement/projection.rs | 17 ++- 4 files changed, 221 insertions(+), 6 deletions(-) diff --git a/pgdog/src/backend/pool/connection/buffer.rs b/pgdog/src/backend/pool/connection/buffer.rs index 278f54507..773d26981 100644 --- a/pgdog/src/backend/pool/connection/buffer.rs +++ b/pgdog/src/backend/pool/connection/buffer.rs @@ -374,6 +374,11 @@ mod test { Field::bigint("__pgdog_having_0"), Field::bigint("__pgdog_having_1"), ])); + let mut noop = Buffer::default(); + noop.having(&ProjectionRewritePlan::default(), &decoder) + .unwrap(); + assert!(noop.is_empty()); + let plan = ProjectionRewritePlan { having: crate::frontend::router::parser::HavingRewritePlan::new( HavingExpr::Compare { diff --git a/pgdog/src/backend/pool/connection/buffer/having.rs b/pgdog/src/backend/pool/connection/buffer/having.rs index dbd6d0d81..01e4e083a 100644 --- a/pgdog/src/backend/pool/connection/buffer/having.rs +++ b/pgdog/src/backend/pool/connection/buffer/having.rs @@ -206,6 +206,7 @@ mod tests { messages::{DataRow, Datum}, }, }; + use pgdog_postgres_types::{Double, Float}; fn compare_expression(operator: CompareOperator) -> HavingExpr { HavingExpr::Compare { @@ -260,6 +261,47 @@ mod tests { value: HavingValue::new("__pgdog_having_0".into()), negated: false, }; + let is_not_null = HavingExpr::IsNull { + value: HavingValue::new("__pgdog_having_1".into()), + negated: true, + }; + let known_false = HavingExpr::Compare { + left: HavingValue::new("__pgdog_having_1".into()), + operator: CompareOperator::Greater, + right: HavingValue::new("__pgdog_having_1".into()), + }; + + assert_eq!( + evaluate_expression(&unknown, &row, &decoder).unwrap(), + TruthValue::Unknown + ); + assert_eq!( + evaluate_expression( + &HavingExpr::And(vec![is_null.clone(), unknown.clone()]), + &row, + &decoder + ) + .unwrap(), + TruthValue::Unknown + ); + assert_eq!( + evaluate_expression( + &HavingExpr::And(vec![unknown.clone(), known_false.clone()]), + &row, + &decoder + ) + .unwrap(), + TruthValue::False + ); + assert_eq!( + evaluate_expression( + &HavingExpr::Or(vec![known_false.clone(), unknown.clone()]), + &row, + &decoder + ) + .unwrap(), + TruthValue::Unknown + ); assert!( evaluate( &HavingExpr::Or(vec![unknown.clone(), is_null]), @@ -268,6 +310,89 @@ mod tests { ) .unwrap() ); + assert!(evaluate(&is_not_null, &row, &decoder).unwrap()); assert!(!evaluate(&HavingExpr::Not(Box::new(unknown)), &row, &decoder).unwrap()); + + let missing = HavingExpr::Compare { + left: HavingValue::new("__pgdog_missing".into()), + operator: CompareOperator::Equal, + right: HavingValue::new("__pgdog_having_1".into()), + }; + assert!(matches!( + evaluate(&missing, &row, &decoder), + Err(Error::InvalidHaving(_)) + )); + } + + #[test] + fn compares_supported_scalar_types_and_all_operators() { + for (operator, left, right, expected) in [ + ( + CompareOperator::Equal, + Datum::Bigint(2), + Datum::Integer(2), + true, + ), + ( + CompareOperator::NotEqual, + Datum::Integer(2), + Datum::SmallInt(3), + true, + ), + ( + CompareOperator::Less, + Datum::SmallInt(2), + Datum::Bigint(3), + true, + ), + ( + CompareOperator::LessOrEqual, + Datum::Bigint(2), + Datum::Integer(2), + true, + ), + ( + CompareOperator::Greater, + Datum::Numeric(Numeric::from(3_i32)), + Datum::SmallInt(2), + true, + ), + ( + CompareOperator::GreaterOrEqual, + Datum::Float(Float(2.0)), + Datum::Double(Double(2.0)), + true, + ), + ( + CompareOperator::Greater, + Datum::Text("a".into()), + Datum::Text("z".into()), + false, + ), + ] { + assert_eq!( + compare(&left, operator, &right).unwrap(), + if expected { + TruthValue::True + } else { + TruthValue::False + } + ); + } + + assert_eq!( + numeric_ordering(&Datum::Double(Double(f64::NAN)), &Datum::Float(Float(1.0))), + Some(Ordering::Greater) + ); + assert_eq!(postgres_float_ordering(f64::NAN, f64::NAN), Ordering::Equal); + assert_eq!(postgres_float_ordering(1.0, f64::NAN), Ordering::Less); + assert!(matches!( + compare( + &Datum::Text("2".into()), + CompareOperator::Equal, + &Datum::Integer(2) + ), + Err(Error::UnsupportedHaving(_)) + )); } } diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/having.rs b/pgdog/src/frontend/router/parser/rewrite/statement/having.rs index 4973d9b5e..6b61afb27 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/having.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/having.rs @@ -309,3 +309,83 @@ impl<'a> Parser<'a> { } } } + +#[cfg(test)] +mod tests { + use super::*; + + fn rewrite(sql: &str) -> (String, HavingRewritePlan) { + let ast = pg_raw_parse::parse(sql).expect("parse HAVING query"); + let Node::SelectStmt(select) = ast.stmts().next().expect("one statement") else { + panic!("expected SELECT") + }; + let mut plan = None; + let rewritten = make::owned(|mem| { + let mut select = mem.make_unique(select); + plan = + Some(rewrite_select(&mut select.as_mut(), mem).expect("rewrite supported HAVING")); + select + }); + ( + pg_raw_parse::deparse(&*rewritten) + .expect("deparse rewritten query") + .as_str() + .to_owned(), + plan.expect("rewrite plan"), + ) + } + + #[test] + fn rewrites_nested_boolean_null_cast_and_parameter_expression() { + let (sql, plan) = rewrite( + "SELECT customer_id, SUM(price)::numeric FROM purchases \ + GROUP BY customer_id \ + HAVING NOT (COUNT(*) <= $1 OR SUM(price) IS NOT NULL) \ + AND MAX(price)::numeric >= 10", + ); + + let Some(HavingExpr::And(expressions)) = plan.expression() else { + panic!("expected top-level AND") + }; + assert!(matches!( + expressions.as_slice(), + [ + HavingExpr::Not(not), + HavingExpr::Compare { + operator: CompareOperator::GreaterOrEqual, + .. + } + ] if matches!( + not.as_ref(), + HavingExpr::Or(or) if matches!( + or.as_slice(), + [ + HavingExpr::Compare { + operator: CompareOperator::LessOrEqual, + .. + }, + HavingExpr::IsNull { negated: true, .. } + ] + ) + ) + )); + assert!(sql.contains("OR true")); + assert_eq!(plan.hidden_aliases().len(), 6); + } + + #[test] + fn parses_every_comparison_operator() { + for (sql, expected) in [ + ("=", CompareOperator::Equal), + ("<>", CompareOperator::NotEqual), + ("!=", CompareOperator::NotEqual), + ("<", CompareOperator::Less), + ("<=", CompareOperator::LessOrEqual), + (">", CompareOperator::Greater), + (">=", CompareOperator::GreaterOrEqual), + ] { + assert_eq!(Parser::comparison_operator(sql), Some(expected)); + } + assert_eq!(Parser::comparison_operator("~"), None); + } +} diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs b/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs index 0ae20da87..981f6d5e4 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs @@ -390,13 +390,18 @@ mod tests { } #[test] - fn unsupported_having_expression_fails_closed() { - let ast = pg_raw_parse::parse( + fn unsupported_having_queries_fail_closed() { + for sql in [ "SELECT customer_id, SUM(price) FROM purchases \ GROUP BY customer_id HAVING SUM(price) / COUNT(*) > 2", - ) - .unwrap(); - let error = build(&ast, &Schema::default(), false).unwrap_err(); - assert!(matches!(error, Error::UnsupportedHaving(_))); + "SELECT customer_id FROM purchases \ + GROUP BY customer_id HAVING customer_id > 2", + "SELECT lower(customer_id), COUNT(*) FROM purchases \ + GROUP BY lower(customer_id) HAVING COUNT(*) > 2", + ] { + let ast = pg_raw_parse::parse(sql).unwrap(); + let error = build(&ast, &Schema::default(), false).unwrap_err(); + assert!(matches!(error, Error::UnsupportedHaving(_)), "{sql}"); + } } } From 718ff441582a85a9f3e5c5f2274a89e7496a381b Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Fri, 16 Oct 2026 22:37:10 +0530 Subject: [PATCH 03/24] use arc for having rewrite plan --- pgdog/src/backend/pool/connection/buffer.rs | 22 +++++++------- .../query_engine/test/rewrite_projection.rs | 1 - pgdog/src/frontend/router/parser/having.rs | 14 ++++----- .../router/parser/rewrite/statement/having.rs | 6 ++-- .../parser/rewrite/statement/projection.rs | 30 +++++++++++++++---- pgdog/src/frontend/router/parser/route.rs | 2 +- 6 files changed, 46 insertions(+), 29 deletions(-) diff --git a/pgdog/src/backend/pool/connection/buffer.rs b/pgdog/src/backend/pool/connection/buffer.rs index 773d26981..d2ffecfc6 100644 --- a/pgdog/src/backend/pool/connection/buffer.rs +++ b/pgdog/src/backend/pool/connection/buffer.rs @@ -164,13 +164,13 @@ impl Buffer { plan: &ProjectionRewritePlan, decoder: &Decoder, ) -> Result<(), super::Error> { - let Some(expression) = plan.having.expression() else { + let Some(having) = plan.having.as_deref() else { return Ok(()); }; let mut filtered = VecDeque::with_capacity(self.buffer.len()); for row in self.buffer.drain(..) { - if having::evaluate(expression, &row, decoder)? { + if having::evaluate(having.expression(), &row, decoder)? { filtered.push_back(row); } } @@ -380,14 +380,16 @@ mod test { assert!(noop.is_empty()); let plan = ProjectionRewritePlan { - having: crate::frontend::router::parser::HavingRewritePlan::new( - HavingExpr::Compare { - left: HavingValue::new("__pgdog_having_0".into()), - operator: CompareOperator::Greater, - right: HavingValue::new("__pgdog_having_1".into()), - }, - vec!["__pgdog_having_0".into(), "__pgdog_having_1".into()], - ), + having: Some(std::sync::Arc::new( + crate::frontend::router::parser::HavingRewritePlan::new( + HavingExpr::Compare { + left: HavingValue::new("__pgdog_having_0".into()), + operator: CompareOperator::Greater, + right: HavingValue::new("__pgdog_having_1".into()), + }, + vec!["__pgdog_having_0".into(), "__pgdog_having_1".into()], + ), + )), ..Default::default() }; diff --git a/pgdog/src/frontend/client/query_engine/test/rewrite_projection.rs b/pgdog/src/frontend/client/query_engine/test/rewrite_projection.rs index 88fa18741..394846a4a 100644 --- a/pgdog/src/frontend/client/query_engine/test/rewrite_projection.rs +++ b/pgdog/src/frontend/client/query_engine/test/rewrite_projection.rs @@ -132,7 +132,6 @@ async fn cross_shard_having_rewrites_and_replaces_route_aggregate() { .route() .projection_rewrite_plan .having - .expression() .is_some() ); } diff --git a/pgdog/src/frontend/router/parser/having.rs b/pgdog/src/frontend/router/parser/having.rs index ca83d0b59..5ad23a267 100644 --- a/pgdog/src/frontend/router/parser/having.rs +++ b/pgdog/src/frontend/router/parser/having.rs @@ -45,29 +45,25 @@ pub(crate) enum HavingExpr { } /// Metadata produced while rewriting a cross-shard `HAVING` clause. -#[derive(Debug, Clone, Default, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq)] pub(crate) struct HavingRewritePlan { - expression: Option, + expression: HavingExpr, hidden_aliases: Vec, } impl HavingRewritePlan { pub(crate) fn new(expression: HavingExpr, hidden_aliases: Vec) -> Self { Self { - expression: Some(expression), + expression, hidden_aliases, } } - pub(crate) fn expression(&self) -> Option<&HavingExpr> { - self.expression.as_ref() + pub(crate) fn expression(&self) -> &HavingExpr { + &self.expression } pub(crate) fn hidden_aliases(&self) -> &[String] { &self.hidden_aliases } - - pub(crate) fn is_noop(&self) -> bool { - self.expression.is_none() - } } diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/having.rs b/pgdog/src/frontend/router/parser/rewrite/statement/having.rs index 6b61afb27..c54514c53 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/having.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/having.rs @@ -47,7 +47,9 @@ pub(super) fn rewrite_select<'a>( ) -> Result { let having = select.having_clause(); if matches!(having, Node::None) { - return Ok(HavingRewritePlan::default()); + return Err(Error::UnsupportedHaving(String::from( + "HAVING rewrite requested without a HAVING clause", + ))); } let having_source = mem.make_unique(having); let having = having_source.as_ref(); @@ -344,7 +346,7 @@ mod tests { AND MAX(price)::numeric >= 10", ); - let Some(HavingExpr::And(expressions)) = plan.expression() else { + let HavingExpr::And(expressions) = plan.expression() else { panic!("expected top-level AND") }; assert!(matches!( diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs b/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs index 981f6d5e4..fbd79b05c 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs @@ -58,14 +58,14 @@ impl OrderByHelper { pub(crate) struct ProjectionRewritePlan { pub(crate) aggregate_helpers: Vec, pub(crate) order_by_helpers: Vec, - pub(crate) having: HavingRewritePlan, + pub(crate) having: Option>, } impl ProjectionRewritePlan { pub(crate) fn is_noop(&self) -> bool { self.aggregate_helpers.is_empty() && self.order_by_helpers.is_empty() - && self.having.is_noop() + && self.having.is_none() } pub(crate) fn drop_columns(&self, row_description: &RowDescription) -> BTreeSet { @@ -78,7 +78,11 @@ impl ProjectionRewritePlan { .filter(|helper| helper.injected) .map(|helper| helper.alias.as_str()), ) - .chain(self.having.hidden_aliases().iter().map(String::as_str)) + .chain( + self.having + .iter() + .flat_map(|plan| plan.hidden_aliases().iter().map(String::as_str)), + ) .filter_map(|alias| row_description.field_index(alias)) .collect() } @@ -201,7 +205,7 @@ fn build( let rewritten = make::try_owned(|mem| { let mut select = mem.make_unique(select); if has_having { - plan.having = having::rewrite_select(&mut select.as_mut(), mem)?; + plan.having = Some(Arc::new(having::rewrite_select(&mut select.as_mut(), mem)?)); let aggregate = Aggregate::parse(&select, schema); if aggregate.is_empty() { return Err(Error::UnsupportedHaving(String::from( @@ -300,7 +304,16 @@ mod tests { assert!(rewrite.sql.contains("__pgdog_having_0")); assert!(rewrite.sql.contains("__pgdog_having_1")); - assert_eq!(rewrite.plan.having.hidden_aliases().len(), 3); + assert_eq!( + rewrite + .plan + .having + .as_deref() + .expect("HAVING plan") + .hidden_aliases() + .len(), + 3 + ); assert_eq!( rewrite .aggregate @@ -325,7 +338,12 @@ mod tests { .expect("HAVING requires a rewrite"); assert_eq!( - rewrite.plan.having.hidden_aliases(), + rewrite + .plan + .having + .as_deref() + .expect("HAVING plan") + .hidden_aliases(), &["__pgdog_having_1", "__pgdog_having_2", "__pgdog_having_3"] ); } diff --git a/pgdog/src/frontend/router/parser/route.rs b/pgdog/src/frontend/router/parser/route.rs index 88699d3a2..c3859c4b5 100644 --- a/pgdog/src/frontend/router/parser/route.rs +++ b/pgdog/src/frontend/router/parser/route.rs @@ -347,7 +347,7 @@ impl Route { pub(crate) fn requires_post_processing(&self) -> bool { !self.order_by().is_empty() || !self.aggregate().is_empty() - || !self.projection_rewrite_plan.having.is_noop() + || self.projection_rewrite_plan.having.is_some() || self.distinct().is_some() || self.limit().limit.is_some() || self.limit().offset.is_some() From 9b68185d55b27c927f4abfb19a8f63c53f0485fb Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Fri, 16 Oct 2026 23:04:25 +0530 Subject: [PATCH 04/24] store rewritten aggreate in post route plan --- pgdog/src/backend/pool/connection/buffer.rs | 6 +- .../parser/rewrite/statement/projection.rs | 87 ++++++++++++------- 2 files changed, 62 insertions(+), 31 deletions(-) diff --git a/pgdog/src/backend/pool/connection/buffer.rs b/pgdog/src/backend/pool/connection/buffer.rs index d2ffecfc6..c0a5a15ce 100644 --- a/pgdog/src/backend/pool/connection/buffer.rs +++ b/pgdog/src/backend/pool/connection/buffer.rs @@ -269,6 +269,7 @@ impl Buffer { #[cfg(test)] mod test { use super::*; + use crate::frontend::router::parser::rewrite::statement::projection::HavingPostRoutePlan; use crate::frontend::router::parser::{CompareOperator, HavingExpr, HavingValue}; use crate::net::{Datum, Field, Format, RowDescription}; use bytes::Bytes; @@ -380,7 +381,7 @@ mod test { assert!(noop.is_empty()); let plan = ProjectionRewritePlan { - having: Some(std::sync::Arc::new( + having: Some(std::sync::Arc::new(HavingPostRoutePlan::new( crate::frontend::router::parser::HavingRewritePlan::new( HavingExpr::Compare { left: HavingValue::new("__pgdog_having_0".into()), @@ -389,7 +390,8 @@ mod test { }, vec!["__pgdog_having_0".into(), "__pgdog_having_1".into()], ), - )), + aggregate.clone(), + ))), ..Default::default() }; diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs b/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs index fbd79b05c..5ad982245 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs @@ -10,7 +10,7 @@ use super::having; use super::offset::{self, OffsetPlan}; use super::order_by; use crate::backend::schema::Schema; -use crate::frontend::router::parser::{Aggregate, HavingRewritePlan, OrderBy}; +use crate::frontend::router::parser::{Aggregate, HavingExpr, HavingRewritePlan, OrderBy}; use crate::frontend::{ClientRequest, PreparedStatements}; use crate::net::{ProtocolMessage, RowDescription}; use pg_raw_parse::{Node, StmtList, make}; @@ -58,7 +58,31 @@ impl OrderByHelper { pub(crate) struct ProjectionRewritePlan { pub(crate) aggregate_helpers: Vec, pub(crate) order_by_helpers: Vec, - pub(crate) having: Option>, + pub(crate) having: Option>, +} + +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct HavingPostRoutePlan { + rewrite: HavingRewritePlan, + aggregate: Aggregate, +} + +impl HavingPostRoutePlan { + pub(crate) fn new(rewrite: HavingRewritePlan, aggregate: Aggregate) -> Self { + Self { rewrite, aggregate } + } + + pub(crate) fn expression(&self) -> &HavingExpr { + self.rewrite.expression() + } + + pub(crate) fn hidden_aliases(&self) -> &[String] { + self.rewrite.hidden_aliases() + } + + pub(crate) fn aggregate(&self) -> &Aggregate { + &self.aggregate + } } impl ProjectionRewritePlan { @@ -92,7 +116,6 @@ impl ProjectionRewritePlan { pub(crate) struct PostRouteRewrite { sql: Arc, plan: ProjectionRewritePlan, - aggregate: Option, } pub(crate) fn finalize_after_route( @@ -158,8 +181,8 @@ pub(crate) fn finalize_after_route( && let Some(route) = request.route.as_mut() { route.projection_rewrite_plan = rewrite.plan.clone(); - if let Some(aggregate) = &rewrite.aggregate { - route.set_aggregate(aggregate.clone()); + if let Some(having) = &rewrite.plan.having { + route.set_aggregate(having.aggregate().clone()); } let mut order_by = route.order_by().to_vec(); for helper in &rewrite.plan.order_by_helpers { @@ -201,25 +224,31 @@ fn build( } let mut plan = ProjectionRewritePlan::default(); - let mut rewritten_aggregate = None; let rewritten = make::try_owned(|mem| { let mut select = mem.make_unique(select); if has_having { - plan.having = Some(Arc::new(having::rewrite_select(&mut select.as_mut(), mem)?)); - let aggregate = Aggregate::parse(&select, schema); - if aggregate.is_empty() { + let having = having::rewrite_select(&mut select.as_mut(), mem)?; + let rewritten_aggregate = Aggregate::parse(&select, schema); + if rewritten_aggregate.is_empty() { return Err(Error::UnsupportedHaving(String::from( "at least one aggregate function is required", ))); } - if aggregate.group_by().len() != select.group_clause().len() { + if rewritten_aggregate.group_by().len() != select.group_clause().len() { return Err(Error::UnsupportedHaving(String::from( "every GROUP BY expression must be a selected column or ordinal", ))); } - rewritten_aggregate = Some(aggregate); + plan.having = Some(Arc::new(HavingPostRoutePlan::new( + having, + rewritten_aggregate, + ))); } - let aggregate = rewritten_aggregate.as_ref().unwrap_or(&aggregate); + let aggregate = plan + .having + .as_deref() + .map(HavingPostRoutePlan::aggregate) + .unwrap_or(&aggregate); if !aggregate.is_empty() { let aggregate_plan = AggregatesRewrite::rewrite_select(&mut select.as_mut(), mem, aggregate); @@ -236,11 +265,7 @@ fn build( } let sql: Arc = pg_raw_parse::deparse(&*rewritten)?.as_str().into(); - Ok(Some(PostRouteRewrite { - sql, - plan, - aggregate: rewritten_aggregate, - })) + Ok(Some(PostRouteRewrite { sql, plan })) } #[cfg(test)] @@ -316,9 +341,11 @@ mod tests { ); assert_eq!( rewrite - .aggregate - .as_ref() - .expect("rewritten aggregate") + .plan + .having + .as_deref() + .expect("HAVING plan") + .aggregate() .targets() .len(), 2, @@ -358,10 +385,10 @@ mod tests { let rewrite = build(&ast, &Schema::default(), false) .unwrap() .expect("HAVING requires a rewrite"); - let aggregate = rewrite.aggregate.expect("rewritten aggregate"); + let aggregate = rewrite.plan.having.expect("HAVING plan"); - assert_eq!(aggregate.group_by().len(), 1); - assert_eq!(aggregate.targets().len(), 2); + assert_eq!(aggregate.aggregate().group_by().len(), 1); + assert_eq!(aggregate.aggregate().targets().len(), 2); } #[test] @@ -377,9 +404,11 @@ mod tests { assert_eq!( rewrite - .aggregate - .as_ref() - .expect("rewritten aggregate") + .plan + .having + .as_deref() + .expect("HAVING plan") + .aggregate() .targets() .len(), 2 @@ -401,10 +430,10 @@ mod tests { let rewrite = build(&ast, &Schema::default(), false) .unwrap() .expect("HAVING requires a rewrite"); - let aggregate = rewrite.aggregate.expect("rewritten aggregate"); + let aggregate = rewrite.plan.having.expect("HAVING plan"); - assert_eq!(aggregate.targets().len(), 1); - assert_eq!(aggregate.group_by().len(), 1); + assert_eq!(aggregate.aggregate().targets().len(), 1); + assert_eq!(aggregate.aggregate().group_by().len(), 1); } #[test] From f24c437661b35c9f05dbced8b4937de42473830d Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Mon, 19 Oct 2026 22:50:34 +0530 Subject: [PATCH 05/24] nits --- pgdog/src/backend/pool/connection/buffer.rs | 4 +- .../backend/pool/connection/buffer/having.rs | 45 ++++-------- .../pool/connection/multi_shard/mod.rs | 24 +++---- pgdog/src/frontend/router/parser/aggregate.rs | 41 ++++++----- .../router/parser/rewrite/statement/error.rs | 2 +- .../router/parser/rewrite/statement/having.rs | 69 ++++++------------- .../parser/rewrite/statement/projection.rs | 40 +++++++---- 7 files changed, 99 insertions(+), 126 deletions(-) diff --git a/pgdog/src/backend/pool/connection/buffer.rs b/pgdog/src/backend/pool/connection/buffer.rs index c0a5a15ce..2bdb04813 100644 --- a/pgdog/src/backend/pool/connection/buffer.rs +++ b/pgdog/src/backend/pool/connection/buffer.rs @@ -269,7 +269,7 @@ impl Buffer { #[cfg(test)] mod test { use super::*; - use crate::frontend::router::parser::rewrite::statement::projection::HavingPostRoutePlan; + use crate::frontend::router::parser::rewrite::statement::projection::HavingRewriteResult; use crate::frontend::router::parser::{CompareOperator, HavingExpr, HavingValue}; use crate::net::{Datum, Field, Format, RowDescription}; use bytes::Bytes; @@ -381,7 +381,7 @@ mod test { assert!(noop.is_empty()); let plan = ProjectionRewritePlan { - having: Some(std::sync::Arc::new(HavingPostRoutePlan::new( + having: Some(std::sync::Arc::new(HavingRewriteResult::new( crate::frontend::router::parser::HavingRewritePlan::new( HavingExpr::Compare { left: HavingValue::new("__pgdog_having_0".into()), diff --git a/pgdog/src/backend/pool/connection/buffer/having.rs b/pgdog/src/backend/pool/connection/buffer/having.rs index 01e4e083a..91261b25c 100644 --- a/pgdog/src/backend/pool/connection/buffer/having.rs +++ b/pgdog/src/backend/pool/connection/buffer/having.rs @@ -1,6 +1,6 @@ use std::cmp::Ordering; -use pgdog_postgres_types::{Datum, Numeric}; +use pgdog_postgres_types::{Datum, Double, Numeric}; use crate::{ frontend::router::parser::{CompareOperator, HavingExpr, HavingValue}, @@ -103,7 +103,7 @@ fn compare(left: &Datum, operator: CompareOperator, right: &Datum) -> Result Result Option { - match (left, right) { - (Datum::Text(_), Datum::Text(_)) - | (Datum::Interval(_), Datum::Interval(_)) - | (Datum::Timestamp(_), Datum::Timestamp(_)) - | (Datum::TimestampTz(_), Datum::TimestampTz(_)) - | (Datum::Uuid(_), Datum::Uuid(_)) - | (Datum::Boolean(_), Datum::Boolean(_)) - | (Datum::Oid(_), Datum::Oid(_)) => left.partial_cmp(right), - _ => None, - } -} - fn numeric_ordering(left: &Datum, right: &Datum) -> Option { if let (Some(left), Some(right)) = (integer(left), integer(right)) { return Some(left.cmp(&right)); @@ -149,23 +136,12 @@ fn numeric_ordering(left: &Datum, right: &Datum) -> Option { } if let (Some(left), Some(right)) = (float(left), float(right)) { - return Some(postgres_float_ordering(left, right)); + return Some(Double(left).cmp(&Double(right))); } None } -fn postgres_float_ordering(left: f64, right: f64) -> Ordering { - match (left.is_nan(), right.is_nan()) { - (true, true) => Ordering::Equal, - (true, false) => Ordering::Greater, - (false, true) => Ordering::Less, - (false, false) => left - .partial_cmp(&right) - .expect("non-NaN floats always have an ordering"), - } -} - fn integer(value: &Datum) -> Option { match value { Datum::Bigint(value) => Some((*value).into()), @@ -206,7 +182,7 @@ mod tests { messages::{DataRow, Datum}, }, }; - use pgdog_postgres_types::{Double, Float}; + use pgdog_postgres_types::Float; fn compare_expression(operator: CompareOperator) -> HavingExpr { HavingExpr::Compare { @@ -384,8 +360,17 @@ mod tests { numeric_ordering(&Datum::Double(Double(f64::NAN)), &Datum::Float(Float(1.0))), Some(Ordering::Greater) ); - assert_eq!(postgres_float_ordering(f64::NAN, f64::NAN), Ordering::Equal); - assert_eq!(postgres_float_ordering(1.0, f64::NAN), Ordering::Less); + assert_eq!( + numeric_ordering( + &Datum::Double(Double(f64::NAN)), + &Datum::Double(Double(f64::NAN)) + ), + Some(Ordering::Equal) + ); + assert_eq!( + numeric_ordering(&Datum::Float(Float(1.0)), &Datum::Double(Double(f64::NAN))), + Some(Ordering::Less) + ); assert!(matches!( compare( &Datum::Text("2".into()), diff --git a/pgdog/src/backend/pool/connection/multi_shard/mod.rs b/pgdog/src/backend/pool/connection/multi_shard/mod.rs index adcdac68b..3942b59f9 100644 --- a/pgdog/src/backend/pool/connection/multi_shard/mod.rs +++ b/pgdog/src/backend/pool/connection/multi_shard/mod.rs @@ -246,20 +246,18 @@ impl MultiShard { if !self.buffer.is_empty() { // Helpers remain in the internal row through aggregation and // sorting, then are removed before client-visible operations. - if let Err(err) = self.buffer.aggregate( - self.route.aggregate(), - &self.decoder, - &self.route.projection_rewrite_plan, - ) { - // Probably an unsupported aggregation feature was requested. - self.buffer.clear(); - return Err(err.into()); - } - - if let Err(err) = self + let result = self .buffer - .having(&self.route.projection_rewrite_plan, &self.decoder) - { + .aggregate( + self.route.aggregate(), + &self.decoder, + &self.route.projection_rewrite_plan, + ) + .and_then(|_| { + self.buffer + .having(&self.route.projection_rewrite_plan, &self.decoder) + }); + if let Err(err) = result { self.buffer.clear(); return Err(err.into()); } diff --git a/pgdog/src/frontend/router/parser/aggregate.rs b/pgdog/src/frontend/router/parser/aggregate.rs index 27bec68ee..d26dbda38 100644 --- a/pgdog/src/frontend/router/parser/aggregate.rs +++ b/pgdog/src/frontend/router/parser/aggregate.rs @@ -39,6 +39,23 @@ pub(crate) enum AggregateFunction { Unrecognized(String), } +impl AggregateFunction { + pub(crate) fn from_builtin_name(name: &str) -> Option { + Some(match name { + "count" => Self::Count, + "max" => Self::Max, + "min" => Self::Min, + "sum" => Self::Sum, + "avg" => Self::Avg, + "stddev" | "stddev_samp" => Self::StddevSamp, + "stddev_pop" => Self::StddevPop, + "variance" | "var_samp" => Self::VarSamp, + "var_pop" => Self::VarPop, + _ => return None, + }) + } +} + impl fmt::Display for AggregateFunction { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> { match self { @@ -120,24 +137,12 @@ impl Aggregate { .enumerate() .filter_map(|(idx, node)| { let func = Function::try_from(node.val()).ok()?; - let function = match func.name { - "count" => AggregateFunction::Count, - "max" => AggregateFunction::Max, - "min" => AggregateFunction::Min, - "sum" => AggregateFunction::Sum, - "avg" => AggregateFunction::Avg, - "stddev" | "stddev_samp" => AggregateFunction::StddevSamp, - "stddev_pop" => AggregateFunction::StddevPop, - "variance" | "var_samp" => AggregateFunction::VarSamp, - "var_pop" => AggregateFunction::VarPop, - fname => { - if schema.aggregate_functions.contains(fname) { - AggregateFunction::Unrecognized(fname.to_owned()) - } else { - return None; - } - } - }; + let function = AggregateFunction::from_builtin_name(func.name).or_else(|| { + schema + .aggregate_functions + .contains(func.name) + .then(|| AggregateFunction::Unrecognized(func.name.to_owned())) + })?; // FIXME: This doesn't correctly handle distinct behind type // casts any other nodes that Function::try_from handles diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/error.rs b/pgdog/src/frontend/router/parser/rewrite/statement/error.rs index f13f3808a..5a0407d87 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/error.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/error.rs @@ -15,7 +15,7 @@ pub(crate) enum Error { UnsupportedShardingKeyUpdate(String), #[error("cross-shard HAVING is unsupported: {0}")] - UnsupportedHaving(String), + UnsupportedHaving(&'static str), #[error("net: {0}")] Net(#[from] crate::net::Error), diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/having.rs b/pgdog/src/frontend/router/parser/rewrite/statement/having.rs index c54514c53..d689edb93 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/having.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/having.rs @@ -6,7 +6,7 @@ use itertools::Itertools; use pg_raw_parse::{ConstValue, Node, make, nodes}; use crate::frontend::router::parser::{ - CompareOperator, Function, HavingExpr, HavingRewritePlan, HavingValue, + AggregateFunction, CompareOperator, Function, HavingExpr, HavingRewritePlan, HavingValue, }; use super::Error; @@ -47,9 +47,9 @@ pub(super) fn rewrite_select<'a>( ) -> Result { let having = select.having_clause(); if matches!(having, Node::None) { - return Err(Error::UnsupportedHaving(String::from( + return Err(Error::UnsupportedHaving( "HAVING rewrite requested without a HAVING clause", - ))); + )); } let having_source = mem.make_unique(having); let having = having_source.as_ref(); @@ -84,9 +84,9 @@ pub(super) fn rewrite_select<'a>( } Node::A_Const(_) => {} _ => { - return Err(Error::UnsupportedHaving(String::from( + return Err(Error::UnsupportedHaving( "GROUP BY expressions other than columns and ordinals are unsupported", - ))); + )); } } } @@ -186,9 +186,9 @@ impl<'a> Parser<'a> { .and_then(Node::as_str) .and_then(Self::comparison_operator) .ok_or_else(|| { - Error::UnsupportedHaving(String::from( + Error::UnsupportedHaving( "only =, <>, !=, <, <=, > and >= comparisons are supported", - )) + ) })?; Ok(HavingTemplate::Compare { left: self.parse_value(expression.lexpr())?, @@ -206,20 +206,13 @@ impl<'a> Parser<'a> { nodes::BoolExprType::AND_EXPR => Ok(HavingTemplate::And(expressions)), nodes::BoolExprType::OR_EXPR => Ok(HavingTemplate::Or(expressions)), nodes::BoolExprType::NOT_EXPR => { - let mut expressions = expressions.into_iter(); - let expression = expressions.next().ok_or_else(|| { - Error::UnsupportedHaving(String::from("NOT requires one expression")) - })?; - if expressions.next().is_some() { - return Err(Error::UnsupportedHaving(String::from( - "NOT requires one expression", - ))); - } + let expression = expressions + .into_iter() + .exactly_one() + .map_err(|_| Error::UnsupportedHaving("NOT requires one expression"))?; Ok(HavingTemplate::Not(Box::new(expression))) } - _ => Err(Error::UnsupportedHaving(String::from( - "unsupported boolean expression", - ))), + _ => Err(Error::UnsupportedHaving("unsupported boolean expression")), } } Node::NullTest(test) => { @@ -227,9 +220,9 @@ impl<'a> Parser<'a> { nodes::NullTestType::IS_NULL => false, nodes::NullTestType::IS_NOT_NULL => true, _ => { - return Err(Error::UnsupportedHaving(String::from( + return Err(Error::UnsupportedHaving( "only IS NULL and IS NOT NULL are supported", - ))); + )); } }; Ok(HavingTemplate::IsNull { @@ -237,23 +230,17 @@ impl<'a> Parser<'a> { negated, }) } - _ => Err(Error::UnsupportedHaving(String::from( + _ => Err(Error::UnsupportedHaving( "supported predicates are comparisons, AND, OR, NOT and NULL tests", - ))), + )), } } fn parse_value(&mut self, node: Node<'a>) -> Result { - let supported = match node { - Node::A_Const(_) | Node::ParamRef(_) | Node::ColumnRef(_) => true, - Node::TypeCast(cast) => Self::supported_value(cast.arg()), - Node::FuncCall(_) => Self::supported_aggregate(node), - _ => false, - }; - if !supported { - return Err(Error::UnsupportedHaving(String::from( + if !Self::supported_value(node) { + return Err(Error::UnsupportedHaving( "operands must be aggregate calls, grouped columns, constants or parameters", - ))); + )); } let alias = self.next_alias(); @@ -271,22 +258,8 @@ impl<'a> Parser<'a> { } fn supported_aggregate(node: Node<'_>) -> bool { - Function::try_from(node).is_ok_and(|function| { - matches!( - function.name, - "count" - | "max" - | "min" - | "sum" - | "avg" - | "stddev" - | "stddev_samp" - | "stddev_pop" - | "variance" - | "var_samp" - | "var_pop" - ) - }) + Function::try_from(node) + .is_ok_and(|function| AggregateFunction::from_builtin_name(function.name).is_some()) } fn comparison_operator(operator: &str) -> Option { diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs b/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs index 5ad982245..75ff85f34 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs @@ -58,16 +58,20 @@ impl OrderByHelper { pub(crate) struct ProjectionRewritePlan { pub(crate) aggregate_helpers: Vec, pub(crate) order_by_helpers: Vec, - pub(crate) having: Option>, + pub(crate) having: Option>, } +/// The filter and aggregate metadata produced by rewriting a HAVING query. +/// +/// Keeping them together behind one `Arc` avoids storing HAVING only aggregate +/// metadata inline on the common query path. #[derive(Debug, Clone, PartialEq)] -pub(crate) struct HavingPostRoutePlan { +pub(crate) struct HavingRewriteResult { rewrite: HavingRewritePlan, aggregate: Aggregate, } -impl HavingPostRoutePlan { +impl HavingRewriteResult { pub(crate) fn new(rewrite: HavingRewritePlan, aggregate: Aggregate) -> Self { Self { rewrite, aggregate } } @@ -218,10 +222,15 @@ fn build( }; let has_having = !matches!(select.having_clause(), Node::None); - let aggregate = Aggregate::parse(select, schema); - if aggregate.is_empty() && !has_having && select.sort_clause().is_empty() && !rewrite_offset { - return Ok(None); - } + let aggregate = if has_having { + None + } else { + let aggregate = Aggregate::parse(select, schema); + if aggregate.is_empty() && select.sort_clause().is_empty() && !rewrite_offset { + return Ok(None); + } + Some(aggregate) + }; let mut plan = ProjectionRewritePlan::default(); let rewritten = make::try_owned(|mem| { @@ -230,16 +239,18 @@ fn build( let having = having::rewrite_select(&mut select.as_mut(), mem)?; let rewritten_aggregate = Aggregate::parse(&select, schema); if rewritten_aggregate.is_empty() { - return Err(Error::UnsupportedHaving(String::from( + return Err(Error::UnsupportedHaving( "at least one aggregate function is required", - ))); + )); } + // Aggregate::parse skips GROUP BY items it cannot resolve. Require + // every item here so unrelated groups are never merged together. if rewritten_aggregate.group_by().len() != select.group_clause().len() { - return Err(Error::UnsupportedHaving(String::from( + return Err(Error::UnsupportedHaving( "every GROUP BY expression must be a selected column or ordinal", - ))); + )); } - plan.having = Some(Arc::new(HavingPostRoutePlan::new( + plan.having = Some(Arc::new(HavingRewriteResult::new( having, rewritten_aggregate, ))); @@ -247,8 +258,9 @@ fn build( let aggregate = plan .having .as_deref() - .map(HavingPostRoutePlan::aggregate) - .unwrap_or(&aggregate); + .map(HavingRewriteResult::aggregate) + .or(aggregate.as_ref()) + .expect("non-HAVING aggregate parsed before projection rewrite"); if !aggregate.is_empty() { let aggregate_plan = AggregatesRewrite::rewrite_select(&mut select.as_mut(), mem, aggregate); From 467946d2b254e91b7f58ef31e8ba9335a52db5d0 Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Tue, 20 Oct 2026 00:18:43 +0530 Subject: [PATCH 06/24] tests --- integration/rust/tests/integration/having.rs | 7 +- pgdog/src/backend/pool/connection/buffer.rs | 7 +- .../backend/pool/connection/buffer/having.rs | 165 +++++++++++------- .../client/query_engine/route_query.rs | 21 ++- .../test/unsupported_aggregation.rs | 148 ++++++++-------- .../router/parser/rewrite/statement/having.rs | 19 +- .../parser/rewrite/statement/projection.rs | 37 ---- 7 files changed, 216 insertions(+), 188 deletions(-) diff --git a/integration/rust/tests/integration/having.rs b/integration/rust/tests/integration/having.rs index a43d95338..4fd80c360 100644 --- a/integration/rust/tests/integration/having.rs +++ b/integration/rust/tests/integration/having.rs @@ -15,7 +15,7 @@ async fn having_filters_merged_groups() -> Result<(), Box seed( &sharded, 1, - "INSERT INTO having_merge_test VALUES (1, 20), (2, 21), (2, 22)", + "INSERT INTO having_merge_test VALUES (1, 20), (2, 21)", ) .await?; @@ -29,11 +29,10 @@ async fn having_filters_merged_groups() -> Result<(), Box ) .await?; - assert_eq!(rows.len(), 2); + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].columns().len(), 2); assert_eq!(rows[0].get::("customer_id"), 1); assert_eq!(rows[0].get::("order_count"), 3); - assert_eq!(rows[1].get::("customer_id"), 2); - assert_eq!(rows[1].get::("order_count"), 3); cleanup(&sharded).await; Ok(()) diff --git a/pgdog/src/backend/pool/connection/buffer.rs b/pgdog/src/backend/pool/connection/buffer.rs index 2bdb04813..32555a520 100644 --- a/pgdog/src/backend/pool/connection/buffer.rs +++ b/pgdog/src/backend/pool/connection/buffer.rs @@ -167,10 +167,11 @@ impl Buffer { let Some(having) = plan.having.as_deref() else { return Ok(()); }; + let evaluator = having::Evaluator::new(having.expression(), decoder); let mut filtered = VecDeque::with_capacity(self.buffer.len()); for row in self.buffer.drain(..) { - if having::evaluate(having.expression(), &row, decoder)? { + if evaluator.evaluate(&row, decoder)? { filtered.push_back(row); } } @@ -375,10 +376,6 @@ mod test { Field::bigint("__pgdog_having_0"), Field::bigint("__pgdog_having_1"), ])); - let mut noop = Buffer::default(); - noop.having(&ProjectionRewritePlan::default(), &decoder) - .unwrap(); - assert!(noop.is_empty()); let plan = ProjectionRewritePlan { having: Some(std::sync::Arc::new(HavingRewriteResult::new( diff --git a/pgdog/src/backend/pool/connection/buffer/having.rs b/pgdog/src/backend/pool/connection/buffer/having.rs index 91261b25c..77f186768 100644 --- a/pgdog/src/backend/pool/connection/buffer/having.rs +++ b/pgdog/src/backend/pool/connection/buffer/having.rs @@ -1,4 +1,4 @@ -use std::cmp::Ordering; +use std::{cmp::Ordering, collections::HashMap}; use pgdog_postgres_types::{Datum, Double, Numeric}; @@ -16,6 +16,11 @@ enum TruthValue { Unknown, } +pub(super) struct Evaluator<'a> { + expression: &'a HavingExpr, + columns: HashMap<&'a str, usize>, +} + impl TruthValue { fn not(self) -> Self { match self { @@ -26,75 +31,101 @@ impl TruthValue { } } -pub(super) fn evaluate( - expression: &HavingExpr, - row: &DataRow, - decoder: &Decoder, -) -> Result { - Ok(matches!( - evaluate_expression(expression, row, decoder)?, - TruthValue::True - )) -} - -fn evaluate_expression( - expression: &HavingExpr, - row: &DataRow, - decoder: &Decoder, -) -> Result { - match expression { - HavingExpr::Compare { - left, - operator, - right, - } => { - let left = value(left, row, decoder)?; - let right = value(right, row, decoder)?; - compare(&left, *operator, &right) +impl<'a> Evaluator<'a> { + pub(super) fn new(expression: &'a HavingExpr, decoder: &'a Decoder) -> Self { + let columns = decoder + .row_description() + .fields + .iter() + .enumerate() + .map(|(index, field)| (field.name.as_str(), index)) + .collect(); + Self { + expression, + columns, } - HavingExpr::And(expressions) => { - let mut result = TruthValue::True; - for expression in expressions { - match evaluate_expression(expression, row, decoder)? { - TruthValue::False => return Ok(TruthValue::False), - TruthValue::Unknown => result = TruthValue::Unknown, - TruthValue::True => {} + } + + pub(super) fn evaluate(&self, row: &DataRow, decoder: &Decoder) -> Result { + Ok(matches!( + self.evaluate_expression(self.expression, row, decoder)?, + TruthValue::True + )) + } + + fn evaluate_expression( + &self, + expression: &HavingExpr, + row: &DataRow, + decoder: &Decoder, + ) -> Result { + match expression { + HavingExpr::Compare { + left, + operator, + right, + } => { + let left = self.value(left, row, decoder)?; + let right = self.value(right, row, decoder)?; + compare(&left, *operator, &right) + } + HavingExpr::And(expressions) => { + let mut result = TruthValue::True; + for expression in expressions { + match self.evaluate_expression(expression, row, decoder)? { + TruthValue::False => return Ok(TruthValue::False), + TruthValue::Unknown => result = TruthValue::Unknown, + TruthValue::True => {} + } } + Ok(result) } - Ok(result) - } - HavingExpr::Or(expressions) => { - let mut result = TruthValue::False; - for expression in expressions { - match evaluate_expression(expression, row, decoder)? { - TruthValue::True => return Ok(TruthValue::True), - TruthValue::Unknown => result = TruthValue::Unknown, - TruthValue::False => {} + HavingExpr::Or(expressions) => { + let mut result = TruthValue::False; + for expression in expressions { + match self.evaluate_expression(expression, row, decoder)? { + TruthValue::True => return Ok(TruthValue::True), + TruthValue::Unknown => result = TruthValue::Unknown, + TruthValue::False => {} + } } + Ok(result) + } + HavingExpr::Not(expression) => { + Ok(self.evaluate_expression(expression, row, decoder)?.not()) + } + HavingExpr::IsNull { value, negated } => { + let is_null = self.value(value, row, decoder)?.is_null(); + Ok(if is_null != *negated { + TruthValue::True + } else { + TruthValue::False + }) } - Ok(result) - } - HavingExpr::Not(expression) => Ok(evaluate_expression(expression, row, decoder)?.not()), - HavingExpr::IsNull { - value: operand, - negated, - } => { - let is_null = value(operand, row, decoder)?.is_null(); - Ok(if is_null != *negated { - TruthValue::True - } else { - TruthValue::False - }) } } + + fn value(&self, value: &HavingValue, row: &DataRow, decoder: &Decoder) -> Result { + let index = + self.columns.get(value.alias()).copied().ok_or_else(|| { + Error::InvalidHaving(format!("missing operand {}", value.alias())) + })?; + Ok(row.get_column_checked(index, decoder)?.value) + } } -fn value(operand: &HavingValue, row: &DataRow, decoder: &Decoder) -> Result { - let index = decoder - .row_description() - .field_index(operand.alias()) - .ok_or_else(|| Error::InvalidHaving(format!("missing operand {}", operand.alias())))?; - Ok(row.get_column_checked(index, decoder)?.value) +#[cfg(test)] +fn evaluate(expression: &HavingExpr, row: &DataRow, decoder: &Decoder) -> Result { + Evaluator::new(expression, decoder).evaluate(row, decoder) +} + +#[cfg(test)] +fn evaluate_expression( + expression: &HavingExpr, + row: &DataRow, + decoder: &Decoder, +) -> Result { + Evaluator::new(expression, decoder).evaluate_expression(expression, row, decoder) } fn compare(left: &Datum, operator: CompareOperator, right: &Datum) -> Result { @@ -205,13 +236,14 @@ mod tests { } #[test] - fn compares_merged_bigint_with_integer_without_losing_precision() { + fn compares_merged_bigint_with_numeric_without_losing_precision() { let decoder = Decoder::from(RowDescription::new(&[ Field::bigint("__pgdog_having_0"), - integer_field("__pgdog_having_1"), + Field::numeric("__pgdog_having_1"), ])); let mut row = DataRow::new(); - row.add(9_007_199_254_740_993_i64).add(2_i32); + row.add(9_007_199_254_740_993_i64) + .add(Numeric::from(9_007_199_254_740_992_i64)); assert!( evaluate( @@ -287,7 +319,10 @@ mod tests { .unwrap() ); assert!(evaluate(&is_not_null, &row, &decoder).unwrap()); - assert!(!evaluate(&HavingExpr::Not(Box::new(unknown)), &row, &decoder).unwrap()); + assert_eq!( + evaluate_expression(&HavingExpr::Not(Box::new(unknown)), &row, &decoder).unwrap(), + TruthValue::Unknown + ); let missing = HavingExpr::Compare { left: HavingValue::new("__pgdog_missing".into()), diff --git a/pgdog/src/frontend/client/query_engine/route_query.rs b/pgdog/src/frontend/client/query_engine/route_query.rs index e79f8df6e..6c603a1b8 100644 --- a/pgdog/src/frontend/client/query_engine/route_query.rs +++ b/pgdog/src/frontend/client/query_engine/route_query.rs @@ -4,6 +4,7 @@ use tracing::trace; use crate::backend::Error as BackendError; use crate::frontend::router::Error as RouterError; use crate::frontend::router::parser::Error as ParserError; +use crate::frontend::router::parser::rewrite::statement::Error as StatementRewriteError; use crate::frontend::router::parser::rewrite::statement::plan::RewriteResult; use crate::frontend::router::parser::rewrite::statement::projection; use crate::frontend::router::sharding::lookup; @@ -154,11 +155,27 @@ impl QueryEngine { client_request.route = Some(command.route().clone()); trace!("routing {:#?} to {:#?}", client_request.messages, command,); - projection::finalize_after_route( + let finalize = projection::finalize_after_route( client_request, &cluster.schema(), rewrite_result.and_then(RewriteResult::offset_plan), - )?; + ); + match finalize { + Err(err @ StatementRewriteError::UnsupportedHaving(_)) => { + if let Some(transaction) = context.transaction.as_mut() { + transaction.set_error(); + } + self.error_response( + context, + client_request, + ErrorResponse::unsupported(&err), + ) + .await?; + return Ok(false); + } + Err(err) => return Err(err.into()), + Ok(()) => {} + } if let Some(rewrite_result) = rewrite_result { rewrite_result.apply_after_route(client_request)?; diff --git a/pgdog/src/frontend/client/query_engine/test/unsupported_aggregation.rs b/pgdog/src/frontend/client/query_engine/test/unsupported_aggregation.rs index fa7af43d2..76d88e14c 100644 --- a/pgdog/src/frontend/client/query_engine/test/unsupported_aggregation.rs +++ b/pgdog/src/frontend/client/query_engine/test/unsupported_aggregation.rs @@ -13,91 +13,101 @@ use crate::{ async fn test_unsupported_aggregation_keeps_connection_usable() { let mut client = SpawnedClient::new_sharded(Parameters::default()).await; - for begin in [None, Some("BEGIN"), Some("BEGIN READ ONLY")] { - timeout(Duration::from_secs(5), async { - if let Some(begin) = begin { - client.send(Query::new(begin)).await; - assert_eq!( - expect_message!(client.read().await, CommandComplete).command(), - "BEGIN" + for (query, expected_error) in [ + ( + "SELECT count(DISTINCT value) FROM sharded", + "unsupported aggregation count: count(DISTINCT ...) is not yet supported", + ), + ( + "SELECT count(*) FROM sharded HAVING count(*) + 1 > 0", + "cross-shard HAVING is unsupported: operands must be aggregate calls, grouped columns, constants or parameters", + ), + ] { + for begin in [None, Some("BEGIN"), Some("BEGIN READ ONLY")] { + timeout(Duration::from_secs(5), async { + if let Some(begin) = begin { + client.send(Query::new(begin)).await; + assert_eq!( + expect_message!(client.read().await, CommandComplete).command(), + "BEGIN" + ); + assert_eq!( + expect_message!(client.read().await, ReadyForQuery).status, + 'T' + ); + } + + client.send(Query::new(query)).await; + + let messages = timeout(Duration::from_secs(2), client.read_until('E')) + .await + .unwrap_or_else(|_| { + panic!("unsupported aggregation did not return an error: {query}") + }); + let error = expect_message!( + messages.last().expect("error response").clone(), + ErrorResponse ); + assert_eq!(error.severity, "ERROR"); + assert_eq!(error.code, "0A000"); + assert_eq!(error.message, expected_error); assert_eq!( - expect_message!(client.read().await, ReadyForQuery).status, - 'T' + expect_message!( + timeout(Duration::from_secs(2), client.read()) + .await + .expect("ReadyForQuery missing after unsupported aggregation"), + ReadyForQuery + ) + .status, + if begin.is_some() { 'E' } else { 'I' } ); - } - - client - .send(Query::new("SELECT count(DISTINCT value) FROM sharded")) - .await; - let messages = timeout(Duration::from_secs(2), client.read_until('E')) - .await - .expect("unsupported aggregation did not return an error"); - let error = expect_message!( - messages.last().expect("error response").clone(), - ErrorResponse - ); - assert_eq!(error.severity, "ERROR"); - assert_eq!(error.code, "0A000"); - assert_eq!( - error.message, - "unsupported aggregation count: count(DISTINCT ...) is not yet supported" - ); - assert_eq!( - expect_message!( - timeout(Duration::from_secs(2), client.read()) - .await - .expect("ReadyForQuery missing after unsupported aggregation"), - ReadyForQuery - ) - .status, - if begin.is_some() { 'E' } else { 'I' } - ); + if begin.is_some() { + // The failed transaction must reject further queries until rollback. + for _ in 0..2 { + client + .send(Query::new("SELECT count(*) FROM sharded WHERE false")) + .await; + let error = expect_message!(client.read().await, ErrorResponse); + assert_eq!(error.severity, "ERROR"); + assert_eq!(error.code, "25P02"); + assert_eq!( + expect_message!(client.read().await, ReadyForQuery).status, + 'E' + ); + } - if begin.is_some() { - // The failed transaction must reject further queries until rollback. - for _ in 0..2 { - client - .send(Query::new("SELECT count(*) FROM sharded WHERE false")) - .await; - let error = expect_message!(client.read().await, ErrorResponse); - assert_eq!(error.severity, "ERROR"); - assert_eq!(error.code, "25P02"); + client.send(Query::new("ROLLBACK")).await; + assert_eq!( + expect_message!(client.read().await, CommandComplete).command(), + "ROLLBACK" + ); assert_eq!( expect_message!(client.read().await, ReadyForQuery).status, - 'E' + 'I' ); } - client.send(Query::new("ROLLBACK")).await; + // Each shard returns a count row, exercising aggregation again without fixtures. + client + .send(Query::new("SELECT count(*) FROM sharded WHERE false")) + .await; + expect_message!(client.read().await, RowDescription); + let row = expect_message!(client.read().await, DataRow); + assert_eq!(row.get_int(0, true), Some(0)); assert_eq!( expect_message!(client.read().await, CommandComplete).command(), - "ROLLBACK" + "SELECT 1" ); assert_eq!( expect_message!(client.read().await, ReadyForQuery).status, 'I' ); - } - - // Each shard returns a count row, exercising aggregation again without fixtures. - client - .send(Query::new("SELECT count(*) FROM sharded WHERE false")) - .await; - expect_message!(client.read().await, RowDescription); - let row = expect_message!(client.read().await, DataRow); - assert_eq!(row.get_int(0, true), Some(0)); - assert_eq!( - expect_message!(client.read().await, CommandComplete).command(), - "SELECT 1" - ); - assert_eq!( - expect_message!(client.read().await, ReadyForQuery).status, - 'I' - ); - }) - .await - .unwrap_or_else(|_| panic!("unsupported aggregation or subsequent query hung: {begin:?}")); + }) + .await + .unwrap_or_else(|_| { + panic!("unsupported aggregation or subsequent query hung: {query}, {begin:?}") + }); + } } } diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/having.rs b/pgdog/src/frontend/router/parser/rewrite/statement/having.rs index d689edb93..d9a015fcf 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/having.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/having.rs @@ -46,11 +46,10 @@ pub(super) fn rewrite_select<'a>( mem: make::MemoryToken<'a>, ) -> Result { let having = select.having_clause(); - if matches!(having, Node::None) { - return Err(Error::UnsupportedHaving( - "HAVING rewrite requested without a HAVING clause", - )); - } + debug_assert!( + !matches!(having, Node::None), + "HAVING rewrite requires a HAVING clause" + ); let having_source = mem.make_unique(having); let having = having_source.as_ref(); @@ -344,7 +343,15 @@ mod tests { ) ) )); - assert!(sql.contains("OR true")); + assert_eq!( + sql, + "SELECT customer_id, sum(price)::numeric, count(*) AS __pgdog_having_0, \ + $1 AS __pgdog_having_1, sum(price) AS __pgdog_having_2, \ + max(price)::numeric AS __pgdog_having_3, 10 AS __pgdog_having_4, \ + customer_id AS __pgdog_having_5 FROM purchases GROUP BY customer_id \ + HAVING (NOT (count(*) <= $1 OR sum(price) IS NOT NULL) AND \ + max(price)::numeric >= 10) OR true" + ); assert_eq!(plan.hidden_aliases().len(), 6); } diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs b/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs index 75ff85f34..db01d2b78 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs @@ -328,43 +328,6 @@ mod tests { assert!(plan.drop_columns(&row_description).is_empty()); } - #[test] - fn having_projects_operands_and_updates_aggregate_plan() { - let ast = pg_raw_parse::parse( - "SELECT customer_id, COUNT(*) AS orders \ - FROM purchases GROUP BY customer_id HAVING COUNT(*) > 2", - ) - .unwrap(); - let rewrite = build(&ast, &Schema::default(), false) - .unwrap() - .expect("HAVING requires a rewrite"); - - assert!(rewrite.sql.contains("__pgdog_having_0")); - assert!(rewrite.sql.contains("__pgdog_having_1")); - assert_eq!( - rewrite - .plan - .having - .as_deref() - .expect("HAVING plan") - .hidden_aliases() - .len(), - 3 - ); - assert_eq!( - rewrite - .plan - .having - .as_deref() - .expect("HAVING plan") - .aggregate() - .targets() - .len(), - 2, - "the selected and hidden COUNT must both be merged" - ); - } - #[test] fn having_internal_aliases_do_not_collide_with_client_aliases() { let ast = pg_raw_parse::parse( From 611347aef5f65e7081d2ff7a58c91348cbf4807e Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Wed, 21 Oct 2026 20:24:19 +0530 Subject: [PATCH 07/24] use option --- .../src/backend/pool/connection/aggregate.rs | 67 +++++++------------ pgdog/src/backend/pool/connection/buffer.rs | 33 +++++---- .../pool/connection/multi_shard/mod.rs | 22 ++---- .../pool/connection/multi_shard/test.rs | 2 +- .../query_engine/test/rewrite_projection.rs | 34 +++++++--- .../parser/rewrite/statement/projection.rs | 48 +++++++------ pgdog/src/frontend/router/parser/route.rs | 11 +-- 7 files changed, 110 insertions(+), 107 deletions(-) diff --git a/pgdog/src/backend/pool/connection/aggregate.rs b/pgdog/src/backend/pool/connection/aggregate.rs index dd48d2e30..9d68d73e0 100644 --- a/pgdog/src/backend/pool/connection/aggregate.rs +++ b/pgdog/src/backend/pool/connection/aggregate.rs @@ -246,7 +246,7 @@ impl<'a> Aggregates<'a> { rows: &'a VecDeque, decoder: &'a Decoder, aggregate: &'a Aggregate, - plan: &ProjectionRewritePlan, + plan: Option<&ProjectionRewritePlan>, ) -> Option { let mut helper_columns: HashMap = HashMap::new(); @@ -263,16 +263,18 @@ impl<'a> Aggregates<'a> { } } - for helper in &plan.aggregate_helpers { - let Some(index) = decoder.row_description().field_index(&helper.alias) else { - continue; - }; + if let Some(plan) = plan { + 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(); - match helper.kind { - HelperKind::Count => entry.count = Some(index), - HelperKind::Sum => entry.sum = Some(index), - HelperKind::SumSquares => entry.sumsq = Some(index), + let entry = helper_columns.entry(helper.target_column).or_default(); + match helper.kind { + HelperKind::Count => entry.count = Some(index), + HelperKind::Sum => entry.sum = Some(index), + HelperKind::SumSquares => entry.sumsq = Some(index), + } } } @@ -524,8 +526,7 @@ mod test { shard1.add("3"); rows.push_back(shard1); - let plan = ProjectionRewritePlan::default(); - let mut result = Aggregates::new(&rows, &decoder, &aggregate, &plan) + let mut result = Aggregates::new(&rows, &decoder, &aggregate, None) .unwrap() .aggregate() .unwrap(); @@ -572,8 +573,7 @@ mod test { ); rows.push_back(row); } - let plan = ProjectionRewritePlan::default(); - let mut result = Aggregates::new(&rows, &decoder, &aggregate, &plan) + let mut result = Aggregates::new(&rows, &decoder, &aggregate, None) .expect("count aggregate") .aggregate() .expect("merge partial counts"); @@ -606,15 +606,10 @@ mod test { shard2.add(20.0_f64).add(4_i64); rows.push_back(shard2); - let mut result = Aggregates::new( - &rows, - &decoder, - &aggregate, - &ProjectionRewritePlan::default(), - ) - .unwrap() - .aggregate() - .unwrap(); + let mut result = Aggregates::new(&rows, &decoder, &aggregate, None) + .unwrap() + .aggregate() + .unwrap(); assert_eq!(result.len(), 2); let mut groups: Vec<(f64, i64)> = result @@ -651,15 +646,10 @@ mod test { shard2.add(Bytes::from_static(b"{{5,6},{7,8}}")).add(1_i64); rows.push_back(shard2); - let mut result = Aggregates::new( - &rows, - &decoder, - &aggregate, - &ProjectionRewritePlan::default(), - ) - .unwrap() - .aggregate() - .unwrap(); + let mut result = Aggregates::new(&rows, &decoder, &aggregate, None) + .unwrap() + .aggregate() + .unwrap(); let mut groups: Vec<(String, i64)> = result .drain(..) @@ -698,15 +688,10 @@ mod test { shard1.add(input.clone()).add(1_i64); rows.push_back(shard1); - let mut result = Aggregates::new( - &rows, - &decoder, - &aggregate, - &ProjectionRewritePlan::default(), - ) - .unwrap() - .aggregate() - .unwrap(); + let mut result = Aggregates::new(&rows, &decoder, &aggregate, None) + .unwrap() + .aggregate() + .unwrap(); assert_eq!(result.len(), 1); let row = result.pop_front().unwrap(); diff --git a/pgdog/src/backend/pool/connection/buffer.rs b/pgdog/src/backend/pool/connection/buffer.rs index 32555a520..609387bca 100644 --- a/pgdog/src/backend/pool/connection/buffer.rs +++ b/pgdog/src/backend/pool/connection/buffer.rs @@ -142,7 +142,7 @@ impl Buffer { &mut self, aggregate: &Aggregate, decoder: &Decoder, - plan: &ProjectionRewritePlan, + plan: Option<&ProjectionRewritePlan>, ) -> Result<(), super::Error> { let buffer: VecDeque = std::mem::take(&mut self.buffer); let rows = if aggregate.is_empty() { @@ -161,10 +161,13 @@ impl Buffer { /// Apply `HAVING` after all partial aggregate rows have been merged. pub(super) fn having( &mut self, - plan: &ProjectionRewritePlan, + plan: Option<&ProjectionRewritePlan>, decoder: &Decoder, ) -> Result<(), super::Error> { - let Some(having) = plan.having.as_deref() else { + let Some(plan) = plan else { + return Ok(()); + }; + let Some(having) = plan.having.as_ref() else { return Ok(()); }; let evaluator = having::Evaluator::new(having.expression(), decoder); @@ -179,10 +182,14 @@ impl Buffer { Ok(()) } - pub(super) fn drop_helper_columns(&mut self, plan: &ProjectionRewritePlan, decoder: &Decoder) { - if plan.is_noop() { + pub(super) fn drop_helper_columns( + &mut self, + plan: Option<&ProjectionRewritePlan>, + decoder: &Decoder, + ) { + let Some(plan) = plan else { return; - } + }; let drop = plan.drop_columns(decoder.row_description()); @@ -318,8 +325,7 @@ mod test { buf.add(dr.message()).unwrap(); } - buf.aggregate(&agg, &Decoder::from(rd), &ProjectionRewritePlan::default()) - .unwrap(); + buf.aggregate(&agg, &Decoder::from(rd), None).unwrap(); buf.mark_full(); assert_eq!(buf.len(), 1); @@ -345,8 +351,7 @@ mod test { } } - buf.aggregate(&agg, &Decoder::from(rd), &ProjectionRewritePlan::default()) - .unwrap(); + buf.aggregate(&agg, &Decoder::from(rd), None).unwrap(); buf.mark_full(); assert_eq!(buf.len(), 2); @@ -378,7 +383,7 @@ mod test { ])); let plan = ProjectionRewritePlan { - having: Some(std::sync::Arc::new(HavingRewriteResult::new( + having: Some(HavingRewriteResult::new( crate::frontend::router::parser::HavingRewritePlan::new( HavingExpr::Compare { left: HavingValue::new("__pgdog_having_0".into()), @@ -388,7 +393,7 @@ mod test { vec!["__pgdog_having_0".into(), "__pgdog_having_1".into()], ), aggregate.clone(), - ))), + )), ..Default::default() }; @@ -402,8 +407,8 @@ mod test { buffer.add(row.message()).unwrap(); } - buffer.aggregate(&aggregate, &decoder, &plan).unwrap(); - buffer.having(&plan, &decoder).unwrap(); + buffer.aggregate(&aggregate, &decoder, Some(&plan)).unwrap(); + buffer.having(Some(&plan), &decoder).unwrap(); assert_eq!(buffer.len(), 1); assert_eq!(buffer.buffer[0].get::(0, Format::Text).unwrap(), 1); diff --git a/pgdog/src/backend/pool/connection/multi_shard/mod.rs b/pgdog/src/backend/pool/connection/multi_shard/mod.rs index 3942b59f9..32da11c96 100644 --- a/pgdog/src/backend/pool/connection/multi_shard/mod.rs +++ b/pgdog/src/backend/pool/connection/multi_shard/mod.rs @@ -246,25 +246,18 @@ impl MultiShard { if !self.buffer.is_empty() { // Helpers remain in the internal row through aggregation and // sorting, then are removed before client-visible operations. + let plan = self.route.projection_rewrite_plan.as_deref(); let result = self .buffer - .aggregate( - self.route.aggregate(), - &self.decoder, - &self.route.projection_rewrite_plan, - ) - .and_then(|_| { - self.buffer - .having(&self.route.projection_rewrite_plan, &self.decoder) - }); + .aggregate(self.route.aggregate(), &self.decoder, plan) + .and_then(|_| self.buffer.having(plan, &self.decoder)); if let Err(err) = result { self.buffer.clear(); return Err(err.into()); } self.buffer.sort(self.route.order_by(), &self.decoder); - self.buffer - .drop_helper_columns(&self.route.projection_rewrite_plan, &self.decoder); + self.buffer.drop_helper_columns(plan, &self.decoder); self.buffer.distinct(self.route.distinct(), &self.decoder); self.buffer.limit(self.route.limit()); } @@ -307,12 +300,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.projection_rewrite_plan; - if plan.is_noop() { - forward = Some(message); - } else { + if let Some(plan) = self.route.projection_rewrite_plan.as_deref() { let client_rd = rd.drop_columns(plan.drop_columns(&rd)); forward = Some(client_rd.message()); + } else { + forward = Some(message); } // The next statement describes a different result set. diff --git a/pgdog/src/backend/pool/connection/multi_shard/test.rs b/pgdog/src/backend/pool/connection/multi_shard/test.rs index 1aaf6805d..f6cc5ccfa 100644 --- a/pgdog/src/backend/pool/connection/multi_shard/test.rs +++ b/pgdog/src/backend/pool/connection/multi_shard/test.rs @@ -79,7 +79,7 @@ fn test_order_by_helper_after_star_expansion_is_dropped_after_sorting() { Default::default(), None, ); - route.projection_rewrite_plan = plan; + route.projection_rewrite_plan = Some(std::sync::Arc::new(plan)); let mut multi_shard = MultiShard::new(2, &route); let row_description = RowDescription::new(&[ diff --git a/pgdog/src/frontend/client/query_engine/test/rewrite_projection.rs b/pgdog/src/frontend/client/query_engine/test/rewrite_projection.rs index 394846a4a..20ed6bd36 100644 --- a/pgdog/src/frontend/client/query_engine/test/rewrite_projection.rs +++ b/pgdog/src/frontend/client/query_engine/test/rewrite_projection.rs @@ -55,7 +55,7 @@ async fn direct_aggregate_keeps_base_sql() { _ => panic!("expected Query"), }; assert_eq!(query.query(), sql); - assert!(client_request.route().projection_rewrite_plan.is_noop()); + assert!(client_request.route().projection_rewrite_plan.is_none()); } #[tokio::test] @@ -89,6 +89,8 @@ async fn cross_shard_aggregate_adds_and_tracks_helpers() { client_request .route() .projection_rewrite_plan + .as_deref() + .expect("projection rewrite plan") .aggregate_helpers .len(), 1 @@ -131,6 +133,8 @@ async fn cross_shard_having_rewrites_and_replaces_route_aggregate() { client_request .route() .projection_rewrite_plan + .as_deref() + .expect("projection rewrite plan") .having .is_some() ); @@ -163,7 +167,7 @@ async fn direct_having_keeps_native_postgres_semantics() { _ => panic!("expected Query"), }; assert_eq!(query.query(), sql); - assert!(client_request.route().projection_rewrite_plan.is_noop()); + assert!(client_request.route().projection_rewrite_plan.is_none()); } #[tokio::test] @@ -326,6 +330,8 @@ async fn cross_shard_order_by_projects_missing_sort_column() { client_request .route() .projection_rewrite_plan + .as_deref() + .expect("projection rewrite plan") .order_by_helpers .len(), 1 @@ -405,7 +411,15 @@ fn aliased_projected_sort_column_remaps_route() { request.route().order_by(), &[OrderBy::AscColumn("item_price".into())] ); - assert!(!request.route().projection_rewrite_plan.order_by_helpers[0].injected); + assert!( + !request + .route() + .projection_rewrite_plan + .as_deref() + .expect("projection rewrite plan") + .order_by_helpers[0] + .injected + ); } #[test] @@ -508,14 +522,12 @@ async fn aggregate_order_by_and_offset_compose_after_route() { route.order_by(), &[OrderBy::AscColumn("__pgdog_order_col0".into())] ); - assert_eq!( - route.projection_rewrite_plan.aggregate_helpers[0].alias, - "__pgdog_count_col0" - ); - assert_eq!( - route.projection_rewrite_plan.order_by_helpers[0].alias, - "__pgdog_order_col0" - ); + let plan = route + .projection_rewrite_plan + .as_deref() + .expect("projection rewrite plan"); + assert_eq!(plan.aggregate_helpers[0].alias, "__pgdog_count_col0"); + assert_eq!(plan.order_by_helpers[0].alias, "__pgdog_order_col0"); assert_eq!( route.limit(), &Limit { diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs b/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs index db01d2b78..eec7a7933 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs @@ -54,18 +54,15 @@ impl OrderByHelper { } } -#[derive(Debug, Clone, Default, PartialEq)] +#[derive(Debug, Default, PartialEq)] pub(crate) struct ProjectionRewritePlan { pub(crate) aggregate_helpers: Vec, pub(crate) order_by_helpers: Vec, - pub(crate) having: Option>, + pub(crate) having: Option, } /// The filter and aggregate metadata produced by rewriting a HAVING query. -/// -/// Keeping them together behind one `Arc` avoids storing HAVING only aggregate -/// metadata inline on the common query path. -#[derive(Debug, Clone, PartialEq)] +#[derive(Debug, PartialEq)] pub(crate) struct HavingRewriteResult { rewrite: HavingRewritePlan, aggregate: Aggregate, @@ -119,7 +116,7 @@ impl ProjectionRewritePlan { #[derive(Debug, Clone)] pub(crate) struct PostRouteRewrite { sql: Arc, - plan: ProjectionRewritePlan, + plan: Option>, } pub(crate) fn finalize_after_route( @@ -181,15 +178,15 @@ pub(crate) fn finalize_after_route( { parse.set_query(&rewrite.sql); } - if !rewrite.plan.is_noop() + if let Some(plan) = &rewrite.plan && let Some(route) = request.route.as_mut() { - route.projection_rewrite_plan = rewrite.plan.clone(); - if let Some(having) = &rewrite.plan.having { + route.projection_rewrite_plan = Some(plan.clone()); + if let Some(having) = &plan.having { route.set_aggregate(having.aggregate().clone()); } let mut order_by = route.order_by().to_vec(); - for helper in &rewrite.plan.order_by_helpers { + for helper in &plan.order_by_helpers { // Prefer the structural position. The source fallback handles a // bind-dependent vector sort omitted from this execution's route. let position = order_by @@ -250,14 +247,11 @@ fn build( "every GROUP BY expression must be a selected column or ordinal", )); } - plan.having = Some(Arc::new(HavingRewriteResult::new( - having, - rewritten_aggregate, - ))); + plan.having = Some(HavingRewriteResult::new(having, rewritten_aggregate)); } let aggregate = plan .having - .as_deref() + .as_ref() .map(HavingRewriteResult::aggregate) .or(aggregate.as_ref()) .expect("non-HAVING aggregate parsed before projection rewrite"); @@ -276,6 +270,7 @@ fn build( return Ok(None); } let sql: Arc = pg_raw_parse::deparse(&*rewritten)?.as_str().into(); + let plan = (!plan.is_noop()).then(|| Arc::new(plan)); Ok(Some(PostRouteRewrite { sql, plan })) } @@ -342,8 +337,10 @@ mod tests { assert_eq!( rewrite .plan - .having .as_deref() + .expect("projection plan") + .having + .as_ref() .expect("HAVING plan") .hidden_aliases(), &["__pgdog_having_1", "__pgdog_having_2", "__pgdog_having_3"] @@ -360,7 +357,8 @@ mod tests { let rewrite = build(&ast, &Schema::default(), false) .unwrap() .expect("HAVING requires a rewrite"); - let aggregate = rewrite.plan.having.expect("HAVING plan"); + let plan = rewrite.plan.expect("projection plan"); + let aggregate = plan.having.as_ref().expect("HAVING plan"); assert_eq!(aggregate.aggregate().group_by().len(), 1); assert_eq!(aggregate.aggregate().targets().len(), 2); @@ -380,8 +378,10 @@ mod tests { assert_eq!( rewrite .plan - .having .as_deref() + .expect("projection plan") + .having + .as_ref() .expect("HAVING plan") .aggregate() .targets() @@ -389,7 +389,12 @@ mod tests { 2 ); assert_eq!( - rewrite.plan.aggregate_helpers.len(), + rewrite + .plan + .as_deref() + .expect("projection plan") + .aggregate_helpers + .len(), 2, "selected and hidden AVG each require their own COUNT" ); @@ -405,7 +410,8 @@ mod tests { let rewrite = build(&ast, &Schema::default(), false) .unwrap() .expect("HAVING requires a rewrite"); - let aggregate = rewrite.plan.having.expect("HAVING plan"); + let plan = rewrite.plan.expect("projection plan"); + let aggregate = plan.having.as_ref().expect("HAVING plan"); assert_eq!(aggregate.aggregate().targets().len(), 1); assert_eq!(aggregate.aggregate().group_by().len(), 1); diff --git a/pgdog/src/frontend/router/parser/route.rs b/pgdog/src/frontend/router/parser/route.rs index c3859c4b5..b270058e1 100644 --- a/pgdog/src/frontend/router/parser/route.rs +++ b/pgdog/src/frontend/router/parser/route.rs @@ -1,4 +1,4 @@ -use std::{fmt::Display, ops::Deref}; +use std::{fmt::Display, ops::Deref, sync::Arc}; use super::{ Aggregate, DistinctBy, Limit, OrderBy, StatementType, advisory_lock::AdvisoryLocks, @@ -119,8 +119,8 @@ pub(crate) struct Route { advisory_locks: AdvisoryLocks, /// `DISTINCT` clause, if set. distinct: Option, - /// Temporary columns projected for cross-shard result processing. - pub(crate) projection_rewrite_plan: ProjectionRewritePlan, + /// Shared metadata for temporary cross-shard projections, absent on the common path. + pub(crate) projection_rewrite_plan: Option>, /// Our query explain plan. We attach /// this to the `EXPLAIN` output. explain: Option, @@ -347,7 +347,10 @@ impl Route { pub(crate) fn requires_post_processing(&self) -> bool { !self.order_by().is_empty() || !self.aggregate().is_empty() - || self.projection_rewrite_plan.having.is_some() + || self + .projection_rewrite_plan + .as_ref() + .is_some_and(|plan| plan.having.is_some()) || self.distinct().is_some() || self.limit().limit.is_some() || self.limit().offset.is_some() From 7f92f25caf4e23d489c5c4f77d5030ead0de88a1 Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Wed, 21 Oct 2026 21:05:54 +0530 Subject: [PATCH 08/24] redo argh --- pgdog/src/backend/pool/connection/buffer.rs | 4 ++-- .../router/parser/rewrite/statement/projection.rs | 9 ++++++--- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/pgdog/src/backend/pool/connection/buffer.rs b/pgdog/src/backend/pool/connection/buffer.rs index 609387bca..12b80fd8e 100644 --- a/pgdog/src/backend/pool/connection/buffer.rs +++ b/pgdog/src/backend/pool/connection/buffer.rs @@ -383,7 +383,7 @@ mod test { ])); let plan = ProjectionRewritePlan { - having: Some(HavingRewriteResult::new( + having: Some(Box::new(HavingRewriteResult::new( crate::frontend::router::parser::HavingRewritePlan::new( HavingExpr::Compare { left: HavingValue::new("__pgdog_having_0".into()), @@ -393,7 +393,7 @@ mod test { vec!["__pgdog_having_0".into(), "__pgdog_having_1".into()], ), aggregate.clone(), - )), + ))), ..Default::default() }; diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs b/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs index eec7a7933..c786e1d54 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs @@ -58,7 +58,7 @@ impl OrderByHelper { pub(crate) struct ProjectionRewritePlan { pub(crate) aggregate_helpers: Vec, pub(crate) order_by_helpers: Vec, - pub(crate) having: Option, + pub(crate) having: Option>, } /// The filter and aggregate metadata produced by rewriting a HAVING query. @@ -247,11 +247,14 @@ fn build( "every GROUP BY expression must be a selected column or ordinal", )); } - plan.having = Some(HavingRewriteResult::new(having, rewritten_aggregate)); + plan.having = Some(Box::new(HavingRewriteResult::new( + having, + rewritten_aggregate, + ))); } let aggregate = plan .having - .as_ref() + .as_deref() .map(HavingRewriteResult::aggregate) .or(aggregate.as_ref()) .expect("non-HAVING aggregate parsed before projection rewrite"); From 9c6fea4c82ba43d21ded15bba8b75cee8420ccd9 Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Tue, 27 Oct 2026 14:10:17 +0530 Subject: [PATCH 09/24] wrap option --- .../src/backend/pool/connection/aggregate.rs | 67 +++++++------------ pgdog/src/backend/pool/connection/buffer.rs | 18 ++--- .../pool/connection/multi_shard/mod.rs | 20 +++--- .../pool/connection/multi_shard/test.rs | 2 +- .../query_engine/test/rewrite_projection.rs | 30 ++++++--- pgdog/src/frontend/router/parser/aggregate.rs | 41 +++++++----- .../parser/rewrite/statement/projection.rs | 11 +-- pgdog/src/frontend/router/parser/route.rs | 6 +- 8 files changed, 98 insertions(+), 97 deletions(-) diff --git a/pgdog/src/backend/pool/connection/aggregate.rs b/pgdog/src/backend/pool/connection/aggregate.rs index dd48d2e30..9d68d73e0 100644 --- a/pgdog/src/backend/pool/connection/aggregate.rs +++ b/pgdog/src/backend/pool/connection/aggregate.rs @@ -246,7 +246,7 @@ impl<'a> Aggregates<'a> { rows: &'a VecDeque, decoder: &'a Decoder, aggregate: &'a Aggregate, - plan: &ProjectionRewritePlan, + plan: Option<&ProjectionRewritePlan>, ) -> Option { let mut helper_columns: HashMap = HashMap::new(); @@ -263,16 +263,18 @@ impl<'a> Aggregates<'a> { } } - for helper in &plan.aggregate_helpers { - let Some(index) = decoder.row_description().field_index(&helper.alias) else { - continue; - }; + if let Some(plan) = plan { + 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(); - match helper.kind { - HelperKind::Count => entry.count = Some(index), - HelperKind::Sum => entry.sum = Some(index), - HelperKind::SumSquares => entry.sumsq = Some(index), + let entry = helper_columns.entry(helper.target_column).or_default(); + match helper.kind { + HelperKind::Count => entry.count = Some(index), + HelperKind::Sum => entry.sum = Some(index), + HelperKind::SumSquares => entry.sumsq = Some(index), + } } } @@ -524,8 +526,7 @@ mod test { shard1.add("3"); rows.push_back(shard1); - let plan = ProjectionRewritePlan::default(); - let mut result = Aggregates::new(&rows, &decoder, &aggregate, &plan) + let mut result = Aggregates::new(&rows, &decoder, &aggregate, None) .unwrap() .aggregate() .unwrap(); @@ -572,8 +573,7 @@ mod test { ); rows.push_back(row); } - let plan = ProjectionRewritePlan::default(); - let mut result = Aggregates::new(&rows, &decoder, &aggregate, &plan) + let mut result = Aggregates::new(&rows, &decoder, &aggregate, None) .expect("count aggregate") .aggregate() .expect("merge partial counts"); @@ -606,15 +606,10 @@ mod test { shard2.add(20.0_f64).add(4_i64); rows.push_back(shard2); - let mut result = Aggregates::new( - &rows, - &decoder, - &aggregate, - &ProjectionRewritePlan::default(), - ) - .unwrap() - .aggregate() - .unwrap(); + let mut result = Aggregates::new(&rows, &decoder, &aggregate, None) + .unwrap() + .aggregate() + .unwrap(); assert_eq!(result.len(), 2); let mut groups: Vec<(f64, i64)> = result @@ -651,15 +646,10 @@ mod test { shard2.add(Bytes::from_static(b"{{5,6},{7,8}}")).add(1_i64); rows.push_back(shard2); - let mut result = Aggregates::new( - &rows, - &decoder, - &aggregate, - &ProjectionRewritePlan::default(), - ) - .unwrap() - .aggregate() - .unwrap(); + let mut result = Aggregates::new(&rows, &decoder, &aggregate, None) + .unwrap() + .aggregate() + .unwrap(); let mut groups: Vec<(String, i64)> = result .drain(..) @@ -698,15 +688,10 @@ mod test { shard1.add(input.clone()).add(1_i64); rows.push_back(shard1); - let mut result = Aggregates::new( - &rows, - &decoder, - &aggregate, - &ProjectionRewritePlan::default(), - ) - .unwrap() - .aggregate() - .unwrap(); + let mut result = Aggregates::new(&rows, &decoder, &aggregate, None) + .unwrap() + .aggregate() + .unwrap(); assert_eq!(result.len(), 1); let row = result.pop_front().unwrap(); diff --git a/pgdog/src/backend/pool/connection/buffer.rs b/pgdog/src/backend/pool/connection/buffer.rs index 786063fbf..d0ba21a3e 100644 --- a/pgdog/src/backend/pool/connection/buffer.rs +++ b/pgdog/src/backend/pool/connection/buffer.rs @@ -140,7 +140,7 @@ impl Buffer { &mut self, aggregate: &Aggregate, decoder: &Decoder, - plan: &ProjectionRewritePlan, + plan: Option<&ProjectionRewritePlan>, ) -> Result<(), super::Error> { let buffer: VecDeque = std::mem::take(&mut self.buffer); let rows = if aggregate.is_empty() { @@ -156,10 +156,14 @@ impl Buffer { Ok(()) } - pub(super) fn drop_helper_columns(&mut self, plan: &ProjectionRewritePlan, decoder: &Decoder) { - if plan.is_noop() { + pub(super) fn drop_helper_columns( + &mut self, + plan: Option<&ProjectionRewritePlan>, + decoder: &Decoder, + ) { + let Some(plan) = plan else { return; - } + }; let drop = plan.drop_columns(decoder.row_description()); @@ -292,8 +296,7 @@ mod test { buf.add(dr.message()).unwrap(); } - buf.aggregate(&agg, &Decoder::from(rd), &ProjectionRewritePlan::default()) - .unwrap(); + buf.aggregate(&agg, &Decoder::from(rd), None).unwrap(); buf.mark_full(); assert_eq!(buf.len(), 1); @@ -319,8 +322,7 @@ mod test { } } - buf.aggregate(&agg, &Decoder::from(rd), &ProjectionRewritePlan::default()) - .unwrap(); + buf.aggregate(&agg, &Decoder::from(rd), None).unwrap(); buf.mark_full(); assert_eq!(buf.len(), 2); diff --git a/pgdog/src/backend/pool/connection/multi_shard/mod.rs b/pgdog/src/backend/pool/connection/multi_shard/mod.rs index 81de53db3..5eafe91fd 100644 --- a/pgdog/src/backend/pool/connection/multi_shard/mod.rs +++ b/pgdog/src/backend/pool/connection/multi_shard/mod.rs @@ -246,19 +246,18 @@ impl MultiShard { if !self.buffer.is_empty() { // Helpers remain in the internal row through aggregation and // sorting, then are removed before client-visible operations. - if let Err(err) = self.buffer.aggregate( - self.route.aggregate(), - &self.decoder, - &self.route.projection_rewrite_plan, - ) { + let plan = self.route.projection_rewrite_plan.as_deref(); + if let Err(err) = self + .buffer + .aggregate(self.route.aggregate(), &self.decoder, plan) + { // Probably an unsupported aggregation feature was requested. self.buffer.clear(); return Err(err.into()); } self.buffer.sort(self.route.order_by(), &self.decoder); - self.buffer - .drop_helper_columns(&self.route.projection_rewrite_plan, &self.decoder); + self.buffer.drop_helper_columns(plan, &self.decoder); self.buffer.distinct(self.route.distinct(), &self.decoder); self.buffer.limit(self.route.limit()); } @@ -301,12 +300,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.projection_rewrite_plan; - if plan.is_noop() { - forward = Some(message); - } else { + if let Some(plan) = self.route.projection_rewrite_plan.as_deref() { let client_rd = rd.drop_columns(plan.drop_columns(&rd)); forward = Some(client_rd.message()); + } else { + forward = Some(message); } // The next statement describes a different result set. diff --git a/pgdog/src/backend/pool/connection/multi_shard/test.rs b/pgdog/src/backend/pool/connection/multi_shard/test.rs index 1aaf6805d..f6cc5ccfa 100644 --- a/pgdog/src/backend/pool/connection/multi_shard/test.rs +++ b/pgdog/src/backend/pool/connection/multi_shard/test.rs @@ -79,7 +79,7 @@ fn test_order_by_helper_after_star_expansion_is_dropped_after_sorting() { Default::default(), None, ); - route.projection_rewrite_plan = plan; + route.projection_rewrite_plan = Some(std::sync::Arc::new(plan)); let mut multi_shard = MultiShard::new(2, &route); let row_description = RowDescription::new(&[ diff --git a/pgdog/src/frontend/client/query_engine/test/rewrite_projection.rs b/pgdog/src/frontend/client/query_engine/test/rewrite_projection.rs index d08c39dd0..6732ddec4 100644 --- a/pgdog/src/frontend/client/query_engine/test/rewrite_projection.rs +++ b/pgdog/src/frontend/client/query_engine/test/rewrite_projection.rs @@ -55,7 +55,7 @@ async fn direct_aggregate_keeps_base_sql() { _ => panic!("expected Query"), }; assert_eq!(query.query(), sql); - assert!(client_request.route().projection_rewrite_plan.is_noop()); + assert!(client_request.route().projection_rewrite_plan.is_none()); } #[tokio::test] @@ -89,6 +89,8 @@ async fn cross_shard_aggregate_adds_and_tracks_helpers() { client_request .route() .projection_rewrite_plan + .as_deref() + .expect("projection rewrite plan") .aggregate_helpers .len(), 1 @@ -255,6 +257,8 @@ async fn cross_shard_order_by_projects_missing_sort_column() { client_request .route() .projection_rewrite_plan + .as_deref() + .expect("projection rewrite plan") .order_by_helpers .len(), 1 @@ -334,7 +338,15 @@ fn aliased_projected_sort_column_remaps_route() { request.route().order_by(), &[OrderBy::AscColumn("item_price".into())] ); - assert!(!request.route().projection_rewrite_plan.order_by_helpers[0].injected); + assert!( + !request + .route() + .projection_rewrite_plan + .as_deref() + .expect("projection rewrite plan") + .order_by_helpers[0] + .injected + ); } #[test] @@ -437,14 +449,12 @@ async fn aggregate_order_by_and_offset_compose_after_route() { route.order_by(), &[OrderBy::AscColumn("__pgdog_order_col0".into())] ); - assert_eq!( - route.projection_rewrite_plan.aggregate_helpers[0].alias, - "__pgdog_count_col0" - ); - assert_eq!( - route.projection_rewrite_plan.order_by_helpers[0].alias, - "__pgdog_order_col0" - ); + let plan = route + .projection_rewrite_plan + .as_deref() + .expect("projection rewrite plan"); + assert_eq!(plan.aggregate_helpers[0].alias, "__pgdog_count_col0"); + assert_eq!(plan.order_by_helpers[0].alias, "__pgdog_order_col0"); assert_eq!( route.limit(), &Limit { diff --git a/pgdog/src/frontend/router/parser/aggregate.rs b/pgdog/src/frontend/router/parser/aggregate.rs index 27bec68ee..d26dbda38 100644 --- a/pgdog/src/frontend/router/parser/aggregate.rs +++ b/pgdog/src/frontend/router/parser/aggregate.rs @@ -39,6 +39,23 @@ pub(crate) enum AggregateFunction { Unrecognized(String), } +impl AggregateFunction { + pub(crate) fn from_builtin_name(name: &str) -> Option { + Some(match name { + "count" => Self::Count, + "max" => Self::Max, + "min" => Self::Min, + "sum" => Self::Sum, + "avg" => Self::Avg, + "stddev" | "stddev_samp" => Self::StddevSamp, + "stddev_pop" => Self::StddevPop, + "variance" | "var_samp" => Self::VarSamp, + "var_pop" => Self::VarPop, + _ => return None, + }) + } +} + impl fmt::Display for AggregateFunction { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> { match self { @@ -120,24 +137,12 @@ impl Aggregate { .enumerate() .filter_map(|(idx, node)| { let func = Function::try_from(node.val()).ok()?; - let function = match func.name { - "count" => AggregateFunction::Count, - "max" => AggregateFunction::Max, - "min" => AggregateFunction::Min, - "sum" => AggregateFunction::Sum, - "avg" => AggregateFunction::Avg, - "stddev" | "stddev_samp" => AggregateFunction::StddevSamp, - "stddev_pop" => AggregateFunction::StddevPop, - "variance" | "var_samp" => AggregateFunction::VarSamp, - "var_pop" => AggregateFunction::VarPop, - fname => { - if schema.aggregate_functions.contains(fname) { - AggregateFunction::Unrecognized(fname.to_owned()) - } else { - return None; - } - } - }; + let function = AggregateFunction::from_builtin_name(func.name).or_else(|| { + schema + .aggregate_functions + .contains(func.name) + .then(|| AggregateFunction::Unrecognized(func.name.to_owned())) + })?; // FIXME: This doesn't correctly handle distinct behind type // casts any other nodes that Function::try_from handles diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs b/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs index 772246586..9b7f55565 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs @@ -53,7 +53,7 @@ impl OrderByHelper { } } -#[derive(Debug, Clone, Default, PartialEq)] +#[derive(Debug, Default, PartialEq)] pub(crate) struct ProjectionRewritePlan { pub(crate) aggregate_helpers: Vec, pub(crate) order_by_helpers: Vec, @@ -82,7 +82,7 @@ impl ProjectionRewritePlan { #[derive(Debug, Clone)] pub(crate) struct PostRouteRewrite { sql: Arc, - plan: ProjectionRewritePlan, + plan: Option>, } pub(crate) fn finalize_after_route( @@ -144,12 +144,12 @@ pub(crate) fn finalize_after_route( { parse.set_query(&rewrite.sql); } - if !rewrite.plan.is_noop() + if let Some(plan) = &rewrite.plan && let Some(route) = request.route.as_mut() { - route.projection_rewrite_plan = rewrite.plan.clone(); + route.projection_rewrite_plan = Some(plan.clone()); let mut order_by = route.order_by().to_vec(); - for helper in &rewrite.plan.order_by_helpers { + for helper in &plan.order_by_helpers { // Prefer the structural position. The source fallback handles a // bind-dependent vector sort omitted from this execution's route. let position = order_by @@ -202,6 +202,7 @@ fn build( return Ok(None); } let sql: Arc = pg_raw_parse::deparse(&*rewritten)?.as_str().into(); + let plan = (!plan.is_noop()).then(|| Arc::new(plan)); Ok(Some(PostRouteRewrite { sql, plan })) } diff --git a/pgdog/src/frontend/router/parser/route.rs b/pgdog/src/frontend/router/parser/route.rs index afc53609a..5e19e0351 100644 --- a/pgdog/src/frontend/router/parser/route.rs +++ b/pgdog/src/frontend/router/parser/route.rs @@ -1,4 +1,4 @@ -use std::{fmt::Display, ops::Deref}; +use std::{fmt::Display, ops::Deref, sync::Arc}; use super::{ Aggregate, DistinctBy, Limit, OrderBy, StatementType, advisory_lock::AdvisoryLocks, @@ -119,8 +119,8 @@ pub(crate) struct Route { advisory_locks: AdvisoryLocks, /// `DISTINCT` clause, if set. distinct: Option, - /// Temporary columns projected for cross-shard result processing. - pub(crate) projection_rewrite_plan: ProjectionRewritePlan, + /// Shared metadata for temporary cross-shard projections, absent on the common path. + pub(crate) projection_rewrite_plan: Option>, /// Our query explain plan. We attach /// this to the `EXPLAIN` output. explain: Option, From 9ac875e08e30732ada881d6360b91a25d68af92b Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Tue, 27 Oct 2026 20:45:34 +0530 Subject: [PATCH 10/24] arc clone --- .../src/frontend/router/parser/rewrite/statement/projection.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs b/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs index 9b7f55565..00b094368 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs @@ -147,7 +147,7 @@ pub(crate) fn finalize_after_route( if let Some(plan) = &rewrite.plan && let Some(route) = request.route.as_mut() { - route.projection_rewrite_plan = Some(plan.clone()); + route.projection_rewrite_plan = Some(Arc::clone(plan)); let mut order_by = route.order_by().to_vec(); for helper in &plan.order_by_helpers { // Prefer the structural position. The source fallback handles a From 773bb04523dbe3477a868ee7c5cb6797d5b358a6 Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Tue, 27 Oct 2026 20:49:47 +0530 Subject: [PATCH 11/24] remove agg --- pgdog/src/frontend/router/parser/aggregate.rs | 41 ++++++++----------- 1 file changed, 18 insertions(+), 23 deletions(-) diff --git a/pgdog/src/frontend/router/parser/aggregate.rs b/pgdog/src/frontend/router/parser/aggregate.rs index d26dbda38..27bec68ee 100644 --- a/pgdog/src/frontend/router/parser/aggregate.rs +++ b/pgdog/src/frontend/router/parser/aggregate.rs @@ -39,23 +39,6 @@ pub(crate) enum AggregateFunction { Unrecognized(String), } -impl AggregateFunction { - pub(crate) fn from_builtin_name(name: &str) -> Option { - Some(match name { - "count" => Self::Count, - "max" => Self::Max, - "min" => Self::Min, - "sum" => Self::Sum, - "avg" => Self::Avg, - "stddev" | "stddev_samp" => Self::StddevSamp, - "stddev_pop" => Self::StddevPop, - "variance" | "var_samp" => Self::VarSamp, - "var_pop" => Self::VarPop, - _ => return None, - }) - } -} - impl fmt::Display for AggregateFunction { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> { match self { @@ -137,12 +120,24 @@ impl Aggregate { .enumerate() .filter_map(|(idx, node)| { let func = Function::try_from(node.val()).ok()?; - let function = AggregateFunction::from_builtin_name(func.name).or_else(|| { - schema - .aggregate_functions - .contains(func.name) - .then(|| AggregateFunction::Unrecognized(func.name.to_owned())) - })?; + let function = match func.name { + "count" => AggregateFunction::Count, + "max" => AggregateFunction::Max, + "min" => AggregateFunction::Min, + "sum" => AggregateFunction::Sum, + "avg" => AggregateFunction::Avg, + "stddev" | "stddev_samp" => AggregateFunction::StddevSamp, + "stddev_pop" => AggregateFunction::StddevPop, + "variance" | "var_samp" => AggregateFunction::VarSamp, + "var_pop" => AggregateFunction::VarPop, + fname => { + if schema.aggregate_functions.contains(fname) { + AggregateFunction::Unrecognized(fname.to_owned()) + } else { + return None; + } + } + }; // FIXME: This doesn't correctly handle distinct behind type // casts any other nodes that Function::try_from handles From 516b41c6dc4c407e4b00e7175bb0ce51f582fbad Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Tue, 27 Oct 2026 21:02:36 +0530 Subject: [PATCH 12/24] fix: reject cross-shard HAVING with an explicit error --- .../client/query_engine/route_query.rs | 21 ++- .../test/unsupported_aggregation.rs | 148 ++++++++++-------- .../router/parser/rewrite/statement/error.rs | 3 + .../parser/rewrite/statement/projection.rs | 5 + 4 files changed, 106 insertions(+), 71 deletions(-) diff --git a/pgdog/src/frontend/client/query_engine/route_query.rs b/pgdog/src/frontend/client/query_engine/route_query.rs index 133035d50..49b05ced9 100644 --- a/pgdog/src/frontend/client/query_engine/route_query.rs +++ b/pgdog/src/frontend/client/query_engine/route_query.rs @@ -4,6 +4,7 @@ use tracing::trace; use crate::backend::Error as BackendError; use crate::frontend::router::Error as RouterError; use crate::frontend::router::parser::Error as ParserError; +use crate::frontend::router::parser::rewrite::statement::Error as StatementRewriteError; use crate::frontend::router::parser::rewrite::statement::plan::RewriteResult; use crate::frontend::router::parser::rewrite::statement::projection; use crate::frontend::router::sharding::lookup; @@ -159,11 +160,27 @@ impl QueryEngine { client_request.route = Some(command.route().clone()); trace!("routing {:#?} to {:#?}", client_request.messages, command,); - projection::finalize_after_route( + let finalize = projection::finalize_after_route( client_request, &cluster.schema(), rewrite_result.and_then(RewriteResult::offset_plan), - )?; + ); + match finalize { + Err(err @ StatementRewriteError::UnsupportedHaving(_)) => { + if let Some(transaction) = context.transaction.as_mut() { + transaction.set_error(); + } + self.error_response( + context, + client_request, + ErrorResponse::unsupported(&err), + ) + .await?; + return Ok(false); + } + Err(err) => return Err(err.into()), + Ok(()) => {} + } if let Some(rewrite_result) = rewrite_result { rewrite_result.apply_after_route(client_request)?; diff --git a/pgdog/src/frontend/client/query_engine/test/unsupported_aggregation.rs b/pgdog/src/frontend/client/query_engine/test/unsupported_aggregation.rs index fa7af43d2..7633ac241 100644 --- a/pgdog/src/frontend/client/query_engine/test/unsupported_aggregation.rs +++ b/pgdog/src/frontend/client/query_engine/test/unsupported_aggregation.rs @@ -13,91 +13,101 @@ use crate::{ async fn test_unsupported_aggregation_keeps_connection_usable() { let mut client = SpawnedClient::new_sharded(Parameters::default()).await; - for begin in [None, Some("BEGIN"), Some("BEGIN READ ONLY")] { - timeout(Duration::from_secs(5), async { - if let Some(begin) = begin { - client.send(Query::new(begin)).await; - assert_eq!( - expect_message!(client.read().await, CommandComplete).command(), - "BEGIN" + for (query, expected_error) in [ + ( + "SELECT count(DISTINCT value) FROM sharded", + "unsupported aggregation count: count(DISTINCT ...) is not yet supported", + ), + ( + "SELECT count(*) FROM sharded HAVING count(*) + 1 > 0", + "cross-shard HAVING is unsupported: aggregates would be filtered on each shard before merging", + ), + ] { + for begin in [None, Some("BEGIN"), Some("BEGIN READ ONLY")] { + timeout(Duration::from_secs(5), async { + if let Some(begin) = begin { + client.send(Query::new(begin)).await; + assert_eq!( + expect_message!(client.read().await, CommandComplete).command(), + "BEGIN" + ); + assert_eq!( + expect_message!(client.read().await, ReadyForQuery).status, + 'T' + ); + } + + client.send(Query::new(query)).await; + + let messages = timeout(Duration::from_secs(2), client.read_until('E')) + .await + .unwrap_or_else(|_| { + panic!("unsupported aggregation did not return an error: {query}") + }); + let error = expect_message!( + messages.last().expect("error response").clone(), + ErrorResponse ); + assert_eq!(error.severity, "ERROR"); + assert_eq!(error.code, "0A000"); + assert_eq!(error.message, expected_error); assert_eq!( - expect_message!(client.read().await, ReadyForQuery).status, - 'T' + expect_message!( + timeout(Duration::from_secs(2), client.read()) + .await + .expect("ReadyForQuery missing after unsupported aggregation"), + ReadyForQuery + ) + .status, + if begin.is_some() { 'E' } else { 'I' } ); - } - - client - .send(Query::new("SELECT count(DISTINCT value) FROM sharded")) - .await; - let messages = timeout(Duration::from_secs(2), client.read_until('E')) - .await - .expect("unsupported aggregation did not return an error"); - let error = expect_message!( - messages.last().expect("error response").clone(), - ErrorResponse - ); - assert_eq!(error.severity, "ERROR"); - assert_eq!(error.code, "0A000"); - assert_eq!( - error.message, - "unsupported aggregation count: count(DISTINCT ...) is not yet supported" - ); - assert_eq!( - expect_message!( - timeout(Duration::from_secs(2), client.read()) - .await - .expect("ReadyForQuery missing after unsupported aggregation"), - ReadyForQuery - ) - .status, - if begin.is_some() { 'E' } else { 'I' } - ); + if begin.is_some() { + // The failed transaction must reject further queries until rollback. + for _ in 0..2 { + client + .send(Query::new("SELECT count(*) FROM sharded WHERE false")) + .await; + let error = expect_message!(client.read().await, ErrorResponse); + assert_eq!(error.severity, "ERROR"); + assert_eq!(error.code, "25P02"); + assert_eq!( + expect_message!(client.read().await, ReadyForQuery).status, + 'E' + ); + } - if begin.is_some() { - // The failed transaction must reject further queries until rollback. - for _ in 0..2 { - client - .send(Query::new("SELECT count(*) FROM sharded WHERE false")) - .await; - let error = expect_message!(client.read().await, ErrorResponse); - assert_eq!(error.severity, "ERROR"); - assert_eq!(error.code, "25P02"); + client.send(Query::new("ROLLBACK")).await; + assert_eq!( + expect_message!(client.read().await, CommandComplete).command(), + "ROLLBACK" + ); assert_eq!( expect_message!(client.read().await, ReadyForQuery).status, - 'E' + 'I' ); } - client.send(Query::new("ROLLBACK")).await; + // Each shard returns a count row, exercising aggregation again without fixtures. + client + .send(Query::new("SELECT count(*) FROM sharded WHERE false")) + .await; + expect_message!(client.read().await, RowDescription); + let row = expect_message!(client.read().await, DataRow); + assert_eq!(row.get_int(0, true), Some(0)); assert_eq!( expect_message!(client.read().await, CommandComplete).command(), - "ROLLBACK" + "SELECT 1" ); assert_eq!( expect_message!(client.read().await, ReadyForQuery).status, 'I' ); - } - - // Each shard returns a count row, exercising aggregation again without fixtures. - client - .send(Query::new("SELECT count(*) FROM sharded WHERE false")) - .await; - expect_message!(client.read().await, RowDescription); - let row = expect_message!(client.read().await, DataRow); - assert_eq!(row.get_int(0, true), Some(0)); - assert_eq!( - expect_message!(client.read().await, CommandComplete).command(), - "SELECT 1" - ); - assert_eq!( - expect_message!(client.read().await, ReadyForQuery).status, - 'I' - ); - }) - .await - .unwrap_or_else(|_| panic!("unsupported aggregation or subsequent query hung: {begin:?}")); + }) + .await + .unwrap_or_else(|_| { + panic!("unsupported aggregation or subsequent query hung: {query}, {begin:?}") + }); + } } } diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/error.rs b/pgdog/src/frontend/router/parser/rewrite/statement/error.rs index 844be94f3..e928b3f3d 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/error.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/error.rs @@ -14,6 +14,9 @@ pub(crate) enum Error { #[error("sharding key assignment unsupported: {0}")] UnsupportedShardingKeyUpdate(String), + #[error("cross-shard HAVING is unsupported: {0}")] + UnsupportedHaving(&'static str), + #[error("net: {0}")] Net(#[from] crate::net::Error), diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs b/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs index 00b094368..ce151ce26 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs @@ -180,6 +180,11 @@ fn build( let Some(Node::SelectStmt(select)) = ast.stmts().next() else { return Ok(None); }; + if !matches!(select.having_clause(), Node::None) { + return Err(Error::UnsupportedHaving( + "aggregates would be filtered on each shard before merging", + )); + } let aggregate = Aggregate::parse(select, schema); if aggregate.is_empty() && select.sort_clause().is_empty() && !rewrite_offset { From f442f6e12a82a8a9ccd1b0eff84a6d7c4283f7a7 Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Tue, 27 Oct 2026 21:32:41 +0530 Subject: [PATCH 13/24] cleanup --- .../backend/pool/connection/buffer/having.rs | 182 +----------------- pgdog/src/frontend/router/parser/aggregate.rs | 41 ++-- pgdog/src/frontend/router/parser/having.rs | 7 - .../router/parser/rewrite/statement/having.rs | 124 ++++-------- 4 files changed, 67 insertions(+), 287 deletions(-) diff --git a/pgdog/src/backend/pool/connection/buffer/having.rs b/pgdog/src/backend/pool/connection/buffer/having.rs index 77f186768..fa3450cd0 100644 --- a/pgdog/src/backend/pool/connection/buffer/having.rs +++ b/pgdog/src/backend/pool/connection/buffer/having.rs @@ -1,6 +1,6 @@ use std::{cmp::Ordering, collections::HashMap}; -use pgdog_postgres_types::{Datum, Double, Numeric}; +use pgdog_postgres_types::Datum; use crate::{ frontend::router::parser::{CompareOperator, HavingExpr, HavingValue}, @@ -21,16 +21,6 @@ pub(super) struct Evaluator<'a> { columns: HashMap<&'a str, usize>, } -impl TruthValue { - fn not(self) -> Self { - match self { - Self::True => Self::False, - Self::False => Self::True, - Self::Unknown => Self::Unknown, - } - } -} - impl<'a> Evaluator<'a> { pub(super) fn new(expression: &'a HavingExpr, decoder: &'a Decoder) -> Self { let columns = decoder @@ -69,39 +59,6 @@ impl<'a> Evaluator<'a> { let right = self.value(right, row, decoder)?; compare(&left, *operator, &right) } - HavingExpr::And(expressions) => { - let mut result = TruthValue::True; - for expression in expressions { - match self.evaluate_expression(expression, row, decoder)? { - TruthValue::False => return Ok(TruthValue::False), - TruthValue::Unknown => result = TruthValue::Unknown, - TruthValue::True => {} - } - } - Ok(result) - } - HavingExpr::Or(expressions) => { - let mut result = TruthValue::False; - for expression in expressions { - match self.evaluate_expression(expression, row, decoder)? { - TruthValue::True => return Ok(TruthValue::True), - TruthValue::Unknown => result = TruthValue::Unknown, - TruthValue::False => {} - } - } - Ok(result) - } - HavingExpr::Not(expression) => { - Ok(self.evaluate_expression(expression, row, decoder)?.not()) - } - HavingExpr::IsNull { value, negated } => { - let is_null = self.value(value, row, decoder)?.is_null(); - Ok(if is_null != *negated { - TruthValue::True - } else { - TruthValue::False - }) - } } } @@ -158,19 +115,8 @@ fn compare(left: &Datum, operator: CompareOperator, right: &Datum) -> Result Option { - if let (Some(left), Some(right)) = (integer(left), integer(right)) { - return Some(left.cmp(&right)); - } - - if let (Some(left), Some(right)) = (exact_numeric(left), exact_numeric(right)) { - return left.partial_cmp(&right); - } - - if let (Some(left), Some(right)) = (float(left), float(right)) { - return Some(Double(left).cmp(&Double(right))); - } - - None + let (left, right) = (integer(left)?, integer(right)?); + Some(left.cmp(&right)) } fn integer(value: &Datum) -> Option { @@ -182,27 +128,6 @@ fn integer(value: &Datum) -> Option { } } -fn exact_numeric(value: &Datum) -> Option { - match value { - Datum::Numeric(value) => Some(*value), - Datum::Bigint(value) => Some(Numeric::from(*value)), - Datum::Integer(value) => Some(Numeric::from(*value)), - Datum::SmallInt(value) => Some(Numeric::from(i32::from(*value))), - _ => None, - } -} - -fn float(value: &Datum) -> Option { - match value { - Datum::Double(value) => Some(value.0), - Datum::Float(value) => Some(value.0.into()), - Datum::Bigint(value) => Some(*value as f64), - Datum::Integer(value) => Some((*value).into()), - Datum::SmallInt(value) => Some((*value).into()), - _ => None, - } -} - #[cfg(test)] mod tests { use super::*; @@ -213,7 +138,6 @@ mod tests { messages::{DataRow, Datum}, }, }; - use pgdog_postgres_types::Float; fn compare_expression(operator: CompareOperator) -> HavingExpr { HavingExpr::Compare { @@ -236,27 +160,7 @@ mod tests { } #[test] - fn compares_merged_bigint_with_numeric_without_losing_precision() { - let decoder = Decoder::from(RowDescription::new(&[ - Field::bigint("__pgdog_having_0"), - Field::numeric("__pgdog_having_1"), - ])); - let mut row = DataRow::new(); - row.add(9_007_199_254_740_993_i64) - .add(Numeric::from(9_007_199_254_740_992_i64)); - - assert!( - evaluate( - &compare_expression(CompareOperator::Greater), - &row, - &decoder - ) - .unwrap() - ); - } - - #[test] - fn implements_sql_three_valued_boolean_logic() { + fn null_comparison_is_unknown_and_filters_the_group() { let decoder = Decoder::from(RowDescription::new(&[ Field::bigint("__pgdog_having_0"), integer_field("__pgdog_having_1"), @@ -265,64 +169,11 @@ mod tests { row.add(Datum::Null).add(2_i32); let unknown = compare_expression(CompareOperator::Greater); - let is_null = HavingExpr::IsNull { - value: HavingValue::new("__pgdog_having_0".into()), - negated: false, - }; - let is_not_null = HavingExpr::IsNull { - value: HavingValue::new("__pgdog_having_1".into()), - negated: true, - }; - let known_false = HavingExpr::Compare { - left: HavingValue::new("__pgdog_having_1".into()), - operator: CompareOperator::Greater, - right: HavingValue::new("__pgdog_having_1".into()), - }; - assert_eq!( evaluate_expression(&unknown, &row, &decoder).unwrap(), TruthValue::Unknown ); - assert_eq!( - evaluate_expression( - &HavingExpr::And(vec![is_null.clone(), unknown.clone()]), - &row, - &decoder - ) - .unwrap(), - TruthValue::Unknown - ); - assert_eq!( - evaluate_expression( - &HavingExpr::And(vec![unknown.clone(), known_false.clone()]), - &row, - &decoder - ) - .unwrap(), - TruthValue::False - ); - assert_eq!( - evaluate_expression( - &HavingExpr::Or(vec![known_false.clone(), unknown.clone()]), - &row, - &decoder - ) - .unwrap(), - TruthValue::Unknown - ); - assert!( - evaluate( - &HavingExpr::Or(vec![unknown.clone(), is_null]), - &row, - &decoder - ) - .unwrap() - ); - assert!(evaluate(&is_not_null, &row, &decoder).unwrap()); - assert_eq!( - evaluate_expression(&HavingExpr::Not(Box::new(unknown)), &row, &decoder).unwrap(), - TruthValue::Unknown - ); + assert!(!evaluate(&unknown, &row, &decoder).unwrap()); let missing = HavingExpr::Compare { left: HavingValue::new("__pgdog_missing".into()), @@ -336,7 +187,7 @@ mod tests { } #[test] - fn compares_supported_scalar_types_and_all_operators() { + fn compares_integers_across_widths_with_all_operators() { for (operator, left, right, expected) in [ ( CompareOperator::Equal, @@ -364,14 +215,14 @@ mod tests { ), ( CompareOperator::Greater, - Datum::Numeric(Numeric::from(3_i32)), + Datum::Bigint(3), Datum::SmallInt(2), true, ), ( CompareOperator::GreaterOrEqual, - Datum::Float(Float(2.0)), - Datum::Double(Double(2.0)), + Datum::Integer(2), + Datum::Bigint(2), true, ), ( @@ -391,21 +242,6 @@ mod tests { ); } - assert_eq!( - numeric_ordering(&Datum::Double(Double(f64::NAN)), &Datum::Float(Float(1.0))), - Some(Ordering::Greater) - ); - assert_eq!( - numeric_ordering( - &Datum::Double(Double(f64::NAN)), - &Datum::Double(Double(f64::NAN)) - ), - Some(Ordering::Equal) - ); - assert_eq!( - numeric_ordering(&Datum::Float(Float(1.0)), &Datum::Double(Double(f64::NAN))), - Some(Ordering::Less) - ); assert!(matches!( compare( &Datum::Text("2".into()), diff --git a/pgdog/src/frontend/router/parser/aggregate.rs b/pgdog/src/frontend/router/parser/aggregate.rs index 27bec68ee..d26dbda38 100644 --- a/pgdog/src/frontend/router/parser/aggregate.rs +++ b/pgdog/src/frontend/router/parser/aggregate.rs @@ -39,6 +39,23 @@ pub(crate) enum AggregateFunction { Unrecognized(String), } +impl AggregateFunction { + pub(crate) fn from_builtin_name(name: &str) -> Option { + Some(match name { + "count" => Self::Count, + "max" => Self::Max, + "min" => Self::Min, + "sum" => Self::Sum, + "avg" => Self::Avg, + "stddev" | "stddev_samp" => Self::StddevSamp, + "stddev_pop" => Self::StddevPop, + "variance" | "var_samp" => Self::VarSamp, + "var_pop" => Self::VarPop, + _ => return None, + }) + } +} + impl fmt::Display for AggregateFunction { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> { match self { @@ -120,24 +137,12 @@ impl Aggregate { .enumerate() .filter_map(|(idx, node)| { let func = Function::try_from(node.val()).ok()?; - let function = match func.name { - "count" => AggregateFunction::Count, - "max" => AggregateFunction::Max, - "min" => AggregateFunction::Min, - "sum" => AggregateFunction::Sum, - "avg" => AggregateFunction::Avg, - "stddev" | "stddev_samp" => AggregateFunction::StddevSamp, - "stddev_pop" => AggregateFunction::StddevPop, - "variance" | "var_samp" => AggregateFunction::VarSamp, - "var_pop" => AggregateFunction::VarPop, - fname => { - if schema.aggregate_functions.contains(fname) { - AggregateFunction::Unrecognized(fname.to_owned()) - } else { - return None; - } - } - }; + let function = AggregateFunction::from_builtin_name(func.name).or_else(|| { + schema + .aggregate_functions + .contains(func.name) + .then(|| AggregateFunction::Unrecognized(func.name.to_owned())) + })?; // FIXME: This doesn't correctly handle distinct behind type // casts any other nodes that Function::try_from handles diff --git a/pgdog/src/frontend/router/parser/having.rs b/pgdog/src/frontend/router/parser/having.rs index 5ad23a267..57ba43a62 100644 --- a/pgdog/src/frontend/router/parser/having.rs +++ b/pgdog/src/frontend/router/parser/having.rs @@ -35,13 +35,6 @@ pub(crate) enum HavingExpr { operator: CompareOperator, right: HavingValue, }, - And(Vec), - Or(Vec), - Not(Box), - IsNull { - value: HavingValue, - negated: bool, - }, } /// Metadata produced while rewriting a cross-shard `HAVING` clause. diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/having.rs b/pgdog/src/frontend/router/parser/rewrite/statement/having.rs index d9a015fcf..33811e0be 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/having.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/having.rs @@ -24,13 +24,6 @@ enum HavingTemplate { operator: CompareOperator, right: ValueTemplate, }, - And(Vec), - Or(Vec), - Not(Box), - IsNull { - value: ValueTemplate, - negated: bool, - }, } struct ValueTemplate { @@ -146,17 +139,6 @@ impl HavingTemplate { operator, right: right.materialize(), }, - Self::And(expressions) => { - HavingExpr::And(expressions.into_iter().map(Self::materialize).collect()) - } - Self::Or(expressions) => { - HavingExpr::Or(expressions.into_iter().map(Self::materialize).collect()) - } - Self::Not(expression) => HavingExpr::Not(Box::new(expression.materialize())), - Self::IsNull { value, negated } => HavingExpr::IsNull { - value: value.materialize(), - negated, - }, } } } @@ -195,42 +177,8 @@ impl<'a> Parser<'a> { right: self.parse_value(expression.rexpr())?, }) } - Node::BoolExpr(expression) => { - let expressions = expression - .args() - .iter() - .map(|node| self.parse_expression(node)) - .collect::, _>>()?; - match expression.boolop { - nodes::BoolExprType::AND_EXPR => Ok(HavingTemplate::And(expressions)), - nodes::BoolExprType::OR_EXPR => Ok(HavingTemplate::Or(expressions)), - nodes::BoolExprType::NOT_EXPR => { - let expression = expressions - .into_iter() - .exactly_one() - .map_err(|_| Error::UnsupportedHaving("NOT requires one expression"))?; - Ok(HavingTemplate::Not(Box::new(expression))) - } - _ => Err(Error::UnsupportedHaving("unsupported boolean expression")), - } - } - Node::NullTest(test) => { - let negated = match test.nulltesttype { - nodes::NullTestType::IS_NULL => false, - nodes::NullTestType::IS_NOT_NULL => true, - _ => { - return Err(Error::UnsupportedHaving( - "only IS NULL and IS NOT NULL are supported", - )); - } - }; - Ok(HavingTemplate::IsNull { - value: self.parse_value(test.arg())?, - negated, - }) - } _ => Err(Error::UnsupportedHaving( - "supported predicates are comparisons, AND, OR, NOT and NULL tests", + "only a single comparison is supported", )), } } @@ -250,7 +198,6 @@ impl<'a> Parser<'a> { fn supported_value(node: Node<'_>) -> bool { match node { Node::A_Const(_) | Node::ParamRef(_) | Node::ColumnRef(_) => true, - Node::TypeCast(cast) => Self::supported_value(cast.arg()), Node::FuncCall(_) => Self::supported_aggregate(node), _ => false, } @@ -310,49 +257,48 @@ mod tests { } #[test] - fn rewrites_nested_boolean_null_cast_and_parameter_expression() { + fn rewrites_single_comparison_with_parameter() { let (sql, plan) = rewrite( - "SELECT customer_id, SUM(price)::numeric FROM purchases \ - GROUP BY customer_id \ - HAVING NOT (COUNT(*) <= $1 OR SUM(price) IS NOT NULL) \ - AND MAX(price)::numeric >= 10", + "SELECT customer_id, COUNT(*) FROM purchases \ + GROUP BY customer_id HAVING COUNT(*) > $1", ); - let HavingExpr::And(expressions) = plan.expression() else { - panic!("expected top-level AND") - }; assert!(matches!( - expressions.as_slice(), - [ - HavingExpr::Not(not), - HavingExpr::Compare { - operator: CompareOperator::GreaterOrEqual, - .. - } - ] if matches!( - not.as_ref(), - HavingExpr::Or(or) if matches!( - or.as_slice(), - [ - HavingExpr::Compare { - operator: CompareOperator::LessOrEqual, - .. - }, - HavingExpr::IsNull { negated: true, .. } - ] - ) - ) + plan.expression(), + HavingExpr::Compare { + operator: CompareOperator::Greater, + .. + } )); assert_eq!( sql, - "SELECT customer_id, sum(price)::numeric, count(*) AS __pgdog_having_0, \ - $1 AS __pgdog_having_1, sum(price) AS __pgdog_having_2, \ - max(price)::numeric AS __pgdog_having_3, 10 AS __pgdog_having_4, \ - customer_id AS __pgdog_having_5 FROM purchases GROUP BY customer_id \ - HAVING (NOT (count(*) <= $1 OR sum(price) IS NOT NULL) AND \ - max(price)::numeric >= 10) OR true" + "SELECT customer_id, count(*), count(*) AS __pgdog_having_0, \ + $1 AS __pgdog_having_1, customer_id AS __pgdog_having_2 \ + FROM purchases GROUP BY customer_id HAVING count(*) > $1 OR true" ); - assert_eq!(plan.hidden_aliases().len(), 6); + assert_eq!(plan.hidden_aliases().len(), 3); + } + + #[test] + fn rejects_boolean_null_test_and_cast_predicates() { + for sql in [ + "SELECT customer_id FROM purchases GROUP BY customer_id \ + HAVING COUNT(*) > 1 AND COUNT(*) < 5", + "SELECT customer_id FROM purchases GROUP BY customer_id \ + HAVING MAX(price) IS NULL", + "SELECT customer_id FROM purchases GROUP BY customer_id \ + HAVING COUNT(*)::numeric > 3", + ] { + let ast = pg_raw_parse::parse(sql).unwrap(); + let Node::SelectStmt(select) = ast.stmts().next().unwrap() else { + panic!("expected SELECT") + }; + let result = make::try_owned(|mem| { + let mut select = mem.make_unique(select); + rewrite_select(&mut select.as_mut(), mem).map(|_| select) + }); + assert!(matches!(result, Err(Error::UnsupportedHaving(_))), "{sql}"); + } } #[test] From d9348d592cff3c80f0481d15bd8cb8bbbd4ba584 Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Tue, 27 Oct 2026 23:33:00 +0530 Subject: [PATCH 14/24] trim --- .../backend/pool/connection/buffer/having.rs | 24 +------- pgdog/src/frontend/router/parser/aggregate.rs | 26 ++------ .../router/parser/rewrite/statement/having.rs | 60 +++---------------- pgdog/src/frontend/router/parser/route.rs | 4 -- 4 files changed, 16 insertions(+), 98 deletions(-) diff --git a/pgdog/src/backend/pool/connection/buffer/having.rs b/pgdog/src/backend/pool/connection/buffer/having.rs index fa3450cd0..b1256660f 100644 --- a/pgdog/src/backend/pool/connection/buffer/having.rs +++ b/pgdog/src/backend/pool/connection/buffer/having.rs @@ -131,13 +131,7 @@ fn integer(value: &Datum) -> Option { #[cfg(test)] mod tests { use super::*; - use crate::{ - frontend::router::parser::HavingValue, - net::{ - Field, RowDescription, - messages::{DataRow, Datum}, - }, - }; + use crate::net::{Field, RowDescription}; fn compare_expression(operator: CompareOperator) -> HavingExpr { HavingExpr::Compare { @@ -147,26 +141,14 @@ mod tests { } } - fn integer_field(name: &str) -> Field { - Field { - name: name.into(), - table_oid: 0, - column: 0, - type_oid: 23, - type_size: 4, - type_modifier: -1, - format: 0, - } - } - #[test] fn null_comparison_is_unknown_and_filters_the_group() { let decoder = Decoder::from(RowDescription::new(&[ Field::bigint("__pgdog_having_0"), - integer_field("__pgdog_having_1"), + Field::bigint("__pgdog_having_1"), ])); let mut row = DataRow::new(); - row.add(Datum::Null).add(2_i32); + row.add(Datum::Null).add(2_i64); let unknown = compare_expression(CompareOperator::Greater); assert_eq!( diff --git a/pgdog/src/frontend/router/parser/aggregate.rs b/pgdog/src/frontend/router/parser/aggregate.rs index 90f218252..853044b20 100644 --- a/pgdog/src/frontend/router/parser/aggregate.rs +++ b/pgdog/src/frontend/router/parser/aggregate.rs @@ -111,20 +111,11 @@ fn columns_match(group_by_names: &[&str], select_names: &[&str]) -> bool { impl AggregateFunction { pub(crate) fn from_name(name: &str, schema: &Schema) -> Option { - Some(match name { - "count" => AggregateFunction::Count, - "max" => AggregateFunction::Max, - "min" => AggregateFunction::Min, - "sum" => AggregateFunction::Sum, - "avg" => AggregateFunction::Avg, - "stddev" | "stddev_samp" => AggregateFunction::StddevSamp, - "stddev_pop" => AggregateFunction::StddevPop, - "variance" | "var_samp" => AggregateFunction::VarSamp, - "var_pop" => AggregateFunction::VarPop, - name if schema.aggregate_functions.contains(name) => { - AggregateFunction::Unrecognized(name.to_owned()) - } - _ => return None, + Self::from_builtin_name(name).or_else(|| { + schema + .aggregate_functions + .contains(name) + .then(|| Self::Unrecognized(name.to_owned())) }) } } @@ -157,12 +148,7 @@ impl Aggregate { .enumerate() .filter_map(|(idx, node)| { let func = Function::try_from(node.val()).ok()?; - let function = AggregateFunction::from_builtin_name(func.name).or_else(|| { - schema - .aggregate_functions - .contains(func.name) - .then(|| AggregateFunction::Unrecognized(func.name.to_owned())) - })?; + let function = AggregateFunction::from_name(func.name, schema)?; // FIXME: This doesn't correctly handle distinct behind type // casts any other nodes that Function::try_from handles diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/having.rs b/pgdog/src/frontend/router/parser/rewrite/statement/having.rs index 33811e0be..fce93adb7 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/having.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/having.rs @@ -13,23 +13,6 @@ use super::Error; const ALIAS_PREFIX: &str = "__pgdog_having_"; -struct ParsedHaving<'a> { - expression: HavingTemplate, - values: Vec<(String, Node<'a>)>, -} - -enum HavingTemplate { - Compare { - left: ValueTemplate, - operator: CompareOperator, - right: ValueTemplate, - }, -} - -struct ValueTemplate { - alias: String, -} - /// Project every supported predicate operand and disable shard-local filtering. /// /// The original expression remains under `OR TRUE` so PostgreSQL still performs @@ -82,17 +65,13 @@ pub(super) fn rewrite_select<'a>( } } } - let parsed = ParsedHaving { - expression, - values: parser.values, - }; - let hidden_aliases = parsed + let hidden_aliases = parser .values .iter() .map(|(alias, _)| alias.clone()) .collect::>(); - let targets = parsed + let targets = parser .values .iter() .map(|(alias, node)| mem.make_res_target(Some(alias), mem.empty(), mem.make_unique(*node))) @@ -112,10 +91,7 @@ pub(super) fn rewrite_select<'a>( .set_args(mem.make_list(&[original, true_value])); select.set_having_clause(neutralized.uncast()); - Ok(HavingRewritePlan::new( - parsed.expression.materialize(), - hidden_aliases, - )) + Ok(HavingRewritePlan::new(expression, hidden_aliases)) } fn target_output_name(node: Node<'_>) -> Option<&str> { @@ -127,28 +103,6 @@ fn target_output_name(node: Node<'_>) -> Option<&str> { } } -impl HavingTemplate { - fn materialize(self) -> HavingExpr { - match self { - Self::Compare { - left, - operator, - right, - } => HavingExpr::Compare { - left: left.materialize(), - operator, - right: right.materialize(), - }, - } - } -} - -impl ValueTemplate { - fn materialize(self) -> HavingValue { - HavingValue::new(self.alias) - } -} - struct Parser<'a> { aliases: HashSet, values: Vec<(String, Node<'a>)>, @@ -156,7 +110,7 @@ struct Parser<'a> { } impl<'a> Parser<'a> { - fn parse_expression(&mut self, node: Node<'a>) -> Result { + fn parse_expression(&mut self, node: Node<'a>) -> Result { match node { Node::A_Expr(expression) if expression.kind == nodes::A_Expr_Kind::AEXPR_OP => { let operator = expression @@ -171,7 +125,7 @@ impl<'a> Parser<'a> { "only =, <>, !=, <, <=, > and >= comparisons are supported", ) })?; - Ok(HavingTemplate::Compare { + Ok(HavingExpr::Compare { left: self.parse_value(expression.lexpr())?, operator, right: self.parse_value(expression.rexpr())?, @@ -183,7 +137,7 @@ impl<'a> Parser<'a> { } } - fn parse_value(&mut self, node: Node<'a>) -> Result { + fn parse_value(&mut self, node: Node<'a>) -> Result { if !Self::supported_value(node) { return Err(Error::UnsupportedHaving( "operands must be aggregate calls, grouped columns, constants or parameters", @@ -192,7 +146,7 @@ impl<'a> Parser<'a> { let alias = self.next_alias(); self.values.push((alias.clone(), node)); - Ok(ValueTemplate { alias }) + Ok(HavingValue::new(alias)) } fn supported_value(node: Node<'_>) -> bool { diff --git a/pgdog/src/frontend/router/parser/route.rs b/pgdog/src/frontend/router/parser/route.rs index b270058e1..624f71cbc 100644 --- a/pgdog/src/frontend/router/parser/route.rs +++ b/pgdog/src/frontend/router/parser/route.rs @@ -347,10 +347,6 @@ impl Route { pub(crate) fn requires_post_processing(&self) -> bool { !self.order_by().is_empty() || !self.aggregate().is_empty() - || self - .projection_rewrite_plan - .as_ref() - .is_some_and(|plan| plan.having.is_some()) || self.distinct().is_some() || self.limit().limit.is_some() || self.limit().offset.is_some() From 8b1d4ccbe3001d70f747ca163b7daaefed0e0cf1 Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Wed, 28 Oct 2026 11:54:40 +0530 Subject: [PATCH 15/24] feat: evaluate HAVING on merged cross-shard rows --- pgdog/src/backend/error.rs | 11 +- .../pool/connection/aggregate/error.rs | 6 + pgdog/src/backend/pool/connection/buffer.rs | 81 ++++++ .../backend/pool/connection/buffer/having.rs | 236 ++++++++++++++++++ .../pool/connection/multi_shard/mod.rs | 6 +- pgdog/src/frontend/router/parser/having.rs | 64 +++++ pgdog/src/frontend/router/parser/mod.rs | 2 + .../parser/rewrite/statement/projection.rs | 39 ++- 8 files changed, 436 insertions(+), 9 deletions(-) create mode 100644 pgdog/src/backend/pool/connection/buffer/having.rs create mode 100644 pgdog/src/frontend/router/parser/having.rs diff --git a/pgdog/src/backend/error.rs b/pgdog/src/backend/error.rs index 5696f3051..c3106a720 100644 --- a/pgdog/src/backend/error.rs +++ b/pgdog/src/backend/error.rs @@ -146,10 +146,13 @@ impl Error { pub(crate) fn is_unsupported(&self) -> bool { matches!( self, - Self::Aggregate(AggregateError::UnsupportedAggregation { .. }) - | Self::MultiShard(MultiShardError::Aggregate( - AggregateError::UnsupportedAggregation { .. } - )) + Self::Aggregate( + AggregateError::UnsupportedAggregation { .. } + | AggregateError::UnsupportedHaving(_) + ) | Self::MultiShard(MultiShardError::Aggregate( + AggregateError::UnsupportedAggregation { .. } + | AggregateError::UnsupportedHaving(_) + )) ) } diff --git a/pgdog/src/backend/pool/connection/aggregate/error.rs b/pgdog/src/backend/pool/connection/aggregate/error.rs index 4015d890e..b45d8578b 100644 --- a/pgdog/src/backend/pool/connection/aggregate/error.rs +++ b/pgdog/src/backend/pool/connection/aggregate/error.rs @@ -14,6 +14,12 @@ pub(crate) enum Error { #[error("unsupported aggregation {function}: {reason}")] UnsupportedAggregation { function: String, reason: String }, + #[error("unsupported cross-shard HAVING: {0}")] + UnsupportedHaving(String), + + #[error("invalid cross-shard HAVING state: {0}")] + InvalidHaving(String), + #[error("unexpected aggregation error: {0}")] Unexpected(String), } diff --git a/pgdog/src/backend/pool/connection/buffer.rs b/pgdog/src/backend/pool/connection/buffer.rs index d0ba21a3e..12b80fd8e 100644 --- a/pgdog/src/backend/pool/connection/buffer.rs +++ b/pgdog/src/backend/pool/connection/buffer.rs @@ -21,6 +21,8 @@ use pgdog_postgres_types::Datum; use super::Aggregates; +mod having; + /// Sort and aggregate rows received from multiple shards. #[derive(Default, Debug, Clone)] pub(super) struct Buffer { @@ -156,6 +158,30 @@ impl Buffer { Ok(()) } + /// Apply `HAVING` after all partial aggregate rows have been merged. + pub(super) fn having( + &mut self, + plan: Option<&ProjectionRewritePlan>, + decoder: &Decoder, + ) -> Result<(), super::Error> { + let Some(plan) = plan else { + return Ok(()); + }; + let Some(having) = plan.having.as_ref() else { + return Ok(()); + }; + let evaluator = having::Evaluator::new(having.expression(), decoder); + + let mut filtered = VecDeque::with_capacity(self.buffer.len()); + for row in self.buffer.drain(..) { + if evaluator.evaluate(&row, decoder)? { + filtered.push_back(row); + } + } + self.buffer = filtered; + Ok(()) + } + pub(super) fn drop_helper_columns( &mut self, plan: Option<&ProjectionRewritePlan>, @@ -251,8 +277,11 @@ impl Buffer { #[cfg(test)] mod test { use super::*; + use crate::frontend::router::parser::rewrite::statement::projection::HavingRewriteResult; + use crate::frontend::router::parser::{CompareOperator, HavingExpr, HavingValue}; use crate::net::{Datum, Field, Format, RowDescription}; use bytes::Bytes; + use pg_raw_parse::Node; #[test] fn test_sort_buffer() { @@ -334,6 +363,58 @@ mod test { } } + #[test] + fn having_filters_after_grouped_aggregate_merge() { + let ast = pg_raw_parse::parse( + "SELECT customer_id, COUNT(*), \ + COUNT(*) AS __pgdog_having_0, 3::bigint AS __pgdog_having_1 \ + FROM purchases GROUP BY customer_id", + ) + .unwrap(); + let Node::SelectStmt(select) = ast.stmts().next().unwrap() else { + panic!("not a select") + }; + let aggregate = Aggregate::parse(select, &Default::default()); + let decoder = Decoder::from(RowDescription::new(&[ + Field::bigint("customer_id"), + Field::bigint("count"), + Field::bigint("__pgdog_having_0"), + Field::bigint("__pgdog_having_1"), + ])); + + let plan = ProjectionRewritePlan { + having: Some(Box::new(HavingRewriteResult::new( + crate::frontend::router::parser::HavingRewritePlan::new( + HavingExpr::Compare { + left: HavingValue::new("__pgdog_having_0".into()), + operator: CompareOperator::Greater, + right: HavingValue::new("__pgdog_having_1".into()), + }, + vec!["__pgdog_having_0".into(), "__pgdog_having_1".into()], + ), + aggregate.clone(), + ))), + ..Default::default() + }; + + let mut buffer = Buffer::default(); + for (customer, partial_count) in [(1_i64, 2_i64), (1, 2), (2, 1), (2, 1)] { + let mut row = DataRow::new(); + row.add(customer) + .add(partial_count) + .add(partial_count) + .add(3_i64); + buffer.add(row.message()).unwrap(); + } + + buffer.aggregate(&aggregate, &decoder, Some(&plan)).unwrap(); + buffer.having(Some(&plan), &decoder).unwrap(); + + assert_eq!(buffer.len(), 1); + assert_eq!(buffer.buffer[0].get::(0, Format::Text).unwrap(), 1); + assert_eq!(buffer.buffer[0].get::(1, Format::Text).unwrap(), 4); + } + #[test] fn test_sort_buffer_with_timestamps() { let mut buf = Buffer::default(); diff --git a/pgdog/src/backend/pool/connection/buffer/having.rs b/pgdog/src/backend/pool/connection/buffer/having.rs new file mode 100644 index 000000000..b1256660f --- /dev/null +++ b/pgdog/src/backend/pool/connection/buffer/having.rs @@ -0,0 +1,236 @@ +use std::{cmp::Ordering, collections::HashMap}; + +use pgdog_postgres_types::Datum; + +use crate::{ + frontend::router::parser::{CompareOperator, HavingExpr, HavingValue}, + net::{Decoder, messages::DataRow}, +}; + +use super::super::aggregate::Error; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum TruthValue { + True, + False, + Unknown, +} + +pub(super) struct Evaluator<'a> { + expression: &'a HavingExpr, + columns: HashMap<&'a str, usize>, +} + +impl<'a> Evaluator<'a> { + pub(super) fn new(expression: &'a HavingExpr, decoder: &'a Decoder) -> Self { + let columns = decoder + .row_description() + .fields + .iter() + .enumerate() + .map(|(index, field)| (field.name.as_str(), index)) + .collect(); + Self { + expression, + columns, + } + } + + pub(super) fn evaluate(&self, row: &DataRow, decoder: &Decoder) -> Result { + Ok(matches!( + self.evaluate_expression(self.expression, row, decoder)?, + TruthValue::True + )) + } + + fn evaluate_expression( + &self, + expression: &HavingExpr, + row: &DataRow, + decoder: &Decoder, + ) -> Result { + match expression { + HavingExpr::Compare { + left, + operator, + right, + } => { + let left = self.value(left, row, decoder)?; + let right = self.value(right, row, decoder)?; + compare(&left, *operator, &right) + } + } + } + + fn value(&self, value: &HavingValue, row: &DataRow, decoder: &Decoder) -> Result { + let index = + self.columns.get(value.alias()).copied().ok_or_else(|| { + Error::InvalidHaving(format!("missing operand {}", value.alias())) + })?; + Ok(row.get_column_checked(index, decoder)?.value) + } +} + +#[cfg(test)] +fn evaluate(expression: &HavingExpr, row: &DataRow, decoder: &Decoder) -> Result { + Evaluator::new(expression, decoder).evaluate(row, decoder) +} + +#[cfg(test)] +fn evaluate_expression( + expression: &HavingExpr, + row: &DataRow, + decoder: &Decoder, +) -> Result { + Evaluator::new(expression, decoder).evaluate_expression(expression, row, decoder) +} + +fn compare(left: &Datum, operator: CompareOperator, right: &Datum) -> Result { + if left.is_null() || right.is_null() { + return Ok(TruthValue::Unknown); + } + + let ordering = numeric_ordering(left, right) + .or_else(|| left.partial_cmp(right)) + .ok_or_else(|| { + Error::UnsupportedHaving(format!( + "cannot compare {} with {}", + left.data_type(), + right.data_type() + )) + })?; + let result = match operator { + CompareOperator::Equal => ordering == Ordering::Equal, + CompareOperator::NotEqual => ordering != Ordering::Equal, + CompareOperator::Less => ordering == Ordering::Less, + CompareOperator::LessOrEqual => ordering != Ordering::Greater, + CompareOperator::Greater => ordering == Ordering::Greater, + CompareOperator::GreaterOrEqual => ordering != Ordering::Less, + }; + Ok(if result { + TruthValue::True + } else { + TruthValue::False + }) +} + +fn numeric_ordering(left: &Datum, right: &Datum) -> Option { + let (left, right) = (integer(left)?, integer(right)?); + Some(left.cmp(&right)) +} + +fn integer(value: &Datum) -> Option { + match value { + Datum::Bigint(value) => Some((*value).into()), + Datum::Integer(value) => Some((*value).into()), + Datum::SmallInt(value) => Some((*value).into()), + _ => None, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::net::{Field, RowDescription}; + + fn compare_expression(operator: CompareOperator) -> HavingExpr { + HavingExpr::Compare { + left: HavingValue::new("__pgdog_having_0".into()), + operator, + right: HavingValue::new("__pgdog_having_1".into()), + } + } + + #[test] + fn null_comparison_is_unknown_and_filters_the_group() { + let decoder = Decoder::from(RowDescription::new(&[ + Field::bigint("__pgdog_having_0"), + Field::bigint("__pgdog_having_1"), + ])); + let mut row = DataRow::new(); + row.add(Datum::Null).add(2_i64); + + let unknown = compare_expression(CompareOperator::Greater); + assert_eq!( + evaluate_expression(&unknown, &row, &decoder).unwrap(), + TruthValue::Unknown + ); + assert!(!evaluate(&unknown, &row, &decoder).unwrap()); + + let missing = HavingExpr::Compare { + left: HavingValue::new("__pgdog_missing".into()), + operator: CompareOperator::Equal, + right: HavingValue::new("__pgdog_having_1".into()), + }; + assert!(matches!( + evaluate(&missing, &row, &decoder), + Err(Error::InvalidHaving(_)) + )); + } + + #[test] + fn compares_integers_across_widths_with_all_operators() { + for (operator, left, right, expected) in [ + ( + CompareOperator::Equal, + Datum::Bigint(2), + Datum::Integer(2), + true, + ), + ( + CompareOperator::NotEqual, + Datum::Integer(2), + Datum::SmallInt(3), + true, + ), + ( + CompareOperator::Less, + Datum::SmallInt(2), + Datum::Bigint(3), + true, + ), + ( + CompareOperator::LessOrEqual, + Datum::Bigint(2), + Datum::Integer(2), + true, + ), + ( + CompareOperator::Greater, + Datum::Bigint(3), + Datum::SmallInt(2), + true, + ), + ( + CompareOperator::GreaterOrEqual, + Datum::Integer(2), + Datum::Bigint(2), + true, + ), + ( + CompareOperator::Greater, + Datum::Text("a".into()), + Datum::Text("z".into()), + false, + ), + ] { + assert_eq!( + compare(&left, operator, &right).unwrap(), + if expected { + TruthValue::True + } else { + TruthValue::False + } + ); + } + + assert!(matches!( + compare( + &Datum::Text("2".into()), + CompareOperator::Equal, + &Datum::Integer(2) + ), + Err(Error::UnsupportedHaving(_)) + )); + } +} diff --git a/pgdog/src/backend/pool/connection/multi_shard/mod.rs b/pgdog/src/backend/pool/connection/multi_shard/mod.rs index 5eafe91fd..32da11c96 100644 --- a/pgdog/src/backend/pool/connection/multi_shard/mod.rs +++ b/pgdog/src/backend/pool/connection/multi_shard/mod.rs @@ -247,11 +247,11 @@ impl MultiShard { // Helpers remain in the internal row through aggregation and // sorting, then are removed before client-visible operations. let plan = self.route.projection_rewrite_plan.as_deref(); - if let Err(err) = self + let result = self .buffer .aggregate(self.route.aggregate(), &self.decoder, plan) - { - // Probably an unsupported aggregation feature was requested. + .and_then(|_| self.buffer.having(plan, &self.decoder)); + if let Err(err) = result { self.buffer.clear(); return Err(err.into()); } diff --git a/pgdog/src/frontend/router/parser/having.rs b/pgdog/src/frontend/router/parser/having.rs new file mode 100644 index 000000000..b007bc341 --- /dev/null +++ b/pgdog/src/frontend/router/parser/having.rs @@ -0,0 +1,64 @@ +//! Cross-shard `HAVING` execution plan. +// Only tests build plans until the HAVING rewriter is wired in. +#![cfg_attr(not(test), allow(dead_code))] + +/// Comparison operators supported by the cross-shard `HAVING` evaluator. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum CompareOperator { + Equal, + NotEqual, + Less, + LessOrEqual, + Greater, + GreaterOrEqual, +} + +/// A value projected by the cross-shard query under an internal alias. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct HavingValue { + alias: String, +} + +impl HavingValue { + pub(crate) fn new(alias: String) -> Self { + Self { alias } + } + + pub(crate) fn alias(&self) -> &str { + &self.alias + } +} + +/// Predicate evaluated after partial aggregate results have been merged. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum HavingExpr { + Compare { + left: HavingValue, + operator: CompareOperator, + right: HavingValue, + }, +} + +/// Metadata produced while rewriting a cross-shard `HAVING` clause. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct HavingRewritePlan { + expression: HavingExpr, + hidden_aliases: Vec, +} + +impl HavingRewritePlan { + pub(crate) fn new(expression: HavingExpr, hidden_aliases: Vec) -> Self { + Self { + expression, + hidden_aliases, + } + } + + pub(crate) fn expression(&self) -> &HavingExpr { + &self.expression + } + + pub(crate) fn hidden_aliases(&self) -> &[String] { + &self.hidden_aliases + } +} diff --git a/pgdog/src/frontend/router/parser/mod.rs b/pgdog/src/frontend/router/parser/mod.rs index a74d7a01f..4eb21227a 100644 --- a/pgdog/src/frontend/router/parser/mod.rs +++ b/pgdog/src/frontend/router/parser/mod.rs @@ -16,6 +16,7 @@ pub(crate) mod error; pub(crate) mod explain_trace; mod from_clause; pub(crate) mod function; +mod having; pub(crate) mod key; mod limit; pub(crate) mod multi_tenant; @@ -45,6 +46,7 @@ pub(crate) use distinct::{Distinct, DistinctBy, DistinctColumn}; pub(crate) use error::Error; pub(crate) use from_clause::FromClause; use function::Function; +pub(crate) use having::{CompareOperator, HavingExpr, HavingRewritePlan, HavingValue}; pub(crate) use key::Key; pub(crate) use limit::{Limit, LimitClause}; pub(crate) use order_by::OrderBy; diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs b/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs index 80fd28029..55c75157c 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs @@ -9,7 +9,7 @@ use super::aggregate::{AggregatesRewrite, HelperKind}; use super::offset::{self, OffsetPlan}; use super::order_by; use crate::backend::schema::Schema; -use crate::frontend::router::parser::{Aggregate, OrderBy}; +use crate::frontend::router::parser::{Aggregate, HavingExpr, HavingRewritePlan, OrderBy}; use crate::frontend::{ClientRequest, PreparedStatements}; use crate::net::{ProtocolMessage, RowDescription}; use pg_raw_parse::{Node, StmtList, make}; @@ -57,11 +57,41 @@ impl OrderByHelper { pub(crate) struct ProjectionRewritePlan { pub(crate) aggregate_helpers: Vec, pub(crate) order_by_helpers: Vec, + pub(crate) having: Option>, +} + +/// The filter and aggregate metadata produced by rewriting a HAVING query. +#[derive(Debug, PartialEq)] +pub(crate) struct HavingRewriteResult { + rewrite: HavingRewritePlan, + aggregate: Aggregate, +} + +impl HavingRewriteResult { + #[cfg_attr(not(test), allow(dead_code))] + pub(crate) fn new(rewrite: HavingRewritePlan, aggregate: Aggregate) -> Self { + Self { rewrite, aggregate } + } + + pub(crate) fn expression(&self) -> &HavingExpr { + self.rewrite.expression() + } + + pub(crate) fn hidden_aliases(&self) -> &[String] { + self.rewrite.hidden_aliases() + } + + #[allow(dead_code)] + pub(crate) fn aggregate(&self) -> &Aggregate { + &self.aggregate + } } impl ProjectionRewritePlan { pub(crate) fn is_noop(&self) -> bool { - self.aggregate_helpers.is_empty() && self.order_by_helpers.is_empty() + self.aggregate_helpers.is_empty() + && self.order_by_helpers.is_empty() + && self.having.is_none() } pub(crate) fn drop_columns(&self, row_description: &RowDescription) -> BTreeSet { @@ -74,6 +104,11 @@ impl ProjectionRewritePlan { .filter(|helper| helper.injected) .map(|helper| helper.alias.as_str()), ) + .chain( + self.having + .iter() + .flat_map(|plan| plan.hidden_aliases().iter().map(String::as_str)), + ) .filter_map(|alias| row_description.field_index(alias)) .collect() } From 3d882251a80caab5d8be8d26570b6cae40700a87 Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Wed, 28 Oct 2026 22:06:09 +0530 Subject: [PATCH 16/24] resolve comments --- pgdog-postgres-types/src/datum.rs | 37 ++++++ .../backend/pool/connection/buffer/having.rs | 17 +-- .../pool/connection/multi_shard/test.rs | 106 +++++++++++++++++- 3 files changed, 143 insertions(+), 17 deletions(-) diff --git a/pgdog-postgres-types/src/datum.rs b/pgdog-postgres-types/src/datum.rs index 8978fda71..75d426a42 100644 --- a/pgdog-postgres-types/src/datum.rs +++ b/pgdog-postgres-types/src/datum.rs @@ -250,6 +250,21 @@ impl Datum { }), } } + + /// Orders numeric values of different widths the way PostgreSQL does. + /// Returns `None` if either side isn't numeric. + pub fn compare_numeric(&self, other: &Datum) -> Option { + Some(self.as_i128()?.cmp(&other.as_i128()?)) + } + + fn as_i128(&self) -> Option { + match self { + Datum::Bigint(value) => Some((*value).into()), + Datum::Integer(value) => Some((*value).into()), + Datum::SmallInt(value) => Some((*value).into()), + _ => None, + } + } } /// PostgreSQL data types. @@ -344,6 +359,28 @@ mod tests { use bytes::{BufMut, BytesMut}; use std::assert_matches; + #[test] + fn test_compare_numeric_across_integer_widths() { + for (left, right, expected) in [ + (Datum::Bigint(2), Datum::Integer(2), Ordering::Equal), + (Datum::Integer(2), Datum::SmallInt(3), Ordering::Less), + (Datum::SmallInt(3), Datum::Bigint(2), Ordering::Greater), + ( + Datum::Bigint(i64::MAX), + Datum::Integer(i32::MAX), + Ordering::Greater, + ), + ] { + assert_eq!(left.compare_numeric(&right), Some(expected)); + } + + assert_eq!( + Datum::Text("2".into()).compare_numeric(&Datum::Integer(2)), + None + ); + assert_eq!(Datum::Null.compare_numeric(&Datum::Integer(2)), None); + } + #[test] fn test_multidimensional_text_array_falls_back_to_unknown() { let input = b"{{1,2},{3,4}}"; diff --git a/pgdog/src/backend/pool/connection/buffer/having.rs b/pgdog/src/backend/pool/connection/buffer/having.rs index b1256660f..f6d8801df 100644 --- a/pgdog/src/backend/pool/connection/buffer/having.rs +++ b/pgdog/src/backend/pool/connection/buffer/having.rs @@ -90,7 +90,8 @@ fn compare(left: &Datum, operator: CompareOperator, right: &Datum) -> Result Result Option { - let (left, right) = (integer(left)?, integer(right)?); - Some(left.cmp(&right)) -} - -fn integer(value: &Datum) -> Option { - match value { - Datum::Bigint(value) => Some((*value).into()), - Datum::Integer(value) => Some((*value).into()), - Datum::SmallInt(value) => Some((*value).into()), - _ => None, - } -} - #[cfg(test)] mod tests { use super::*; diff --git a/pgdog/src/backend/pool/connection/multi_shard/test.rs b/pgdog/src/backend/pool/connection/multi_shard/test.rs index f6cc5ccfa..733d3f191 100644 --- a/pgdog/src/backend/pool/connection/multi_shard/test.rs +++ b/pgdog/src/backend/pool/connection/multi_shard/test.rs @@ -1,7 +1,12 @@ +use std::sync::Arc; + use crate::{ frontend::router::parser::{ - DistinctBy, OrderBy, Shard, ShardWithPriority, - rewrite::statement::projection::{OrderByHelper, OrderBySource, ProjectionRewritePlan}, + Aggregate, CompareOperator, DistinctBy, HavingExpr, HavingRewritePlan, HavingValue, + OrderBy, Shard, ShardWithPriority, + rewrite::statement::projection::{ + HavingRewriteResult, OrderByHelper, OrderBySource, ProjectionRewritePlan, + }, }, net::{BindComplete, DataRow, Field, Format}, }; @@ -137,6 +142,103 @@ fn test_order_by_helper_after_star_expansion_is_dropped_after_sorting() { } } +/// `SELECT customer_id, COUNT(*) ... GROUP BY customer_id HAVING > 3`, +/// with `3` projected as the hidden `__pgdog_having_0` column. +fn having_route(left: &str) -> Route { + let aggregate = Aggregate::new_count_group_by(1, &[0]); + let plan = ProjectionRewritePlan { + having: Some(Box::new(HavingRewriteResult::new( + HavingRewritePlan::new( + HavingExpr::Compare { + left: HavingValue::new(left.into()), + operator: CompareOperator::Greater, + right: HavingValue::new("__pgdog_having_0".into()), + }, + vec!["__pgdog_having_0".into()], + ), + aggregate.clone(), + ))), + ..Default::default() + }; + let mut route = Route::select( + ShardWithPriority::new_default_unset(Shard::All), + vec![], + aggregate, + Default::default(), + None, + ); + route.projection_rewrite_plan = Some(Arc::new(plan)); + route +} + +/// Sends a row description, two groups per shard and `CommandComplete` from both shards. +/// Customer 1 has a partial count of 2 on each shard, customer 2 has 1. +fn send_having_shard_results(multi_shard: &mut MultiShard) -> Result, Error> { + let row_description = RowDescription::new(&[ + Field::bigint("customer_id"), + Field::bigint("count"), + Field::bigint("__pgdog_having_0"), + ]); + let mut forwarded = Vec::new(); + for _ in 0..2 { + forwarded.extend(multi_shard.handle_server_message(row_description.message())?); + } + for _ in 0..2 { + for (customer, count) in [(1_i64, 2_i64), (2, 1)] { + let mut row = DataRow::new(); + row.add(customer).add(count).add(3_i64); + forwarded.extend(multi_shard.handle_server_message(row.message())?); + } + } + for _ in 0..2 { + forwarded.extend( + multi_shard.handle_server_message(CommandComplete::from_str("SELECT 2").message())?, + ); + } + Ok(forwarded) +} + +#[test] +fn test_having_filters_merged_groups_and_drops_hidden_columns() { + let route = having_route("count"); + let mut multi_shard = MultiShard::new(2, &route); + + let forwarded = send_having_shard_results(&mut multi_shard).unwrap(); + let [client_description] = forwarded.as_slice() else { + panic!("expected only the row description to be forwarded") + }; + let client_description = RowDescription::from_bytes(client_description.to_bytes()).unwrap(); + assert_eq!( + client_description + .fields + .iter() + .map(|field| field.name.as_str()) + .collect::>(), + ["customer_id", "count"] + ); + + // No shard passes COUNT(*) > 3 on its own; customer 1 does only after merging (2 + 2). + let row = DataRow::from_bytes(multi_shard.get_server_message().unwrap().to_bytes()).unwrap(); + assert_eq!(row.len(), 2); + assert_eq!(row.get::(0, Format::Text).unwrap(), 1); + assert_eq!(row.get::(1, Format::Text).unwrap(), 4); + + let command_complete = + CommandComplete::from_bytes(multi_shard.get_server_message().unwrap().to_bytes()).unwrap(); + assert_eq!(command_complete.rows().unwrap(), Some(1)); + assert!(multi_shard.get_server_message().is_none()); +} + +#[test] +fn test_having_evaluation_error_clears_buffer() { + let route = having_route("__pgdog_missing"); + let mut multi_shard = MultiShard::new(2, &route); + + let error = send_having_shard_results(&mut multi_shard).unwrap_err(); + assert!(error.to_string().contains("__pgdog_missing"), "{error}"); + assert!(!multi_shard.buffer.can_take()); +} + #[test] fn test_rd_before_dr() { let mut multi_shard = MultiShard::new( From c35e0f976a88abd02496bf1907a1497efad71f59 Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Wed, 28 Oct 2026 22:30:50 +0530 Subject: [PATCH 17/24] fmt --- pgdog/src/backend/error.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/pgdog/src/backend/error.rs b/pgdog/src/backend/error.rs index 76b446adf..68c18461c 100644 --- a/pgdog/src/backend/error.rs +++ b/pgdog/src/backend/error.rs @@ -150,10 +150,10 @@ impl Error { Self::Aggregate( AggregateError::UnsupportedAggregation { .. } | AggregateError::UnsupportedHaving(_) - ) | Self::MultiShard(MultiShardError::Aggregate( - AggregateError::UnsupportedAggregation { .. } - | AggregateError::UnsupportedHaving(_) - )) | Self::Pool(PoolError::PubSubDisabled) + ) | Self::MultiShard(MultiShardError::Aggregate( + AggregateError::UnsupportedAggregation { .. } + | AggregateError::UnsupportedHaving(_) + )) | Self::Pool(PoolError::PubSubDisabled) ) } From 9c9786af292796ac2b96da543f9c99f588e803b3 Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Sun, 1 Nov 2026 00:12:45 +0530 Subject: [PATCH 18/24] fixes --- integration/rust/tests/integration/having.rs | 1 + .../parser/rewrite/statement/projection.rs | 3 --- pgdog/src/frontend/router/parser/route.rs | 16 +++++++++------- 3 files changed, 10 insertions(+), 10 deletions(-) diff --git a/integration/rust/tests/integration/having.rs b/integration/rust/tests/integration/having.rs index 4fd80c360..5186a6d0d 100644 --- a/integration/rust/tests/integration/having.rs +++ b/integration/rust/tests/integration/having.rs @@ -68,6 +68,7 @@ async fn prepared_having_parameter_is_typed_by_postgres() -> Result<(), Box("customer_id"), 1); assert_eq!(rows[0].get::("order_count"), 3); diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs b/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs index 47941296d..41d846f93 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs @@ -182,9 +182,6 @@ pub(crate) fn finalize_after_route( && let Some(route) = request.route.as_mut() { route.projection_rewrite_plan = Some(Arc::clone(plan)); - if let Some(having) = &plan.having { - route.set_aggregate(having.aggregate().clone()); - } let mut order_by = route.order_by().to_vec(); for helper in &plan.order_by_helpers { // Prefer the structural position. The source fallback handles a diff --git a/pgdog/src/frontend/router/parser/route.rs b/pgdog/src/frontend/router/parser/route.rs index 624f71cbc..897e7f176 100644 --- a/pgdog/src/frontend/router/parser/route.rs +++ b/pgdog/src/frontend/router/parser/route.rs @@ -1,8 +1,10 @@ use std::{fmt::Display, ops::Deref, sync::Arc}; use super::{ - Aggregate, DistinctBy, Limit, OrderBy, StatementType, advisory_lock::AdvisoryLocks, - explain_trace::ExplainTrace, rewrite::statement::projection::ProjectionRewritePlan, + Aggregate, DistinctBy, Limit, OrderBy, StatementType, + advisory_lock::AdvisoryLocks, + explain_trace::ExplainTrace, + rewrite::statement::projection::{HavingRewriteResult, ProjectionRewritePlan}, }; use crate::frontend::{client::query_engine::TempTableChange, router::sharding::PendingLookup}; use lazy_static::lazy_static; @@ -270,11 +272,11 @@ impl Route { } pub(crate) fn aggregate(&self) -> &Aggregate { - &self.aggregate - } - - pub(crate) fn set_aggregate(&mut self, aggregate: Aggregate) { - self.aggregate = aggregate; + self.projection_rewrite_plan + .as_deref() + .and_then(|plan| plan.having.as_deref()) + .map(HavingRewriteResult::aggregate) + .unwrap_or(&self.aggregate) } /// Set shard on this route, along with reasoning From 99059badcd132ce45af2d264fe498132cad062e7 Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Sun, 1 Nov 2026 00:37:15 +0530 Subject: [PATCH 19/24] helper merge --- .../router/parser/rewrite/statement/having.rs | 18 ++------------ .../parser/rewrite/statement/order_by.rs | 24 +++++++++++++++---- 2 files changed, 21 insertions(+), 21 deletions(-) diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/having.rs b/pgdog/src/frontend/router/parser/rewrite/statement/having.rs index fce93adb7..5e78c94c1 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/having.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/having.rs @@ -9,7 +9,7 @@ use crate::frontend::router::parser::{ AggregateFunction, CompareOperator, Function, HavingExpr, HavingRewritePlan, HavingValue, }; -use super::Error; +use super::{Error, order_by::target_output_name}; const ALIAS_PREFIX: &str = "__pgdog_having_"; @@ -32,12 +32,7 @@ pub(super) fn rewrite_select<'a>( let aliases = select .target_list() .iter() - .filter_map(|target| { - target - .name() - .or_else(|| target_output_name(target.val())) - .map(str::to_owned) - }) + .filter_map(|target| target_output_name(target).map(str::to_owned)) .collect(); let mut parser = Parser { aliases, @@ -94,15 +89,6 @@ pub(super) fn rewrite_select<'a>( Ok(HavingRewritePlan::new(expression, hidden_aliases)) } -fn target_output_name(node: Node<'_>) -> Option<&str> { - match node { - Node::ColumnRef(column) => column.fields().iter().next_back()?.as_str(), - Node::FuncCall(function) => function.funcname().iter().next_back()?.as_str(), - Node::TypeCast(cast) => target_output_name(cast.arg()), - _ => None, - } -} - struct Parser<'a> { aliases: HashSet, values: Vec<(String, Node<'a>)>, diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/order_by.rs b/pgdog/src/frontend/router/parser/rewrite/statement/order_by.rs index d1c4bf534..a81ceb161 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/order_by.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/order_by.rs @@ -59,12 +59,18 @@ fn star_covers(star: &nodes::ColumnRef, column: &nodes::ColumnRef) -> bool { .map(Node::as_str)))) } -fn target_output_name(target: &nodes::ResTarget) -> Option<&str> { - if let Some(alias) = target.name() { - return Some(alias); - } - match target.val() { +/// Column name PostgreSQL assigns to a SELECT target in the RowDescription. +pub(super) fn target_output_name(target: &nodes::ResTarget) -> Option<&str> { + target + .name() + .or_else(|| expression_output_name(target.val())) +} + +fn expression_output_name(node: Node<'_>) -> Option<&str> { + match node { Node::ColumnRef(column) if !is_star(column) => column_name(column), + Node::FuncCall(function) => function.funcname().iter().next_back()?.as_str(), + Node::TypeCast(cast) => expression_output_name(cast.arg()), _ => None, } } @@ -404,6 +410,14 @@ mod tests { assert!(plan.order_by_helpers[0].injected); } + #[test] + fn injects_helper_when_cast_target_shares_output_name() { + let (sql, plan) = rewrite("SELECT price::int, price FROM purchases ORDER BY price"); + + assert!(sql.contains("price AS __pgdog_order_col0"), "{sql}"); + assert!(plan.order_by_helpers[0].injected); + } + #[test] fn injects_helper_for_unqualified_star_join() { let (sql, plan) = rewrite("SELECT * FROM a JOIN b ON a.id = b.a_id ORDER BY b.price"); From 31a63edd20bd39958a4d30738655af174c51025e Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Sun, 1 Nov 2026 00:51:33 +0530 Subject: [PATCH 20/24] refactor --- .../router/parser/rewrite/statement/having.rs | 2 +- .../parser/rewrite/statement/order_by.rs | 35 ++----------------- .../parser/rewrite/statement/projection.rs | 33 ++++++++++++++++- 3 files changed, 36 insertions(+), 34 deletions(-) diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/having.rs b/pgdog/src/frontend/router/parser/rewrite/statement/having.rs index 5e78c94c1..4ac3939d5 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/having.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/having.rs @@ -9,7 +9,7 @@ use crate::frontend::router::parser::{ AggregateFunction, CompareOperator, Function, HavingExpr, HavingRewritePlan, HavingValue, }; -use super::{Error, order_by::target_output_name}; +use super::{Error, projection::target_output_name}; const ALIAS_PREFIX: &str = "__pgdog_having_"; diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/order_by.rs b/pgdog/src/frontend/router/parser/rewrite/statement/order_by.rs index a81ceb161..60aea9925 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/order_by.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/order_by.rs @@ -4,7 +4,9 @@ use crate::backend::schema::Schema; use crate::frontend::router::parser::aggregate::AggregateFunction; use crate::frontend::router::parser::{Column, Function, OrderBy}; -use super::projection::{OrderByHelper, OrderBySource, ProjectionRewritePlan}; +use super::projection::{ + OrderByHelper, OrderBySource, ProjectionRewritePlan, column_name, is_star, target_output_name, +}; impl OrderBySource { fn name(&self) -> &str { @@ -14,21 +16,6 @@ impl OrderBySource { } } -fn column_name(column: &nodes::ColumnRef) -> Option<&str> { - column - .fields() - .into_iter() - .next_back() - .and_then(Node::as_str) -} - -fn is_star(column: &nodes::ColumnRef) -> bool { - matches!( - column.fields().into_iter().next_back(), - Some(Node::A_Star(_)) - ) -} - /// Compare complete column references so `a.price` and `b.price` remain distinct. fn same_column(left: &nodes::ColumnRef, right: &nodes::ColumnRef) -> bool { left.fields() @@ -59,22 +46,6 @@ fn star_covers(star: &nodes::ColumnRef, column: &nodes::ColumnRef) -> bool { .map(Node::as_str)))) } -/// Column name PostgreSQL assigns to a SELECT target in the RowDescription. -pub(super) fn target_output_name(target: &nodes::ResTarget) -> Option<&str> { - target - .name() - .or_else(|| expression_output_name(target.val())) -} - -fn expression_output_name(node: Node<'_>) -> Option<&str> { - match node { - Node::ColumnRef(column) if !is_star(column) => column_name(column), - Node::FuncCall(function) => function.funcname().iter().next_back()?.as_str(), - Node::TypeCast(cast) => expression_output_name(cast.arg()), - _ => None, - } -} - fn single_relation(select: &nodes::SelectStmtMut<'_, '_>) -> bool { let from = select.from_clause(); from.len() == 1 && matches!(from.first(), Some(Node::RangeVar(_))) diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs b/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs index 41d846f93..1a0288054 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs @@ -13,7 +13,7 @@ use crate::backend::schema::Schema; use crate::frontend::router::parser::{Aggregate, HavingExpr, HavingRewritePlan, OrderBy}; use crate::frontend::{ClientRequest, PreparedStatements}; use crate::net::{ProtocolMessage, RowDescription}; -use pg_raw_parse::{Node, StmtList, make}; +use pg_raw_parse::{Node, StmtList, make, nodes}; use std::collections::BTreeSet; use std::sync::Arc; @@ -86,6 +86,37 @@ impl HavingRewriteResult { } } +pub(super) fn column_name(column: &nodes::ColumnRef) -> Option<&str> { + column + .fields() + .into_iter() + .next_back() + .and_then(Node::as_str) +} + +pub(super) fn is_star(column: &nodes::ColumnRef) -> bool { + matches!( + column.fields().into_iter().next_back(), + Some(Node::A_Star(_)) + ) +} + +/// Column name PostgreSQL assigns to a SELECT target in the RowDescription. +pub(super) fn target_output_name(target: &nodes::ResTarget) -> Option<&str> { + target + .name() + .or_else(|| expression_output_name(target.val())) +} + +fn expression_output_name(node: Node<'_>) -> Option<&str> { + match node { + Node::ColumnRef(column) if !is_star(column) => column_name(column), + Node::FuncCall(function) => function.funcname().iter().next_back()?.as_str(), + Node::TypeCast(cast) => expression_output_name(cast.arg()), + _ => None, + } +} + impl ProjectionRewritePlan { pub(crate) fn is_noop(&self) -> bool { self.aggregate_helpers.is_empty() From 8796c372109c8da9af0257b21a6caab70678360b Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Sun, 1 Nov 2026 01:47:32 +0530 Subject: [PATCH 21/24] testing --- integration/rust/tests/integration/having.rs | 65 ++++++++++++++++++- .../frontend/router/parser/query/prepare.rs | 18 ++++- .../router/parser/rewrite/statement/having.rs | 60 ++++++++++++++--- .../parser/rewrite/statement/projection.rs | 38 +++++++++++ 4 files changed, 167 insertions(+), 14 deletions(-) diff --git a/integration/rust/tests/integration/having.rs b/integration/rust/tests/integration/having.rs index 5186a6d0d..a34988724 100644 --- a/integration/rust/tests/integration/having.rs +++ b/integration/rust/tests/integration/having.rs @@ -1,4 +1,4 @@ -use crate::setup::{admin_sqlx, connections_sqlx}; +use crate::setup::{admin_sqlx, connections_sqlx, connections_tokio}; use sqlx::{Executor, Row, postgres::PgPool}; #[tokio::test] @@ -34,13 +34,22 @@ async fn having_filters_merged_groups() -> Result<(), Box assert_eq!(rows[0].get::("customer_id"), 1); assert_eq!(rows[0].get::("order_count"), 3); + // An untyped literal takes the type of the aggregate it is compared with. + let rows = sharded + .fetch_all( + "SELECT customer_id FROM having_merge_test \ + GROUP BY customer_id HAVING MAX(order_id) > '20'", + ) + .await?; + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].get::("customer_id"), 2); + cleanup(&sharded).await; Ok(()) } #[tokio::test] -async fn prepared_having_parameter_is_typed_by_postgres() -> Result<(), Box> -{ +async fn prepared_having_parameter_typed_and_untyped() -> Result<(), Box> { let sharded = connections_sqlx().await.get(1).cloned().unwrap(); reset_table(&sharded).await?; @@ -72,6 +81,56 @@ async fn prepared_having_parameter_is_typed_by_postgres() -> Result<(), Box("customer_id"), 1); assert_eq!(rows[0].get::("order_count"), 3); + // Parse without parameter types, so PostgreSQL has to infer $1. + let client = connections_tokio().await.remove(1); + let statement = client + .prepare( + "SELECT customer_id FROM having_merge_test \ + GROUP BY customer_id HAVING COUNT(*) > $1", + ) + .await?; + let rows = client.query(&statement, &[&2_i64]).await?; + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].get::<_, i64>(0), 1); + + cleanup(&sharded).await; + Ok(()) +} + +#[tokio::test] +async fn having_rejects_shard_local_steps() -> Result<(), Box> { + let sharded = connections_sqlx().await.get(1).cloned().unwrap(); + reset_table(&sharded).await?; + let client = connections_tokio().await.remove(1); + + // Each of these would run on shards before PgDog merges the groups. + for sql in [ + "SELECT * FROM having_merge_test GROUP BY customer_id, order_id HAVING COUNT(*) > 0", + "SELECT customer_id FROM having_merge_test GROUP BY customer_id \ + HAVING COUNT(*) > 0 ORDER BY customer_id LIMIT 1", + "SELECT customer_id FROM having_merge_test GROUP BY customer_id \ + HAVING COUNT(*) > 0 OFFSET 1", + "SELECT DISTINCT ON (customer_id) customer_id FROM having_merge_test \ + GROUP BY customer_id HAVING COUNT(*) > 0", + "EXECUTE having_prepared(0)", + ] { + if sql.starts_with("EXECUTE") { + client + .simple_query( + "PREPARE having_prepared AS SELECT customer_id FROM having_merge_test \ + GROUP BY customer_id HAVING COUNT(*) > $1", + ) + .await?; + } + let error = client.simple_query(sql).await.unwrap_err(); + let message = error.as_db_error().map(|error| error.message()); + assert!( + message.is_some_and(|message| message.contains("cross-shard HAVING is unsupported")), + "{sql}: {error}" + ); + } + client.simple_query("SELECT 1").await?; + cleanup(&sharded).await; Ok(()) } diff --git a/pgdog/src/frontend/router/parser/query/prepare.rs b/pgdog/src/frontend/router/parser/query/prepare.rs index 315dd5172..8297d3876 100644 --- a/pgdog/src/frontend/router/parser/query/prepare.rs +++ b/pgdog/src/frontend/router/parser/query/prepare.rs @@ -1,7 +1,10 @@ use pg_raw_parse::nodes::{ExecuteStmt, PrepareStmt}; use crate::{ - frontend::{BufferedQuery, PreparedStatements}, + frontend::{ + BufferedQuery, PreparedStatements, + router::parser::rewrite::statement::Error as RewriteError, + }, net::{PREPARE_TEMPLATE_NAME, Query}, }; @@ -68,7 +71,18 @@ impl QueryParser { }; match stmt { - Some(Node::SelectStmt(stmt)) => self.select(stmt, &mut context), + Some(Node::SelectStmt(stmt)) => { + let command = self.select(stmt, &mut context)?; + // The cross-shard HAVING rewrite never sees EXECUTE, so each + // shard would filter its partial groups before the merge. + if !matches!(stmt.having_clause(), Node::None) && command.route().is_cross_shard() { + return Err(RewriteError::UnsupportedHaving( + "PREPARE and EXECUTE statements are unsupported", + ) + .into()); + } + Ok(command) + } Some(Node::InsertStmt(stmt)) => self.insert(stmt.into(), &mut context), Some(Node::UpdateStmt(stmt)) => self.update(stmt.into(), &mut context), Some(Node::DeleteStmt(stmt)) => self.delete(stmt.into(), &mut context), diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/having.rs b/pgdog/src/frontend/router/parser/rewrite/statement/having.rs index 4ac3939d5..25cc50c13 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/having.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/having.rs @@ -50,7 +50,11 @@ pub(super) fn rewrite_select<'a>( match group { Node::ColumnRef(_) => { let alias = parser.next_alias(); - parser.values.push((alias, group)); + parser.values.push(Projected { + alias, + node: group, + typed_like: None, + }); } Node::A_Const(_) => {} _ => { @@ -64,12 +68,18 @@ pub(super) fn rewrite_select<'a>( let hidden_aliases = parser .values .iter() - .map(|(alias, _)| alias.clone()) + .map(|value| value.alias.clone()) .collect::>(); let targets = parser .values .iter() - .map(|(alias, node)| mem.make_res_target(Some(alias), mem.empty(), mem.make_unique(*node))) + .map(|value| { + let projected = match value.typed_like { + Some(like) => typed_like(mem, value.node, like), + None => mem.make_unique(value.node), + }; + mem.make_res_target(Some(&value.alias), mem.empty(), projected) + }) .collect::>(); select .target_list_mut() @@ -89,9 +99,34 @@ pub(super) fn rewrite_select<'a>( Ok(HavingRewritePlan::new(expression, hidden_aliases)) } +/// `CASE WHEN false THEN ELSE END` evaluates to `value` with the +/// type PostgreSQL infers for it in a comparison with `like`. A bare untyped +/// literal or parameter in the target list would be resolved as `text` instead. +fn typed_like<'a>( + mem: make::MemoryToken<'a>, + value: Node<'_>, + like: Node<'_>, +) -> make::Unique<'a, Node<'a>> { + let mut when = mem.make_node::(); + when.as_mut() + .set_expr(mem.make_a_const(ConstValue::Boolean(false)).uncast()); + when.as_mut().set_result(mem.make_unique(like)); + let mut case = mem.make_node::(); + case.as_mut().set_args(mem.make_list(&[when.uncast()])); + case.as_mut().set_defresult(mem.make_unique(value)); + case.uncast() +} + +struct Projected<'a> { + alias: String, + node: Node<'a>, + /// Operand whose type an untyped literal or parameter should take. + typed_like: Option>, +} + struct Parser<'a> { aliases: HashSet, - values: Vec<(String, Node<'a>)>, + values: Vec>, next_alias: usize, } @@ -111,10 +146,11 @@ impl<'a> Parser<'a> { "only =, <>, !=, <, <=, > and >= comparisons are supported", ) })?; + let (left, right) = (expression.lexpr(), expression.rexpr()); Ok(HavingExpr::Compare { - left: self.parse_value(expression.lexpr())?, + left: self.parse_value(left, right)?, operator, - right: self.parse_value(expression.rexpr())?, + right: self.parse_value(right, left)?, }) } _ => Err(Error::UnsupportedHaving( @@ -123,15 +159,20 @@ impl<'a> Parser<'a> { } } - fn parse_value(&mut self, node: Node<'a>) -> Result { + fn parse_value(&mut self, node: Node<'a>, other: Node<'a>) -> Result { if !Self::supported_value(node) { return Err(Error::UnsupportedHaving( "operands must be aggregate calls, grouped columns, constants or parameters", )); } + let untyped = |node| matches!(node, Node::A_Const(_) | Node::ParamRef(_)); let alias = self.next_alias(); - self.values.push((alias.clone(), node)); + self.values.push(Projected { + alias: alias.clone(), + node, + typed_like: (untyped(node) && !untyped(other)).then_some(other), + }); Ok(HavingValue::new(alias)) } @@ -213,7 +254,8 @@ mod tests { assert_eq!( sql, "SELECT customer_id, count(*), count(*) AS __pgdog_having_0, \ - $1 AS __pgdog_having_1, customer_id AS __pgdog_having_2 \ + CASE WHEN false THEN count(*) ELSE $1 END AS __pgdog_having_1, \ + customer_id AS __pgdog_having_2 \ FROM purchases GROUP BY customer_id HAVING count(*) > $1 OR true" ); assert_eq!(plan.hidden_aliases().len(), 3); diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs b/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs index 1a0288054..5b3d99385 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs @@ -237,6 +237,33 @@ pub(crate) fn finalize_after_route( Ok(()) } +/// Query shapes that would return wrong rows once HAVING moves after the merge. +fn unsupported_with_having(select: &nodes::SelectStmt) -> Option<&'static str> { + // Aggregates are merged by SELECT list position, which star expansion shifts. + let expands_star = select + .target_list() + .iter() + .any(|target| match target.val() { + Node::ColumnRef(column) => is_star(column), + Node::A_Indirection(indirection) => matches!( + indirection.indirection().into_iter().next_back(), + Some(Node::A_Star(_)) + ), + _ => false, + }); + if expands_star { + return Some("SELECT * is unsupported"); + } + // Shards would drop groups before PgDog merges and filters them. + if !matches!(select.limit_count(), Node::None) || !matches!(select.limit_offset(), Node::None) { + return Some("LIMIT and OFFSET are unsupported"); + } + if !matches!(select.distinct_clause().first(), None | Some(Node::None)) { + return Some("DISTINCT ON is unsupported"); + } + None +} + fn build( ast: &StmtList, schema: &Schema, @@ -246,6 +273,9 @@ fn build( return Ok(None); }; let has_having = !matches!(select.having_clause(), Node::None); + if has_having && let Some(reason) = unsupported_with_having(select) { + return Err(Error::UnsupportedHaving(reason)); + } let aggregate = if has_having { None } else { @@ -483,6 +513,14 @@ mod tests { GROUP BY customer_id HAVING customer_id > 2", "SELECT lower(customer_id), COUNT(*) FROM purchases \ GROUP BY lower(customer_id) HAVING COUNT(*) > 2", + "SELECT * FROM purchases GROUP BY customer_id HAVING COUNT(*) > 2", + "SELECT p.*, COUNT(*) FROM purchases p GROUP BY customer_id HAVING COUNT(*) > 2", + "SELECT customer_id FROM purchases GROUP BY customer_id \ + HAVING COUNT(*) > 2 LIMIT 1", + "SELECT customer_id FROM purchases GROUP BY customer_id \ + HAVING COUNT(*) > 2 OFFSET 1", + "SELECT DISTINCT ON (customer_id) customer_id FROM purchases \ + GROUP BY customer_id HAVING COUNT(*) > 2", ] { let ast = pg_raw_parse::parse(sql).unwrap(); let error = build(&ast, &Schema::default(), false).unwrap_err(); From a06c0ee5aaf1e2dcccc7b42dcf69428e7f8c24b7 Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Sun, 1 Nov 2026 02:20:05 +0530 Subject: [PATCH 22/24] more tests --- .../src/frontend/router/parser/rewrite/statement/having.rs | 2 ++ .../frontend/router/parser/rewrite/statement/projection.rs | 6 ++++-- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/having.rs b/pgdog/src/frontend/router/parser/rewrite/statement/having.rs index 25cc50c13..b1dca17e5 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/having.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/having.rs @@ -270,6 +270,8 @@ mod tests { HAVING MAX(price) IS NULL", "SELECT customer_id FROM purchases GROUP BY customer_id \ HAVING COUNT(*)::numeric > 3", + "SELECT customer_id FROM purchases GROUP BY customer_id \ + HAVING MAX(name) ~ 'a'", ] { let ast = pg_raw_parse::parse(sql).unwrap(); let Node::SelectStmt(select) = ast.stmts().next().unwrap() else { diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs b/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs index 5b3d99385..7ee4b60f9 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs @@ -490,8 +490,8 @@ mod tests { #[test] fn aggregate_used_only_by_having_is_merged_and_hidden() { let ast = pg_raw_parse::parse( - "SELECT customer_id FROM purchases \ - GROUP BY customer_id HAVING COUNT(*) > 2", + "SELECT customer_id, 'total' FROM purchases \ + GROUP BY 1 HAVING COUNT(*) > 2", ) .unwrap(); let rewrite = build(&ast, &Schema::default(), false) @@ -515,6 +515,8 @@ mod tests { GROUP BY lower(customer_id) HAVING COUNT(*) > 2", "SELECT * FROM purchases GROUP BY customer_id HAVING COUNT(*) > 2", "SELECT p.*, COUNT(*) FROM purchases p GROUP BY customer_id HAVING COUNT(*) > 2", + "SELECT (p).*, COUNT(*) FROM purchases p GROUP BY customer_id HAVING COUNT(*) > 2", + "SELECT customer_id, COUNT(*) FROM purchases GROUP BY 'x' HAVING COUNT(*) > 2", "SELECT customer_id FROM purchases GROUP BY customer_id \ HAVING COUNT(*) > 2 LIMIT 1", "SELECT customer_id FROM purchases GROUP BY customer_id \ From 0b290882e26f156c72083c5089fd0635c5146a77 Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Sun, 1 Nov 2026 03:19:20 +0530 Subject: [PATCH 23/24] more cases --- integration/rust/tests/integration/having.rs | 3 +++ .../router/parser/rewrite/statement/having.rs | 23 ++++++++++++++++++- .../parser/rewrite/statement/projection.rs | 20 +++++++++++++--- 3 files changed, 42 insertions(+), 4 deletions(-) diff --git a/integration/rust/tests/integration/having.rs b/integration/rust/tests/integration/having.rs index a34988724..ad9683e27 100644 --- a/integration/rust/tests/integration/having.rs +++ b/integration/rust/tests/integration/having.rs @@ -112,6 +112,9 @@ async fn having_rejects_shard_local_steps() -> Result<(), Box 0 OFFSET 1", "SELECT DISTINCT ON (customer_id) customer_id FROM having_merge_test \ GROUP BY customer_id HAVING COUNT(*) > 0", + "SELECT customer_id AS c FROM having_merge_test GROUP BY c HAVING COUNT(*) > 0", + "SELECT customer_id, COUNT(*) + 1 FROM having_merge_test \ + GROUP BY customer_id HAVING COUNT(*) > 0", "EXECUTE having_prepared(0)", ] { if sql.starts_with("EXECUTE") { diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/having.rs b/pgdog/src/frontend/router/parser/rewrite/statement/having.rs index b1dca17e5..1f710a44e 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/having.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/having.rs @@ -9,7 +9,10 @@ use crate::frontend::router::parser::{ AggregateFunction, CompareOperator, Function, HavingExpr, HavingRewritePlan, HavingValue, }; -use super::{Error, projection::target_output_name}; +use super::{ + Error, + projection::{column_name, target_output_name}, +}; const ALIAS_PREFIX: &str = "__pgdog_having_"; @@ -48,6 +51,12 @@ pub(super) fn rewrite_select<'a>( for group_source in &group_sources { let group = group_source.as_ref(); match group { + // Shards can't resolve a SELECT alias inside the hidden projection. + Node::ColumnRef(column) if groups_by_output_alias(select, column) => { + return Err(Error::UnsupportedHaving( + "GROUP BY on a SELECT alias is unsupported", + )); + } Node::ColumnRef(_) => { let alias = parser.next_alias(); parser.values.push(Projected { @@ -99,6 +108,18 @@ pub(super) fn rewrite_select<'a>( Ok(HavingRewritePlan::new(expression, hidden_aliases)) } +fn groups_by_output_alias( + select: &nodes::SelectStmtMut<'_, '_>, + column: &nodes::ColumnRef, +) -> bool { + column.fields().len() == 1 + && select.target_list().iter().any(|target| { + target + .name() + .is_some_and(|name| Some(name) == column_name(column)) + }) +} + /// `CASE WHEN false THEN ELSE END` evaluates to `value` with the /// type PostgreSQL infers for it in a comparison with `like`. A bare untyped /// literal or parameter in the target list would be resolved as `text` instead. diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs b/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs index 7ee4b60f9..d1c6c0339 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs @@ -10,7 +10,9 @@ use super::having; use super::offset::{self, OffsetPlan}; use super::order_by; use crate::backend::schema::Schema; -use crate::frontend::router::parser::{Aggregate, HavingExpr, HavingRewritePlan, OrderBy}; +use crate::frontend::router::parser::{ + Aggregate, AggregateFunction, Function, HavingExpr, HavingRewritePlan, OrderBy, +}; use crate::frontend::{ClientRequest, PreparedStatements}; use crate::net::{ProtocolMessage, RowDescription}; use pg_raw_parse::{Node, StmtList, make, nodes}; @@ -238,7 +240,7 @@ pub(crate) fn finalize_after_route( } /// Query shapes that would return wrong rows once HAVING moves after the merge. -fn unsupported_with_having(select: &nodes::SelectStmt) -> Option<&'static str> { +fn unsupported_with_having(select: &nodes::SelectStmt, schema: &Schema) -> Option<&'static str> { // Aggregates are merged by SELECT list position, which star expansion shifts. let expands_star = select .target_list() @@ -254,6 +256,15 @@ fn unsupported_with_having(select: &nodes::SelectStmt) -> Option<&'static str> { if expands_star { return Some("SELECT * is unsupported"); } + // Only whole aggregate calls are merged; an expression over them keeps one shard's value. + let wraps_aggregate = select.target_list().iter().any(|target| { + order_by::contains_aggregate(target.val(), schema) + && !Function::try_from(target.val()) + .is_ok_and(|function| AggregateFunction::from_name(function.name, schema).is_some()) + }); + if wraps_aggregate { + return Some("expressions over aggregates in the SELECT list are unsupported"); + } // Shards would drop groups before PgDog merges and filters them. if !matches!(select.limit_count(), Node::None) || !matches!(select.limit_offset(), Node::None) { return Some("LIMIT and OFFSET are unsupported"); @@ -273,7 +284,7 @@ fn build( return Ok(None); }; let has_having = !matches!(select.having_clause(), Node::None); - if has_having && let Some(reason) = unsupported_with_having(select) { + if has_having && let Some(reason) = unsupported_with_having(select, schema) { return Err(Error::UnsupportedHaving(reason)); } let aggregate = if has_having { @@ -517,6 +528,9 @@ mod tests { "SELECT p.*, COUNT(*) FROM purchases p GROUP BY customer_id HAVING COUNT(*) > 2", "SELECT (p).*, COUNT(*) FROM purchases p GROUP BY customer_id HAVING COUNT(*) > 2", "SELECT customer_id, COUNT(*) FROM purchases GROUP BY 'x' HAVING COUNT(*) > 2", + "SELECT name AS n, COUNT(*) FROM purchases GROUP BY n HAVING COUNT(*) > 2", + "SELECT customer_id, SUM(price) / COUNT(*) FROM purchases \ + GROUP BY customer_id HAVING COUNT(*) > 2", "SELECT customer_id FROM purchases GROUP BY customer_id \ HAVING COUNT(*) > 2 LIMIT 1", "SELECT customer_id FROM purchases GROUP BY customer_id \ From 70b0f5e9de5640993f14293afccf5ac06982a772 Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Sun, 1 Nov 2026 13:20:35 +0530 Subject: [PATCH 24/24] type fixes --- .../backend/pool/connection/buffer/having.rs | 45 +++++++++++++++---- 1 file changed, 37 insertions(+), 8 deletions(-) diff --git a/pgdog/src/backend/pool/connection/buffer/having.rs b/pgdog/src/backend/pool/connection/buffer/having.rs index f6d8801df..c0ebb1077 100644 --- a/pgdog/src/backend/pool/connection/buffer/having.rs +++ b/pgdog/src/backend/pool/connection/buffer/having.rs @@ -89,6 +89,21 @@ fn compare(left: &Datum, operator: CompareOperator, right: &Datum) -> Result are supported for text".into(), + )); + } let ordering = left .compare_numeric(right) @@ -195,10 +210,10 @@ mod tests { true, ), ( - CompareOperator::Greater, + CompareOperator::NotEqual, Datum::Text("a".into()), Datum::Text("z".into()), - false, + true, ), ] { assert_eq!( @@ -211,13 +226,27 @@ mod tests { ); } - assert!(matches!( - compare( - &Datum::Text("2".into()), + for (left, operator, right) in [ + ( + Datum::Text("2".into()), + CompareOperator::Equal, + Datum::Integer(2), + ), + ( + Datum::Text("a".into()), + CompareOperator::Greater, + Datum::Text("B".into()), + ), + ( + Datum::Unknown("2026-01-02".into()), CompareOperator::Equal, - &Datum::Integer(2) + Datum::Unknown("2026-01-02".into()), ), - Err(Error::UnsupportedHaving(_)) - )); + ] { + assert!(matches!( + compare(&left, operator, &right), + Err(Error::UnsupportedHaving(_)) + )); + } } }