Improve case expr constant handling, Add .slt test (apache#14159)

jonahgao · Jan 20, 2025 · 0a85bb4 · 0a85bb4
1 parent 98e8942
commit 0a85bb4
Show file tree

Hide file tree

Showing 2 changed files with 90 additions and 24 deletions.
diff --git a/datafusion/physical-expr/src/expressions/case.rs b/datafusion/physical-expr/src/expressions/case.rs
@@ -345,34 +345,37 @@ impl CaseExpr {
         let when_expr = &self.when_then_expr[0].0;
         let then_expr = &self.when_then_expr[0].1;
 
-        let when_expr_value = when_expr.evaluate(batch)?;
-        let when_expr_value = match when_expr_value {
+        match when_expr.evaluate(batch)? {
+            // WHEN true --> column
+            ColumnarValue::Scalar(ScalarValue::Boolean(Some(true))) => {
+                then_expr.evaluate(batch)
+            }
+            // WHEN [false | null] --> NULL
             ColumnarValue::Scalar(_) => {
-                ColumnarValue::Array(when_expr_value.into_array(batch.num_rows())?)
+                // return scalar NULL value
+                ScalarValue::try_from(self.data_type(&batch.schema())?)
+                    .map(ColumnarValue::Scalar)
             }
-            other => other,
-        };
-
-        if let ColumnarValue::Array(bit_mask) = when_expr_value {
-            let bit_mask = bit_mask
-                .as_any()
-                .downcast_ref::<BooleanArray>()
-                .expect("predicate should evaluate to a boolean array");
-            // invert the bitmask
-            let bit_mask = match bit_mask.null_count() {
-                0 => not(bit_mask)?,
-                _ => not(&prep_null_mask_filter(bit_mask))?,
-            };
-            match then_expr.evaluate(batch)? {
-                ColumnarValue::Array(array) => {
-                    Ok(ColumnarValue::Array(nullif(&array, &bit_mask)?))
-                }
-                ColumnarValue::Scalar(_) => {
-                    internal_err!("expression did not evaluate to an array")
+            // WHEN column --> column
+            ColumnarValue::Array(bit_mask) => {
+                let bit_mask = bit_mask
+                    .as_any()
+                    .downcast_ref::<BooleanArray>()
+                    .expect("predicate should evaluate to a boolean array");
+                // invert the bitmask
+                let bit_mask = match bit_mask.null_count() {
+                    0 => not(bit_mask)?,
+                    _ => not(&prep_null_mask_filter(bit_mask))?,
+                };
+                match then_expr.evaluate(batch)? {
+                    ColumnarValue::Array(array) => {
+                        Ok(ColumnarValue::Array(nullif(&array, &bit_mask)?))
+                    }
+                    ColumnarValue::Scalar(_) => {
+                        internal_err!("expression did not evaluate to an array")
+                    }
                 }
             }
-        } else {
-            internal_err!("predicate did not evaluate to an array")
         }
     }
 

diff --git a/datafusion/sqllogictest/test_files/case.slt b/datafusion/sqllogictest/test_files/case.slt
@@ -235,3 +235,66 @@ SELECT CASE WHEN a < 5 THEN a + b ELSE b - NVL(a, 0) END FROM foo
 NULL
 NULL
 7
+
+# Reproducer for
+# https://github.com/apache/datafusion/issues/14099
+query I
+SELECT - 79 * + 91 * - COUNT ( * ) * + - 2 * + - NULLIF ( - 49, - COALESCE ( - + 69, - COALESCE ( + COALESCE ( - 20, ( - 18 ) * + COUNT ( * ) + - 93, - CASE 51 WHEN + COUNT ( * ) + 28 THEN 0 ELSE + 29 * + CASE ( 50 ) WHEN - ( - ( CASE WHEN NOT + 37 IS NULL THEN + COUNT ( * ) END ) ) THEN NULL WHEN - 46 + 87 * - 28 THEN 85 WHEN - COUNT ( * ) THEN NULL END END ), COUNT ( * ) - 39 ) * + 22 ) / - COUNT ( * ) )
+----
+-704522
+
+
+query B
+select case when true then false end from foo;
+----
+false
+false
+false
+false
+false
+false
+
+query I
+select case when true then a end from foo;
+----
+1
+3
+5
+NULL
+6
+NULL
+
+query I
+select case when false then a end from foo;
+----
+NULL
+NULL
+NULL
+NULL
+NULL
+NULL
+
+query I
+select case when null then a end from foo;
+----
+NULL
+NULL
+NULL
+NULL
+NULL
+NULL
+
+
+query B
+select case when a=1 then false end from foo;
+----
+false
+false
+false
+false
+false
+false
+
+
+statement ok
+drop table foo