Support Expr::Case in aggregate module (#945)

gluesql · Oct 15, 2022 · 24c30da · 24c30da
1 parent 2d18de8
commit 24c30da
Show file tree

Hide file tree

Showing 6 changed files with 62 additions and 3 deletions.
diff --git a/core/src/executor/aggregate/mod.rs b/core/src/executor/aggregate/mod.rs
@@ -179,6 +179,7 @@ async fn aggregate<'a>(
     expr: &'a Expr,
 ) -> Result<State<'a>> {
     let aggr = |state, expr| aggregate(state, filter_context.as_ref().map(Rc::clone), expr);
+
     match expr {
         Expr::Between {
             expr, low, high, ..
@@ -200,6 +201,26 @@ async fn aggregate<'a>(
         }
         Expr::UnaryOp { expr, .. } => aggr(state, expr).await,
         Expr::Nested(expr) => aggr(state, expr).await,
+        Expr::Case {
+            operand,
+            when_then,
+            else_result,
+        } => {
+            let operand = std::iter::once(operand.as_ref())
+                .filter_map(|operand| operand.map(|operand| &**operand));
+            let when_then = when_then
+                .iter()
+                .flat_map(|(when, then)| std::iter::once(when).chain(std::iter::once(then)));
+            let else_result = std::iter::once(else_result.as_ref())
+                .filter_map(|else_result| else_result.map(|else_result| &**else_result));
+
+            stream::iter(operand.chain(when_then).chain(else_result))
+                .fold(
+                    Ok(state),
+                    |state, expr| async move { aggr(state?, expr).await },
+                )
+                .await
+        }
         Expr::Aggregate(aggr) => state.accumulate(filter_context, aggr.as_ref()).await,
         _ => Ok(state),
     }
@@ -213,6 +234,21 @@ fn check(expr: &Expr) -> bool {
         Expr::BinaryOp { left, right, .. } => check(left) || check(right),
         Expr::UnaryOp { expr, .. } => check(expr),
         Expr::Nested(expr) => check(expr),
+        Expr::Case {
+            operand,
+            when_then,
+            else_result,
+        } => {
+            operand.as_ref().map(|expr| check(&*expr)).unwrap_or(false)
+                || when_then
+                    .iter()
+                    .map(|(when, then)| check(when) || check(then))
+                    .any(identity)
+                || else_result
+                    .as_ref()
+                    .map(|expr| check(&*expr))
+                    .unwrap_or(false)
+        }
         Expr::Aggregate(_) => true,
         _ => false,
     }

diff --git a/core/src/executor/aggregate/state.rs b/core/src/executor/aggregate/state.rs
@@ -11,9 +11,11 @@ use {
     std::{cmp::Ordering, rc::Rc},
     utils::{IndexMap, Vector},
 };
+
 type Group = Rc<Vec<Key>>;
 type ValuesMap<'a> = HashMap<&'a Aggregate, Value>;
 type Context<'a> = Rc<BlendContext<'a>>;
+
 enum AggrValue {
     Count {
         wildcard: bool,

diff --git a/test-suite/src/aggregate/min.rs b/test-suite/src/aggregate/min.rs
@@ -46,8 +46,20 @@ test_case!(min, async move {
             ),
         ),
         (
-            "SELECT SUM(quantity) * 2 + MIN(quantity) - 3 / 1 FROM Item",
-            select!("SUM(quantity) * 2 + MIN(quantity) - 3 / 1"; I64; 91),
+            "SELECT SUM(quantity) * 2 + MIN(quantity) - 3 / 1 FROM Item;",
+            select!(
+                "SUM(quantity) * 2 + MIN(quantity) - 3 / 1"
+                I64;
+                91
+            ),
+        ),
+        (
+            "SELECT MIN(CASE WHEN quantity > 5 THEN id END) FROM Item;",
+            select!(
+                "MIN(CASE WHEN quantity > 5 THEN id END)"
+                I64;
+                1
+            ),
         ),
     ];
 

diff --git a/test-suite/src/aggregate/sum.rs b/test-suite/src/aggregate/sum.rs
@@ -70,6 +70,14 @@ test_case!(sum, async move {
                 174
             ),
         ),
+        (
+            "SELECT SUM(CASE WHEN id > 3 THEN quantity ELSE 0 END) FROM Item;",
+            select!(
+                "SUM(CASE WHEN id > 3 THEN quantity ELSE 0 END)"
+                I64;
+                28
+            ),
+        ),
     ];
 
     for (sql, expected) in test_cases {

diff --git a/test-suite/src/arithmetic/error.rs b/test-suite/src/arithmetic/error.rs
@@ -4,7 +4,7 @@ use {
     gluesql_core::{
         data::{Literal, LiteralError, NumericBinaryOperator, ValueError},
         executor::{EvaluateError, UpdateError},
-        prelude::Value::{self},
+        prelude::Value,
     },
     std::borrow::Cow,
 };

diff --git a/test-suite/src/inline_view.rs b/test-suite/src/inline_view.rs
@@ -6,6 +6,7 @@ use {
         translate::TranslateError,
     },
 };
+
 test_case!(inline_view, async move {
     let test_cases = [
         (
-Original file line number
+Diff line change
@@ Expand Up / @@ -6,6 +6,7 @@ use { @@
             translate::TranslateError,
         },
     };
     test_case!(inline_view, async move {
         let test_cases = [
             (
@@ Expand Down @@