feat: support AND, OR, NOT and NULL in cross-shard HAVING

integration/rust/tests/integration/having.rs

@@ -75,6 +75,90 @@ async fn prepared_having_parameter_is_typed_by_postgres() -> Result<(), Box<dyn
     Ok(())
 }
 
+#[tokio::test]
+async fn having_boolean_expression_filters_merged_groups() -> 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, 10), (1, 11), (2, 12), (3, 5)",
+    )
+    .await?;
+    seed(
+        &sharded,
+        1,
+        "INSERT INTO having_merge_test VALUES (1, 20), (2, 21), (2, 22), (3, 6)",
+    )
+    .await?;
+
+    let rows = sharded
+        .fetch_all(
+            "SELECT customer_id, COUNT(*) AS order_count \
+             FROM having_merge_test \
+             GROUP BY customer_id \
+             HAVING (COUNT(*) > 2 AND NOT (MAX(order_id) > 21)) OR MIN(order_id) < 6 \
+             ORDER BY customer_id",
+        )
+        .await?;
+
+    let groups: Vec<(i64, i64)> = rows
+        .iter()
+        .map(|row| (row.get("customer_id"), row.get("order_count")))
+        .collect();
+    assert_eq!(groups, vec![(1, 3), (3, 2)]);
+    assert!(rows.iter().all(|row| row.columns().len() == 2));
+
+    cleanup(&sharded).await;
+    Ok(())
+}
+
+#[tokio::test]
+async fn having_null_semantics_on_merged_groups() -> 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, 10), (2, 5), (3, NULL)",
+    )
+    .await?;
+    seed(
+        &sharded,
+        1,
+        "INSERT INTO having_merge_test VALUES (1, 20), (3, NULL)",
+    )
+    .await?;
+
+    let null_groups = sharded
+        .fetch_all(
+            "SELECT customer_id FROM having_merge_test \
+             GROUP BY customer_id \
+             HAVING MAX(order_id) IS NULL",
+        )
+        .await?;
+    assert_eq!(null_groups.len(), 1);
+    assert_eq!(null_groups[0].get::<i64, _>("customer_id"), 3);
+
+    // NOT of an unknown comparison is still unknown, so customer 3 must be filtered out.
+    let negated = sharded
+        .fetch_all(
+            "SELECT customer_id FROM having_merge_test \
+             GROUP BY customer_id \
+             HAVING NOT (MAX(order_id) > 15) \
+             ORDER BY customer_id",
+        )
+        .await?;
+    let customers: Vec<i64> = negated.iter().map(|row| row.get("customer_id")).collect();
+    assert_eq!(customers, vec![2]);
+
+    cleanup(&sharded).await;
+    Ok(())
+}
+
 async fn reset_table(pool: &PgPool) -> Result<(), sqlx::Error> {
     cleanup(pool).await;
     for shard in [0, 1] {

pgdog/src/backend/pool/connection/buffer/having.rs

@@ -21,6 +21,16 @@ 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
@@ -59,6 +69,39 @@ 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
+                })
+            }
         }
     }
 
@@ -168,6 +211,72 @@ mod tests {
         ));
     }
 
+    #[test]
+    fn implements_sql_three_valued_boolean_logic() {
+        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);
+        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(
+                &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
+        );
+    }
+
     #[test]
     fn compares_integers_across_widths_with_all_operators() {
         for (operator, left, right, expected) in [

pgdog/src/frontend/router/parser/having.rs

@@ -35,6 +35,13 @@ pub(crate) enum HavingExpr {
         operator: CompareOperator,
         right: HavingValue,
     },
+    And(Vec<HavingExpr>),
+    Or(Vec<HavingExpr>),
+    Not(Box<HavingExpr>),
+    IsNull {
+        value: HavingValue,
+        negated: bool,
+    },
 }
 
 /// Metadata produced while rewriting a cross-shard `HAVING` clause.

pgdog/src/frontend/router/parser/rewrite/statement/having.rs

@@ -131,8 +131,42 @@ 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::<Result<Vec<_>, _>>()?;
+                match expression.boolop {
+                    nodes::BoolExprType::AND_EXPR => Ok(HavingExpr::And(expressions)),
+                    nodes::BoolExprType::OR_EXPR => Ok(HavingExpr::Or(expressions)),
+                    nodes::BoolExprType::NOT_EXPR => {
+                        let expression = expressions
+                            .into_iter()
+                            .exactly_one()
+                            .map_err(|_| Error::UnsupportedHaving("NOT requires one expression"))?;
+                        Ok(HavingExpr::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(HavingExpr::IsNull {
+                    value: self.parse_value(test.arg())?,
+                    negated,
+                })
+            }
             _ => Err(Error::UnsupportedHaving(
-                "only a single comparison is supported",
+                "supported predicates are comparisons, AND, OR, NOT and NULL tests",
             )),
         }
     }
@@ -234,25 +268,66 @@ mod tests {
     }
 
     #[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",
+    fn rewrites_nested_boolean_null_and_parameter_expression() {
+        let (sql, plan) = rewrite(
+            "SELECT customer_id, SUM(price) FROM purchases \
+             GROUP BY customer_id \
+             HAVING NOT (COUNT(*) <= $1 OR SUM(price) IS NOT NULL) \
+             AND MAX(price) >= 10",
+        );
+
+        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, .. }
+                    ]
+                )
+            )
+        ));
+        assert_eq!(
+            sql,
+            "SELECT customer_id, sum(price), count(*) AS __pgdog_having_0, \
+             $1 AS __pgdog_having_1, sum(price) AS __pgdog_having_2, \
+             max(price) 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) >= 10) OR true"
+        );
+        assert_eq!(plan.hidden_aliases().len(), 6);
+    }
+
+    #[test]
+    fn rejects_cast_operands() {
+        let ast = pg_raw_parse::parse(
             "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}");
-        }
+        )
+        .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(_))));
     }
 
     #[test]