feat: support numeric, float and cast operands in cross-shard HAVING

integration/rust/tests/integration/having.rs

@@ -159,13 +159,64 @@ async fn having_null_semantics_on_merged_groups() -> Result<(), Box<dyn std::err
     Ok(())
 }
 
+#[tokio::test]
+async fn having_numeric_float_and_cast_operands_use_merged_values()
+-> Result<(), Box<dyn std::error::Error>> {
+    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<Vec<i64>, 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(),
         )

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)] {

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<Tru
 }
 
 fn numeric_ordering(left: &Datum, right: &Datum) -> Option<Ordering> {
-    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<i128> {
@@ -171,10 +182,33 @@ fn integer(value: &Datum) -> Option<i128> {
     }
 }
 
+fn exact_numeric(value: &Datum) -> Option<Numeric> {
+    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<f64> {
+    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)
+        );
+    }
 }

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]