diff --git a/integration/rust/tests/integration/having.rs b/integration/rust/tests/integration/having.rs index 32fd1e034..faccec655 100644 --- a/integration/rust/tests/integration/having.rs +++ b/integration/rust/tests/integration/having.rs @@ -159,13 +159,64 @@ async fn having_null_semantics_on_merged_groups() -> Result<(), Box 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, 1, 50.00, 1.5), (2, 2, 40.00, 0.5), (3, 3, 1.00, 4.0)", + ) + .await?; + seed( + &sharded, + 1, + "INSERT INTO having_merge_test VALUES (1, 4, 49.995, 1.5), (2, 5, 10.00, 0.5)", + ) + .await?; + + assert_eq!( + matching_customers(&sharded, "SUM(price) > 99.99").await?, + vec![1] + ); + assert_eq!( + matching_customers(&sharded, "SUM(weight) >= 3.0").await?, + vec![1, 3] + ); + assert_eq!( + matching_customers(&sharded, "COUNT(*)::numeric > 1.5").await?, + vec![1, 2] + ); + + cleanup(&sharded).await; + Ok(()) +} + +async fn matching_customers(pool: &PgPool, having: &str) -> Result, sqlx::Error> { + let rows = pool + .fetch_all( + format!( + "SELECT customer_id FROM having_merge_test \ + GROUP BY customer_id HAVING {having} ORDER BY customer_id" + ) + .as_str(), + ) + .await?; + Ok(rows.iter().map(|row| row.get("customer_id")).collect()) +} + 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)" + CREATE TABLE having_merge_test(\ + customer_id BIGINT, order_id BIGINT, price NUMERIC, weight DOUBLE PRECISION)" ) .as_str(), ) diff --git a/pgdog/src/backend/pool/connection/aggregate/count.rs b/pgdog/src/backend/pool/connection/aggregate/count.rs index 313da853a..2bde6f960 100644 --- a/pgdog/src/backend/pool/connection/aggregate/count.rs +++ b/pgdog/src/backend/pool/connection/aggregate/count.rs @@ -1,6 +1,7 @@ use super::TypeError; use crate::net::messages::Datum; -use pgdog_postgres_types::{DataType, Float}; +use pgdog_postgres_types::{DataType, Float, Numeric}; +use rust_decimal::prelude::ToPrimitive; #[derive(Debug)] pub(super) struct Count { @@ -23,9 +24,16 @@ impl Count { return Ok(()); } + let partial = match &value { + Datum::Numeric(numeric) => numeric + .as_decimal() + .and_then(|decimal| decimal.to_i64()) + .ok_or(TypeError::NumericOutOfRange(DataType::Bigint))?, + value => value.as_i64()?, + }; self.total = self .total - .checked_add(value.as_i64()?) + .checked_add(partial) .ok_or(TypeError::NumericOutOfRange(DataType::Bigint))?; self.data_type = value.data_type(); Ok(()) @@ -42,6 +50,7 @@ impl Count { DataType::Bigint => Ok(self.total.into()), DataType::Real => Ok(Datum::Float(Float(self.total as f32))), DataType::DoublePrecision => Ok(Datum::from(self.total as f64)), + DataType::Numeric => Ok(Datum::Numeric(Numeric::from(self.total))), to => Err(TypeError::InvalidCast { from: DataType::Bigint, to, @@ -72,6 +81,21 @@ fn count_with_only_null() { assert_eq!(state.finalize().expect("valid count"), 0i64.into()); } +#[test] +fn count_cast_to_numeric_merges_as_numeric() { + let mut state = Count::new(0); + state + .accumulate(Datum::Numeric(Numeric::from(2_i64))) + .unwrap(); + state + .accumulate(Datum::Numeric(Numeric::from(3_i64))) + .unwrap(); + assert_eq!( + state.finalize().expect("valid count"), + Datum::Numeric(Numeric::from(5_i64)) + ); +} + #[test] fn count_cast_rejects_out_of_range_results() { for value in [Datum::SmallInt(i16::MAX), Datum::Integer(i32::MAX)] { diff --git a/pgdog/src/backend/pool/connection/buffer/having.rs b/pgdog/src/backend/pool/connection/buffer/having.rs index 281ff31f9..7856f3c31 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; +use pgdog_postgres_types::{Datum, Double, Numeric}; use crate::{ frontend::router::parser::{CompareOperator, HavingExpr, HavingValue}, @@ -158,8 +158,19 @@ fn compare(left: &Datum, operator: CompareOperator, right: &Datum) -> Result Option { - let (left, right) = (integer(left)?, integer(right)?); - Some(left.cmp(&right)) + 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 } fn integer(value: &Datum) -> Option { @@ -171,10 +182,33 @@ 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::Numeric(value) => value.to_f64(), + 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::net::{Field, RowDescription}; + use pgdog_postgres_types::Float; fn compare_expression(operator: CompareOperator) -> HavingExpr { HavingExpr::Compare { @@ -342,4 +376,68 @@ mod tests { Err(Error::UnsupportedHaving(_)) )); } + + #[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 compares_numeric_and_float_values() { + for (operator, left, right) in [ + ( + CompareOperator::Greater, + Datum::Numeric(Numeric::from(3_i32)), + Datum::SmallInt(2), + ), + ( + CompareOperator::GreaterOrEqual, + Datum::Float(Float(2.0)), + Datum::Double(Double(2.0)), + ), + ( + CompareOperator::Less, + Datum::Integer(2), + Datum::Double(Double(2.5)), + ), + ( + CompareOperator::GreaterOrEqual, + Datum::Double(Double(3.0)), + Datum::Numeric(Numeric::from(3_i32)), + ), + ] { + assert_eq!(compare(&left, operator, &right).unwrap(), TruthValue::True); + } + + 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) + ); + } } diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/having.rs b/pgdog/src/frontend/router/parser/rewrite/statement/having.rs index 4ac694fc0..44a192599 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/having.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/having.rs @@ -186,6 +186,7 @@ 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, } @@ -314,20 +315,19 @@ mod tests { } #[test] - fn rejects_cast_operands() { - let ast = pg_raw_parse::parse( + fn rewrites_cast_operands() { + let (sql, plan) = rewrite( "SELECT customer_id FROM purchases GROUP BY customer_id \ HAVING COUNT(*)::numeric > 3", - ) - .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(_)))); + ); + + assert_eq!( + sql, + "SELECT customer_id, count(*)::numeric AS __pgdog_having_0, \ + 3 AS __pgdog_having_1, customer_id AS __pgdog_having_2 \ + FROM purchases GROUP BY customer_id HAVING count(*)::numeric > 3 OR true" + ); + assert_eq!(plan.hidden_aliases().len(), 3); } #[test]