From 4b57cb9d1501806da408e320a14c403e50cc66d2 Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Thu, 9 Jul 2026 00:26:57 +0530 Subject: [PATCH 1/2] feat: add support for having clause --- integration/rust/tests/integration/having.rs | 91 ++++ integration/rust/tests/integration/mod.rs | 1 + pgdog/src/backend/pool/connection/buffer.rs | 208 ++++++++- .../pool/connection/multi_shard/mod.rs | 26 +- pgdog/src/frontend/router/parser/having.rs | 406 ++++++++++++++++++ pgdog/src/frontend/router/parser/mod.rs | 4 + .../router/parser/rewrite/statement/having.rs | 315 ++++++++++++++ .../router/parser/rewrite/statement/mod.rs | 1 + .../router/parser/rewrite/statement/plan.rs | 8 +- pgdog/src/frontend/router/parser/route.rs | 28 +- 10 files changed, 1059 insertions(+), 29 deletions(-) create mode 100644 integration/rust/tests/integration/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..ea52925e6 --- /dev/null +++ b/integration/rust/tests/integration/having.rs @@ -0,0 +1,91 @@ +use crate::setup::{admin_sqlx, connections_sqlx}; +use sqlx::{Executor, Row, postgres::PgPool}; + +#[tokio::test] +async fn having_count_merges_across_shards() -> Result<(), Box> { + let sharded = connections_sqlx().await.get(1).cloned().unwrap(); + + reset_table( + &sharded, + "having_merge_test", + "(customer_id BIGINT, order_id BIGINT)", + ) + .await?; + + seed( + &sharded, + 0, + "INSERT INTO having_merge_test(customer_id, order_id) VALUES (1, 10), (1, 11), (2, 12)", + ) + .await?; + seed( + &sharded, + 1, + "INSERT INTO having_merge_test(customer_id, order_id) 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, + "global HAVING should keep both merged groups" + ); + 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, "having_merge_test").await; + Ok(()) +} + +async fn reset_table(pool: &PgPool, table: &str, schema: &str) -> Result<(), sqlx::Error> { + for shard in [0, 1] { + pool.execute( + format!( + "/* pgdog_shard: {} */ DROP TABLE IF EXISTS {}", + shard, table + ) + .as_str(), + ) + .await + .ok(); + } + + for shard in [0, 1] { + pool.execute( + format!( + "/* pgdog_shard: {} */ CREATE TABLE {} {}", + shard, table, schema + ) + .as_str(), + ) + .await?; + } + admin_sqlx().await.execute("RELOAD").await?; + Ok(()) +} + +async fn seed(pool: &PgPool, shard: usize, stmt: &str) -> Result<(), sqlx::Error> { + pool.execute(format!("/* pgdog_shard: {} */ {}", shard, stmt).as_str()) + .await?; + Ok(()) +} + +async fn cleanup(pool: &PgPool, table: &str) { + for shard in [0, 1] { + pool.execute(format!("/* pgdog_shard: {} */ DROP TABLE {}", shard, table).as_str()) + .await + .ok(); + } +} diff --git a/integration/rust/tests/integration/mod.rs b/integration/rust/tests/integration/mod.rs index 60132d930..5bf39295a 100644 --- a/integration/rust/tests/integration/mod.rs +++ b/integration/rust/tests/integration/mod.rs @@ -12,6 +12,7 @@ pub mod cross_shard_disabled; 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/pool/connection/buffer.rs b/pgdog/src/backend/pool/connection/buffer.rs index edea8fd85..cf14a4193 100644 --- a/pgdog/src/backend/pool/connection/buffer.rs +++ b/pgdog/src/backend/pool/connection/buffer.rs @@ -2,12 +2,13 @@ use std::{ cmp::Ordering, - collections::{HashSet, VecDeque}, + collections::{BTreeSet, HashSet, VecDeque}, }; use crate::{ frontend::router::parser::{ - Aggregate, DistinctBy, DistinctColumn, Limit, OrderBy, + Aggregate, CompareOp, DistinctBy, DistinctColumn, HavingExpr, HavingLiteral, + HavingRewritePlan, HavingValue, Limit, OrderBy, rewrite::statement::aggregate::AggregateRewritePlan, }, net::{ @@ -146,32 +147,18 @@ impl Buffer { plan: &AggregateRewritePlan, ) -> Result<(), super::Error> { let buffer: VecDeque = std::mem::take(&mut self.buffer); - let mut rows = if aggregate.is_empty() { + let rows = if aggregate.is_empty() { buffer } else if let Some(aggregates) = Aggregates::new(&buffer, decoder, aggregate, plan) { aggregates.aggregate()? } else { buffer }; - - Self::drop_helper_columns(&mut rows, plan); self.buffer = rows; Ok(()) } - fn drop_helper_columns(rows: &mut VecDeque, plan: &AggregateRewritePlan) { - if plan.is_noop() { - return; - } - - let drop = plan.drop_columns().collect(); - - for row in rows.iter_mut() { - row.drop_columns(&drop); - } - } - pub(super) fn distinct(&mut self, distinct: &Option, decoder: &Decoder) { if let Some(distinct) = distinct { match distinct { @@ -207,6 +194,48 @@ impl Buffer { } } + pub(super) fn having( + &mut self, + having: &Option, + decoder: &Decoder, + ) -> Result<(), super::Error> { + let Some(having) = having else { + return Ok(()); + }; + + let mut filtered = VecDeque::with_capacity(self.buffer.len()); + for row in self.buffer.drain(..) { + if eval_having(having, &row, decoder)? { + filtered.push_back(row); + } + } + self.buffer = filtered; + Ok(()) + } + + pub(super) fn drop_hidden_columns( + &mut self, + aggregate: &AggregateRewritePlan, + having: &HavingRewritePlan, + decoder: &Decoder, + ) { + let mut drop = aggregate.drop_columns().collect::>(); + for alias in having.hidden_aliases() { + if let Some(index) = decoder.rd().field_index(alias) { + drop.push(index); + } + } + drop.sort_unstable(); + drop.dedup(); + if drop.is_empty() { + return; + } + let drop: BTreeSet<_> = drop.into_iter().collect(); + for row in &mut self.buffer { + row.drop_columns(&drop); + } + } + /// Take messages from buffer. pub(super) fn take(&mut self) -> Option { if self.full { @@ -236,9 +265,127 @@ impl Buffer { } } +fn eval_having( + having: &HavingExpr, + row: &DataRow, + decoder: &Decoder, +) -> Result { + match having { + HavingExpr::Compare { left, op, right } => { + let left = eval_value(left, row, decoder)?; + let right = eval_value(right, row, decoder)?; + compare_values(left.as_ref(), op, right.as_ref()) + } + HavingExpr::And(children) => { + for child in children { + if !eval_having(child, row, decoder)? { + return Ok(false); + } + } + Ok(true) + } + HavingExpr::Or(children) => { + for child in children { + if eval_having(child, row, decoder)? { + return Ok(true); + } + } + Ok(false) + } + HavingExpr::IsNull { value, not } => { + let value = eval_value(value, row, decoder)?; + let result = value.as_ref().map(|datum| datum.is_null()).unwrap_or(true); + Ok(if *not { !result } else { result }) + } + } +} + +fn eval_value( + value: &HavingValue, + row: &DataRow, + decoder: &Decoder, +) -> Result, super::Error> { + match value { + HavingValue::ColumnIndex(index) => Ok(row + .get_column(*index, decoder)? + .map(|column| column.value) + .or(Some(Datum::Null))), + HavingValue::ColumnName(name) => Ok(decoder + .rd() + .field_index(name) + .and_then(|index| row.get_column(index, decoder).ok().flatten()) + .map(|column| column.value) + .or(Some(Datum::Null))), + HavingValue::Literal(literal) => Ok(Some(match literal { + HavingLiteral::Integer(value) => Datum::Bigint(*value), + HavingLiteral::Float(value) => Datum::Double((*value).into()), + HavingLiteral::Boolean(value) => Datum::Boolean(*value), + HavingLiteral::String(value) => Datum::Text(value.clone()), + HavingLiteral::Null => Datum::Null, + })), + } +} + +fn compare_values( + left: Option<&Datum>, + op: &CompareOp, + right: Option<&Datum>, +) -> Result { + let (Some(left), Some(right)) = (left, right) else { + return Ok(false); + }; + + if left.is_null() || right.is_null() { + return Ok(false); + } + + if let Some(ordering) = compare_numeric(left, right) { + return Ok(match op { + CompareOp::Eq => ordering == Ordering::Equal, + CompareOp::NotEq => ordering != Ordering::Equal, + CompareOp::Lt => ordering == Ordering::Less, + CompareOp::Lte => ordering != Ordering::Greater, + CompareOp::Gt => ordering == Ordering::Greater, + CompareOp::Gte => ordering != Ordering::Less, + }); + } + + let Some(ordering) = left.partial_cmp(right) else { + return Ok(false); + }; + Ok(match op { + CompareOp::Eq => ordering == Ordering::Equal, + CompareOp::NotEq => ordering != Ordering::Equal, + CompareOp::Lt => ordering == Ordering::Less, + CompareOp::Lte => ordering != Ordering::Greater, + CompareOp::Gt => ordering == Ordering::Greater, + CompareOp::Gte => ordering != Ordering::Less, + }) +} + +fn compare_numeric(left: &Datum, right: &Datum) -> Option { + let left = datum_to_f64(left)?; + let right = datum_to_f64(right)?; + left.partial_cmp(&right) +} + +fn datum_to_f64(value: &Datum) -> Option { + match value { + Datum::Bigint(value) => Some(*value as f64), + Datum::Integer(value) => Some(*value as f64), + Datum::SmallInt(value) => Some(*value as f64), + Datum::Float(value) => Some(value.0 as f64), + Datum::Double(value) => Some(value.0), + Datum::Numeric(value) => value.to_f64(), + Datum::Text(value) => value.parse::().ok(), + _ => None, + } +} + #[cfg(test)] mod test { use super::*; + use crate::frontend::router::parser::{CompareOp, HavingExpr, HavingLiteral, HavingValue}; use crate::net::{Datum, Field, Format, RowDescription}; use bytes::Bytes; @@ -627,4 +774,31 @@ mod test { assert_eq!(buf.buffer.len(), 3); } + + #[test] + fn test_having_filters_after_aggregate() { + let mut buf = Buffer::default(); + let rd = RowDescription::new(&[Field::bigint("customer_id"), Field::bigint("order_count")]); + let decoder = Decoder::from(&rd); + + let mut row1 = DataRow::new(); + row1.add(1_i64).add(3_i64); + buf.add(row1.message().unwrap()).unwrap(); + + let mut row2 = DataRow::new(); + row2.add(2_i64).add(2_i64); + buf.add(row2.message().unwrap()).unwrap(); + + let having = HavingExpr::Compare { + left: HavingValue::ColumnName("order_count".into()), + op: CompareOp::Gt, + right: HavingValue::Literal(HavingLiteral::Integer(2)), + }; + buf.having(&Some(having), &decoder).unwrap(); + + assert_eq!(buf.len(), 1); + let remaining = buf.buffer.front().unwrap(); + assert_eq!(remaining.get::(0, Format::Text).unwrap(), 1); + assert_eq!(remaining.get::(1, Format::Text).unwrap(), 3); + } } diff --git a/pgdog/src/backend/pool/connection/multi_shard/mod.rs b/pgdog/src/backend/pool/connection/multi_shard/mod.rs index 8c00ee2a2..675465032 100644 --- a/pgdog/src/backend/pool/connection/multi_shard/mod.rs +++ b/pgdog/src/backend/pool/connection/multi_shard/mod.rs @@ -164,17 +164,20 @@ impl MultiShard { self.buffer.full(); if !self.buffer.is_empty() { + let aggregate_plan = self.route.aggregate_rewrite_plan(); + let having_plan = self.route.having_rewrite_plan(); self.buffer - .aggregate( - self.route.aggregate(), - &self.decoder, - self.route.aggregate_rewrite_plan(), - ) + .aggregate(self.route.aggregate(), &self.decoder, aggregate_plan) + .map_err(Error::from)?; + self.buffer + .having(self.route.having(), &self.decoder) .map_err(Error::from)?; self.buffer.sort(self.route.order_by(), &self.decoder); self.buffer.distinct(self.route.distinct(), &self.decoder); self.buffer.limit(self.route.limit()); + self.buffer + .drop_hidden_columns(aggregate_plan, having_plan, &self.decoder); } if has_rows { @@ -207,10 +210,19 @@ impl MultiShard { // Only send it to the client once all shards sent it, // so we don't get early requests from clients. let plan = self.route.aggregate_rewrite_plan(); - if plan.is_noop() { + let having_plan = self.route.having_rewrite_plan(); + if plan.is_noop() && having_plan.is_noop() { forward = Some(message); } else { - let client_rd = rd.drop_columns(plan.drop_columns()); + let mut drop = plan.drop_columns().collect::>(); + for alias in having_plan.hidden_aliases() { + if let Some(index) = rd.field_index(alias) { + drop.push(index); + } + } + drop.sort_unstable(); + drop.dedup(); + let client_rd = rd.drop_columns(drop.into_iter()); forward = Some(client_rd.message()?); } } diff --git a/pgdog/src/frontend/router/parser/having.rs b/pgdog/src/frontend/router/parser/having.rs new file mode 100644 index 000000000..84aab6007 --- /dev/null +++ b/pgdog/src/frontend/router/parser/having.rs @@ -0,0 +1,406 @@ +//! HAVING clause parser for cross-shard post-processing. + +use pg_query::{ + NodeEnum, + protobuf::{AExprKind, BoolExprType, NullTestType, SelectStmt}, +}; + +use super::{Aggregate, AggregateFunction, Function, Value}; + +#[derive(Debug, Clone, PartialEq)] +pub(crate) enum CompareOp { + Eq, + NotEq, + Lt, + Lte, + Gt, + Gte, +} + +#[derive(Debug, Clone, PartialEq)] +pub(crate) enum HavingLiteral { + Integer(i64), + Float(f64), + Boolean(bool), + String(String), + Null, +} + +impl From> for HavingLiteral { + fn from(value: Value<'_>) -> Self { + match value { + Value::String(s) => Self::String(s.to_owned()), + Value::Integer(i) => Self::Integer(i), + Value::Float(f) => Self::Float(f), + Value::Boolean(b) => Self::Boolean(b), + Value::Null => Self::Null, + Value::Placeholder(_) => Self::Null, + Value::Vector(_) => Self::Null, + } + } +} + +#[derive(Debug, Clone, PartialEq)] +pub(crate) enum HavingValueTemplate { + ColumnIndex(usize), + ColumnName(String), + Param(usize), + Literal(HavingLiteral), + Expression(pg_query::protobuf::Node), +} + +#[derive(Debug, Clone, PartialEq)] +pub(crate) enum HavingValue { + ColumnIndex(usize), + ColumnName(String), + Literal(HavingLiteral), +} + +#[derive(Debug, Clone, PartialEq)] +pub(crate) enum HavingExprTemplate { + Compare { + left: HavingValueTemplate, + op: CompareOp, + right: HavingValueTemplate, + }, + And(Vec), + Or(Vec), + IsNull { + value: HavingValueTemplate, + not: bool, + }, +} + +#[derive(Debug, Clone, PartialEq)] +pub(crate) enum HavingExpr { + Compare { + left: HavingValue, + op: CompareOp, + right: HavingValue, + }, + And(Vec), + Or(Vec), + IsNull { + value: HavingValue, + not: bool, + }, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum HavingParseError { + Unsupported, +} + +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub(crate) struct HavingRewritePlan { + hidden_aliases: Vec, +} + +impl HavingRewritePlan { + pub(crate) fn is_noop(&self) -> bool { + self.hidden_aliases.is_empty() + } + + pub(crate) fn hidden_aliases(&self) -> &[String] { + &self.hidden_aliases + } + + pub(crate) fn add_hidden_alias(&mut self, alias: String) { + self.hidden_aliases.push(alias); + } +} + +impl HavingExprTemplate { + pub(crate) fn parse( + stmt: &SelectStmt, + aggregate: &Aggregate, + ) -> Result, HavingParseError> { + let Some(having) = &stmt.having_clause else { + return Ok(None); + }; + Self::parse_node(stmt, aggregate, having).map(Some) + } + + pub(crate) fn param_indices(&self) -> Vec { + let mut params = vec![]; + self.collect_params(&mut params); + params.sort_unstable(); + params.dedup(); + params + } + + fn collect_params(&self, params: &mut Vec) { + match self { + Self::Compare { left, right, .. } => { + Self::collect_value_params(left, params); + Self::collect_value_params(right, params); + } + Self::And(children) | Self::Or(children) => { + for child in children { + child.collect_params(params); + } + } + Self::IsNull { value, .. } => Self::collect_value_params(value, params), + } + } + + fn collect_value_params(value: &HavingValueTemplate, params: &mut Vec) { + if let HavingValueTemplate::Param(index) = value { + params.push(*index); + } + } + + pub(crate) fn resolve_with( + &self, + lookup: &impl Fn(usize) -> Option, + ) -> Option { + match self { + HavingExprTemplate::Compare { left, op, right } => Some(HavingExpr::Compare { + left: Self::resolve_value(left, lookup)?, + op: op.clone(), + right: Self::resolve_value(right, lookup)?, + }), + HavingExprTemplate::And(children) => { + let children = children + .iter() + .map(|child| child.resolve_with(lookup)) + .collect::>>()?; + Some(HavingExpr::And(children)) + } + HavingExprTemplate::Or(children) => { + let children = children + .iter() + .map(|child| child.resolve_with(lookup)) + .collect::>>()?; + Some(HavingExpr::Or(children)) + } + HavingExprTemplate::IsNull { value, not } => Some(HavingExpr::IsNull { + value: Self::resolve_value(value, lookup)?, + not: *not, + }), + } + } + + fn resolve_value( + value: &HavingValueTemplate, + lookup: &impl Fn(usize) -> Option, + ) -> Option { + match value { + HavingValueTemplate::ColumnIndex(index) => Some(HavingValue::ColumnIndex(*index)), + HavingValueTemplate::ColumnName(name) => Some(HavingValue::ColumnName(name.clone())), + HavingValueTemplate::Literal(value) => Some(HavingValue::Literal(value.clone())), + HavingValueTemplate::Param(index) => lookup(*index).map(HavingValue::Literal), + HavingValueTemplate::Expression(_) => None, + } + } + + fn parse_node( + stmt: &SelectStmt, + aggregate: &Aggregate, + node: &pg_query::protobuf::Node, + ) -> Result { + match node.node.as_ref() { + Some(NodeEnum::BoolExpr(bool_expr)) => { + let parsed = bool_expr + .args + .iter() + .map(|arg| Self::parse_node(stmt, aggregate, arg)) + .collect::, _>>()?; + match bool_expr.boolop() { + BoolExprType::AndExpr => Ok(Self::And(parsed)), + BoolExprType::OrExpr => Ok(Self::Or(parsed)), + _ => Err(HavingParseError::Unsupported), + } + } + Some(NodeEnum::AExpr(expr)) => { + if expr.kind() != AExprKind::AexprOp { + return Err(HavingParseError::Unsupported); + } + let op = Self::parse_compare_op(expr.name.first())?; + let left = expr + .lexpr + .as_ref() + .ok_or(HavingParseError::Unsupported) + .and_then(|node| Self::parse_value(stmt, aggregate, node))?; + let right = expr + .rexpr + .as_ref() + .ok_or(HavingParseError::Unsupported) + .and_then(|node| Self::parse_value(stmt, aggregate, node))?; + Ok(Self::Compare { left, op, right }) + } + Some(NodeEnum::NullTest(test)) => { + let value = test + .arg + .as_ref() + .ok_or(HavingParseError::Unsupported) + .and_then(|node| Self::parse_value(stmt, aggregate, node))?; + let not = matches!(test.nulltesttype(), NullTestType::IsNotNull); + if matches!( + test.nulltesttype(), + NullTestType::IsNull | NullTestType::IsNotNull + ) { + Ok(Self::IsNull { value, not }) + } else { + Err(HavingParseError::Unsupported) + } + } + _ => Err(HavingParseError::Unsupported), + } + } + + fn parse_value( + stmt: &SelectStmt, + aggregate: &Aggregate, + node: &pg_query::protobuf::Node, + ) -> Result { + match node.node.as_ref() { + Some(NodeEnum::TypeCast(cast)) => cast + .arg + .as_ref() + .ok_or(HavingParseError::Unsupported) + .and_then(|arg| Self::parse_value(stmt, aggregate, arg)), + Some(NodeEnum::AConst(_)) => { + let value = Value::try_from(node).map_err(|_| HavingParseError::Unsupported)?; + Ok(HavingValueTemplate::Literal(value.into())) + } + Some(NodeEnum::ParamRef(param)) => { + if param.number <= 0 { + return Err(HavingParseError::Unsupported); + } + Ok(HavingValueTemplate::Param((param.number - 1) as usize)) + } + Some(NodeEnum::ColumnRef(column_ref)) => { + let fields = column_ref + .fields + .iter() + .filter_map(|field| match field.node.as_ref() { + Some(NodeEnum::String(string)) => Some(string.sval.as_str()), + _ => None, + }) + .collect::>(); + + if let Some(name) = fields.last() { + if let Some(index) = Self::target_index_by_alias(stmt, name) { + return Ok(HavingValueTemplate::ColumnIndex(index)); + } + return Ok(HavingValueTemplate::Expression(node.clone())); + } + Err(HavingParseError::Unsupported) + } + Some(NodeEnum::FuncCall(_)) => { + let func = Function::try_from(node).map_err(|_| HavingParseError::Unsupported)?; + let agg_fn = Self::aggregate_function_from_name(func.name) + .ok_or(HavingParseError::Unsupported)?; + let candidates = aggregate + .targets() + .iter() + .filter(|target| target.function() == &agg_fn) + .map(|target| target.column()) + .collect::>(); + + if candidates.len() == 1 { + Ok(HavingValueTemplate::ColumnIndex(candidates[0])) + } else { + Ok(HavingValueTemplate::Expression(node.clone())) + } + } + Some(NodeEnum::ResTarget(res_target)) => res_target + .val + .as_ref() + .ok_or(HavingParseError::Unsupported) + .and_then(|value| Self::parse_value(stmt, aggregate, value)), + _ => Err(HavingParseError::Unsupported), + } + } + + fn parse_compare_op( + node: Option<&pg_query::protobuf::Node>, + ) -> Result { + let Some(NodeEnum::String(string)) = node.and_then(|node| node.node.as_ref()) else { + return Err(HavingParseError::Unsupported); + }; + match string.sval.as_str() { + "=" => Ok(CompareOp::Eq), + "<>" | "!=" => Ok(CompareOp::NotEq), + "<" => Ok(CompareOp::Lt), + "<=" => Ok(CompareOp::Lte), + ">" => Ok(CompareOp::Gt), + ">=" => Ok(CompareOp::Gte), + _ => Err(HavingParseError::Unsupported), + } + } + + fn target_index_by_alias(stmt: &SelectStmt, alias: &str) -> Option { + stmt.target_list + .iter() + .position(|target| match target.node.as_ref() { + Some(NodeEnum::ResTarget(res)) => res.name == alias, + _ => false, + }) + } + + fn aggregate_function_from_name(name: &str) -> Option { + match name { + "count" => Some(AggregateFunction::Count), + "max" => Some(AggregateFunction::Max), + "min" => Some(AggregateFunction::Min), + "sum" => Some(AggregateFunction::Sum), + "avg" => Some(AggregateFunction::Avg), + "stddev" | "stddev_samp" => Some(AggregateFunction::StddevSamp), + "stddev_pop" => Some(AggregateFunction::StddevPop), + "variance" | "var_samp" => Some(AggregateFunction::VarSamp), + "var_pop" => Some(AggregateFunction::VarPop), + _ => None, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::backend::schema::Schema; + use pg_query::{NodeEnum, parse}; + + fn parse_select(sql: &str) -> SelectStmt { + let stmt = parse(sql).unwrap().protobuf.stmts.remove(0).stmt.unwrap(); + match stmt.node.unwrap() { + NodeEnum::SelectStmt(stmt) => *stmt, + _ => panic!("not a select"), + } + } + + #[test] + fn parses_simple_count_having() { + let stmt = parse_select( + "SELECT user_id, COUNT(*) AS c FROM t GROUP BY user_id HAVING COUNT(*) > 2", + ); + let aggregate = Aggregate::parse(&stmt, &Schema::default()); + let having = HavingExprTemplate::parse(&stmt, &aggregate) + .unwrap() + .expect("having"); + match having { + HavingExprTemplate::Compare { .. } => {} + _ => panic!("expected compare expression"), + } + } + + #[test] + fn resolves_params() { + let stmt = parse_select( + "SELECT user_id, COUNT(*) AS c FROM t GROUP BY user_id HAVING COUNT(*) > $1", + ); + let aggregate = Aggregate::parse(&stmt, &Schema::default()); + let having = HavingExprTemplate::parse(&stmt, &aggregate) + .unwrap() + .expect("having"); + let resolved = having.resolve_with(&|idx| { + if idx == 0 { + Some(HavingLiteral::Integer(2)) + } else { + None + } + }); + assert!(resolved.is_some()); + } +} diff --git a/pgdog/src/frontend/router/parser/mod.rs b/pgdog/src/frontend/router/parser/mod.rs index b6cefd16c..94c8cb6c2 100644 --- a/pgdog/src/frontend/router/parser/mod.rs +++ b/pgdog/src/frontend/router/parser/mod.rs @@ -15,6 +15,7 @@ pub mod error; pub mod explain_trace; mod from_clause; pub mod function; +pub mod having; pub mod key; mod limit; pub mod multi_tenant; @@ -44,6 +45,9 @@ pub use distinct::{Distinct, DistinctBy, DistinctColumn}; pub use error::Error; pub(crate) use from_clause::FromClause; use function::Function; +pub(crate) use having::{ + CompareOp, HavingExpr, HavingExprTemplate, HavingLiteral, HavingRewritePlan, HavingValue, +}; pub use key::Key; pub(crate) use limit::{Limit, LimitClause}; pub use order_by::OrderBy; 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..dc9334328 --- /dev/null +++ b/pgdog/src/frontend/router/parser/rewrite/statement/having.rs @@ -0,0 +1,315 @@ +use pg_query::{ + NodeEnum, + protobuf::{ + AConst, BoolExpr, BoolExprType, Boolean, Node, ParseResult as ProtobufParseResult, + ResTarget, a_const::Val, + }, +}; + +use crate::{ + backend::schema::Schema, + frontend::ClientRequest, + frontend::router::parser::{ + Aggregate, HavingExprTemplate, HavingLiteral, HavingRewritePlan, + having::HavingValueTemplate, + }, + net::{ProtocolMessage, messages::bind::Bind}, +}; + +use super::Error; + +pub(super) fn apply_having_after_parser(request: &mut ClientRequest) -> Result<(), Error> { + let aggregate = match request.route.as_ref() { + Some(route) if route.is_cross_shard() => route.aggregate().clone(), + _ => return Ok(()), + }; + + let ast = match request.ast.as_ref() { + Some(ast) => ast, + None => return Ok(()), + }; + + let mut source_protobuf = ast.ast.protobuf.clone(); + let Some(select) = select_stmt_mut(&mut source_protobuf) else { + return Ok(()); + }; + + let mut template = match HavingExprTemplate::parse(select, &aggregate) { + Ok(Some(template)) => template, + Ok(None) | Err(_) => return Ok(()), + }; + + let bind = first_bind(request); + if bind.is_none() && !template.param_indices().is_empty() { + return Ok(()); + } + + let Some(current_sql) = current_sql(request) else { + return Ok(()); + }; + + let mut rewritten = pg_query::parse(current_sql)?.protobuf; + let Some(select) = select_stmt_mut(&mut rewritten) else { + return Ok(()); + }; + let having_plan = materialize_hidden_expressions(select, &mut template); + let resolved = template.resolve_with(&|idx| { + let bind = bind?; + decode_bind_param(bind, idx) + }); + let Some(resolved) = resolved else { + return Ok(()); + }; + if !rewrite_having_clause_true(select) { + return Ok(()); + } + + let rewritten_aggregate = if having_plan.is_noop() { + None + } else { + Some(Aggregate::parse(select, &Schema::default())) + }; + let new_sql = pg_query::ParseResult::new(rewritten, String::new()).deparse()?; + for message in request.messages.iter_mut() { + match message { + ProtocolMessage::Query(query) => query.set_query(&new_sql), + ProtocolMessage::Parse(parse) => parse.set_query(&new_sql), + _ => {} + } + } + + if let Some(route) = request.route.as_mut() { + route.set_having(Some(resolved)); + route.set_having_rewrite_plan(having_plan); + if let Some(rewritten_aggregate) = rewritten_aggregate { + *route.aggregate_mut() = rewritten_aggregate; + } + } + Ok(()) +} + +fn first_bind(request: &ClientRequest) -> Option<&Bind> { + request.messages.iter().find_map(|message| match message { + ProtocolMessage::Bind(bind) => Some(bind), + _ => None, + }) +} + +fn current_sql(request: &ClientRequest) -> Option<&str> { + request.messages.iter().find_map(|message| match message { + ProtocolMessage::Query(query) => Some(query.query()), + ProtocolMessage::Parse(parse) => Some(parse.query()), + _ => None, + }) +} + +fn select_stmt_mut(ast: &mut ProtobufParseResult) -> Option<&mut pg_query::protobuf::SelectStmt> { + let stmt = ast.stmts.first_mut()?.stmt.as_mut()?; + match stmt.node.as_mut()? { + NodeEnum::SelectStmt(select) => Some(select.as_mut()), + _ => None, + } +} + +fn decode_bind_param(bind: &Bind, index: usize) -> Option { + let parameter = bind.parameter(index).ok()??; + if parameter.is_null() { + return Some(HavingLiteral::Null); + } + + if let Some(text) = parameter.text() { + if let Ok(boolean) = text.parse::() { + return Some(HavingLiteral::Boolean(boolean)); + } + if let Ok(integer) = text.parse::() { + return Some(HavingLiteral::Integer(integer)); + } + if let Ok(float) = text.parse::() { + return Some(HavingLiteral::Float(float)); + } + return Some(HavingLiteral::String(text.to_owned())); + } + + if let Some(integer) = parameter.decode::() { + return Some(HavingLiteral::Integer(integer)); + } + if let Some(boolean) = parameter.decode::() { + return Some(HavingLiteral::Boolean(boolean)); + } + if let Some(text) = parameter.decode::() { + return Some(HavingLiteral::String(text)); + } + + Some(HavingLiteral::String(parameter.text_debug())) +} + +fn materialize_hidden_expressions( + select: &mut pg_query::protobuf::SelectStmt, + template: &mut HavingExprTemplate, +) -> HavingRewritePlan { + let mut plan = HavingRewritePlan::default(); + let mut next = 0usize; + rewrite_template_values(template, &mut |value| { + let HavingValueTemplate::Expression(expr) = value else { + return; + }; + let alias = format!("__pgdog_having_expr{}", next); + next += 1; + + let target = Node { + node: Some(NodeEnum::ResTarget(Box::new(ResTarget { + name: alias.clone(), + val: Some(Box::new(expr.clone())), + ..Default::default() + }))), + }; + select.target_list.push(target); + plan.add_hidden_alias(alias.clone()); + *value = HavingValueTemplate::ColumnName(alias); + }); + plan +} + +fn rewrite_template_values( + expr: &mut HavingExprTemplate, + callback: &mut impl FnMut(&mut HavingValueTemplate), +) { + match expr { + HavingExprTemplate::Compare { left, right, .. } => { + callback(left); + callback(right); + } + HavingExprTemplate::And(children) | HavingExprTemplate::Or(children) => { + for child in children { + rewrite_template_values(child, callback); + } + } + HavingExprTemplate::IsNull { value, .. } => callback(value), + } +} + +fn rewrite_having_clause_true(select: &mut pg_query::protobuf::SelectStmt) -> bool { + let Some(existing) = select.having_clause.take() else { + return false; + }; + + let true_node = Node { + node: Some(NodeEnum::AConst(AConst { + val: Some(Val::Boolval(Boolean { boolval: true })), + isnull: false, + location: -1, + })), + }; + + select.having_clause = Some(Box::new(Node { + node: Some(NodeEnum::BoolExpr(Box::new(BoolExpr { + boolop: BoolExprType::OrExpr.into(), + args: vec![*existing, true_node], + location: -1, + ..Default::default() + }))), + })); + + true +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::frontend::router::parser::cache::ast::Ast; + use crate::frontend::router::parser::{Aggregate, DistinctBy}; + use crate::frontend::router::parser::{Limit, Route, Shard, ShardWithPriority}; + use crate::net::{Parse, messages::Query}; + + fn cross_shard_route(aggregate: Aggregate) -> Route { + Route::select( + ShardWithPriority::new_table(Shard::All), + vec![], + aggregate, + Limit::default(), + Option::::None, + ) + } + + #[test] + fn apply_having_rewrites_query_and_sets_route_having() { + let sql = "SELECT user_id, COUNT(*) AS c FROM t GROUP BY user_id HAVING COUNT(*) > 2"; + let ast = Ast::new_record(sql, pgdog_config::QueryParserEngine::PgQueryProtobuf).unwrap(); + let aggregate = crate::frontend::router::parser::Aggregate::parse( + match ast + .ast + .protobuf + .stmts + .first() + .and_then(|stmt| stmt.stmt.as_ref()) + .and_then(|stmt| stmt.node.as_ref()) + { + Some(NodeEnum::SelectStmt(select)) => select, + _ => panic!("not a select"), + }, + &Default::default(), + ); + + let mut request = ClientRequest::from(vec![ + ProtocolMessage::Query(Query::new(sql)), + ProtocolMessage::Parse(Parse::named("s", sql)), + ]); + request.ast = Some(ast); + request.route = Some(cross_shard_route(aggregate)); + + apply_having_after_parser(&mut request).unwrap(); + + let query = match &request.messages[0] { + ProtocolMessage::Query(query) => query.query(), + _ => panic!("expected query"), + }; + assert!(query.contains("HAVING")); + assert!(query.contains("true")); + let route = request.route.unwrap(); + assert!(route.having().is_some()); + assert!(route.having_rewrite_plan().is_noop()); + } + + #[test] + fn apply_having_injects_hidden_target_for_missing_aggregate() { + let sql = "SELECT user_id FROM t GROUP BY user_id HAVING COUNT(*) > 1"; + let ast = Ast::new_record(sql, pgdog_config::QueryParserEngine::PgQueryProtobuf).unwrap(); + let aggregate = crate::frontend::router::parser::Aggregate::parse( + match ast + .ast + .protobuf + .stmts + .first() + .and_then(|stmt| stmt.stmt.as_ref()) + .and_then(|stmt| stmt.node.as_ref()) + { + Some(NodeEnum::SelectStmt(select)) => select, + _ => panic!("not a select"), + }, + &Default::default(), + ); + + let mut request = ClientRequest::from(vec![ + ProtocolMessage::Query(Query::new(sql)), + ProtocolMessage::Parse(Parse::named("s", sql)), + ]); + request.ast = Some(ast); + request.route = Some(cross_shard_route(aggregate)); + + apply_having_after_parser(&mut request).unwrap(); + + let query = match &request.messages[0] { + ProtocolMessage::Query(query) => query.query(), + _ => panic!("expected query"), + }; + assert!(query.contains("__pgdog_having_expr0")); + assert!(query.contains("HAVING")); + assert!(query.contains("true")); + + let route = request.route.unwrap(); + assert!(route.having().is_some()); + assert!(!route.having_rewrite_plan().is_noop()); + assert_eq!(route.having_rewrite_plan().hidden_aliases().len(), 1); + assert!(!route.aggregate().targets().is_empty()); + } +} diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs b/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs index 3ef5728a2..2258960b7 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs @@ -16,6 +16,7 @@ use crate::net::parameter::ParameterValue; pub mod aggregate; pub mod auto_id; pub mod error; +mod having; pub mod insert; pub mod offset; pub mod plan; diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs b/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs index 627316696..288e79fa9 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs @@ -3,6 +3,7 @@ use crate::net::messages::bind::{Format, Parameter}; use crate::net::{Bind, Parse, ProtocolMessage, Query}; use crate::unique_id::UniqueId; +use super::having::apply_having_after_parser; use super::insert::build_split_requests; use super::offset::OffsetPlan; use super::{Error, InsertSplit, ShardingKeyUpdate, aggregate::AggregateRewritePlan}; @@ -60,9 +61,10 @@ impl RewriteResult { match self { Self::InPlace { offset: Some(offset), - } => offset.apply_after_parser(request), - _ => Ok(()), - } + } => offset.apply_after_parser(request)?, + _ => {} + }; + apply_having_after_parser(request) } } diff --git a/pgdog/src/frontend/router/parser/route.rs b/pgdog/src/frontend/router/parser/route.rs index 6496f40c2..623422e0f 100644 --- a/pgdog/src/frontend/router/parser/route.rs +++ b/pgdog/src/frontend/router/parser/route.rs @@ -3,8 +3,9 @@ use std::{fmt::Display, ops::Deref}; use lazy_static::lazy_static; use super::{ - Aggregate, DistinctBy, Limit, OrderBy, explain_trace::ExplainTrace, - rewrite::statement::aggregate::AggregateRewritePlan, statement::AdvisoryLocks, + Aggregate, DistinctBy, HavingExpr, HavingRewritePlan, Limit, OrderBy, + explain_trace::ExplainTrace, rewrite::statement::aggregate::AggregateRewritePlan, + statement::AdvisoryLocks, }; /// The shard destination for a query. @@ -105,6 +106,11 @@ pub struct Route { advisory_locks: AdvisoryLocks, /// `DISTINCT` clause, if set. distinct: Option, + /// Parsed `HAVING` clause. + having: Option, + /// Rewrites performed by the HAVING rewriter; adds hidden + /// select targets needed for post-aggregation filtering. + having_rewrite_plan: HavingRewritePlan, /// Rewrites performed by the aggregate rewriter; adds /// helper columns to this query so we can compute things /// like avg() or variance(). @@ -278,11 +284,13 @@ impl Route { /// 2. `GROUP BY` clause /// 3. `DISTINCT` clause /// 4. `LIMIT` or `OFFSET` clause + /// 5. `HAVING` clause /// pub fn should_buffer(&self) -> bool { !self.order_by().is_empty() || !self.aggregate().is_empty() || self.distinct().is_some() + || self.having().is_some() || self.limit().offset.is_some() } @@ -364,6 +372,22 @@ impl Route { &self.distinct } + pub(crate) fn having(&self) -> &Option { + &self.having + } + + pub(crate) fn set_having(&mut self, having: Option) { + self.having = having; + } + + pub(crate) fn having_rewrite_plan(&self) -> &HavingRewritePlan { + &self.having_rewrite_plan + } + + pub(crate) fn set_having_rewrite_plan(&mut self, plan: HavingRewritePlan) { + self.having_rewrite_plan = plan; + } + pub fn should_2pc(&self) -> bool { self.is_cross_shard() && self.is_write() } From eaf00db8d459a992e3a62e83fdde837d221b6d65 Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Thu, 9 Jul 2026 00:40:09 +0530 Subject: [PATCH 2/2] fmt fixes --- .../pool/connection/multi_shard/mod.rs | 2 +- pgdog/src/frontend/router/parser/having.rs | 35 +++++++++++-------- .../router/parser/rewrite/statement/having.rs | 8 ++--- .../router/parser/rewrite/statement/plan.rs | 12 +++---- 4 files changed, 32 insertions(+), 25 deletions(-) diff --git a/pgdog/src/backend/pool/connection/multi_shard/mod.rs b/pgdog/src/backend/pool/connection/multi_shard/mod.rs index 675465032..f9b831b50 100644 --- a/pgdog/src/backend/pool/connection/multi_shard/mod.rs +++ b/pgdog/src/backend/pool/connection/multi_shard/mod.rs @@ -222,7 +222,7 @@ impl MultiShard { } drop.sort_unstable(); drop.dedup(); - let client_rd = rd.drop_columns(drop.into_iter()); + let client_rd = rd.drop_columns(drop); forward = Some(client_rd.message()?); } } diff --git a/pgdog/src/frontend/router/parser/having.rs b/pgdog/src/frontend/router/parser/having.rs index 84aab6007..28956473d 100644 --- a/pgdog/src/frontend/router/parser/having.rs +++ b/pgdog/src/frontend/router/parser/having.rs @@ -46,7 +46,7 @@ pub(crate) enum HavingValueTemplate { ColumnName(String), Param(usize), Literal(HavingLiteral), - Expression(pg_query::protobuf::Node), + Expression(Box), } #[derive(Debug, Clone, PartialEq)] @@ -59,14 +59,14 @@ pub(crate) enum HavingValue { #[derive(Debug, Clone, PartialEq)] pub(crate) enum HavingExprTemplate { Compare { - left: HavingValueTemplate, + left: Box, op: CompareOp, - right: HavingValueTemplate, + right: Box, }, And(Vec), Or(Vec), IsNull { - value: HavingValueTemplate, + value: Box, not: bool, }, } @@ -132,15 +132,15 @@ impl HavingExprTemplate { fn collect_params(&self, params: &mut Vec) { match self { Self::Compare { left, right, .. } => { - Self::collect_value_params(left, params); - Self::collect_value_params(right, params); + Self::collect_value_params(left.as_ref(), params); + Self::collect_value_params(right.as_ref(), params); } Self::And(children) | Self::Or(children) => { for child in children { child.collect_params(params); } } - Self::IsNull { value, .. } => Self::collect_value_params(value, params), + Self::IsNull { value, .. } => Self::collect_value_params(value.as_ref(), params), } } @@ -156,9 +156,9 @@ impl HavingExprTemplate { ) -> Option { match self { HavingExprTemplate::Compare { left, op, right } => Some(HavingExpr::Compare { - left: Self::resolve_value(left, lookup)?, + left: Self::resolve_value(left.as_ref(), lookup)?, op: op.clone(), - right: Self::resolve_value(right, lookup)?, + right: Self::resolve_value(right.as_ref(), lookup)?, }), HavingExprTemplate::And(children) => { let children = children @@ -175,7 +175,7 @@ impl HavingExprTemplate { Some(HavingExpr::Or(children)) } HavingExprTemplate::IsNull { value, not } => Some(HavingExpr::IsNull { - value: Self::resolve_value(value, lookup)?, + value: Self::resolve_value(value.as_ref(), lookup)?, not: *not, }), } @@ -227,7 +227,11 @@ impl HavingExprTemplate { .as_ref() .ok_or(HavingParseError::Unsupported) .and_then(|node| Self::parse_value(stmt, aggregate, node))?; - Ok(Self::Compare { left, op, right }) + Ok(Self::Compare { + left: Box::new(left), + op, + right: Box::new(right), + }) } Some(NodeEnum::NullTest(test)) => { let value = test @@ -240,7 +244,10 @@ impl HavingExprTemplate { test.nulltesttype(), NullTestType::IsNull | NullTestType::IsNotNull ) { - Ok(Self::IsNull { value, not }) + Ok(Self::IsNull { + value: Box::new(value), + not, + }) } else { Err(HavingParseError::Unsupported) } @@ -284,7 +291,7 @@ impl HavingExprTemplate { if let Some(index) = Self::target_index_by_alias(stmt, name) { return Ok(HavingValueTemplate::ColumnIndex(index)); } - return Ok(HavingValueTemplate::Expression(node.clone())); + return Ok(HavingValueTemplate::Expression(Box::new(node.clone()))); } Err(HavingParseError::Unsupported) } @@ -302,7 +309,7 @@ impl HavingExprTemplate { if candidates.len() == 1 { Ok(HavingValueTemplate::ColumnIndex(candidates[0])) } else { - Ok(HavingValueTemplate::Expression(node.clone())) + Ok(HavingValueTemplate::Expression(Box::new(node.clone()))) } } Some(NodeEnum::ResTarget(res_target)) => res_target diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/having.rs b/pgdog/src/frontend/router/parser/rewrite/statement/having.rs index dc9334328..08dd0dfe5 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/having.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/having.rs @@ -159,7 +159,7 @@ fn materialize_hidden_expressions( let target = Node { node: Some(NodeEnum::ResTarget(Box::new(ResTarget { name: alias.clone(), - val: Some(Box::new(expr.clone())), + val: Some(expr.clone()), ..Default::default() }))), }; @@ -176,15 +176,15 @@ fn rewrite_template_values( ) { match expr { HavingExprTemplate::Compare { left, right, .. } => { - callback(left); - callback(right); + callback(left.as_mut()); + callback(right.as_mut()); } HavingExprTemplate::And(children) | HavingExprTemplate::Or(children) => { for child in children { rewrite_template_values(child, callback); } } - HavingExprTemplate::IsNull { value, .. } => callback(value), + HavingExprTemplate::IsNull { value, .. } => callback(value.as_mut()), } } diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs b/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs index 288e79fa9..1c4f61877 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs @@ -58,12 +58,12 @@ pub(crate) enum RewriteResult { impl RewriteResult { pub(crate) fn apply_after_parser(&self, request: &mut ClientRequest) -> Result<(), Error> { - match self { - Self::InPlace { - offset: Some(offset), - } => offset.apply_after_parser(request)?, - _ => {} - }; + if let Self::InPlace { + offset: Some(offset), + } = self + { + offset.apply_after_parser(request)?; + } apply_having_after_parser(request) } }