Restore lazy evaluation of fallible CASE

Commit 0f4b8b1 introduced a new
optimized evaluation mode for CASE expression with exactly one WHEN-THEN
clause and ELSE clause being present. Apparently, the new mode did not
take into account expensive or fallible expressions, as it
unconditionally evaluated both branches. This is definitely incorrect
for fallible expressions. For these, fallback to the original execution
mode.

Author: findepi

datafusion/physical-expr/src/expressions/case.rs

@@ -156,7 +156,10 @@ impl CaseExpr {
                 && else_expr.as_ref().unwrap().as_any().is::<Literal>()
             {
                 EvalMethod::ScalarOrScalar
-            } else if when_then_expr.len() == 1 && else_expr.is_some() {
+            } else if when_then_expr.len() == 1
+                && is_cheap_and_infallible(&(when_then_expr[0].1))
+                && else_expr.as_ref().is_some_and(is_cheap_and_infallible)
+            {
                 EvalMethod::ExpressionOrExpression
             } else {
                 EvalMethod::NoExpression
@@ -813,9 +816,18 @@ mod tests {
     }
 
     fn case_test_batch1() -> Result<RecordBatch> {
-        let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]);
+        let schema = Schema::new(vec![
+            Field::new("a", DataType::Int32, true),
+            Field::new("b", DataType::Int32, true),
+            Field::new("c", DataType::Int32, true),
+        ]);
         let a = Int32Array::from(vec![Some(1), Some(0), None, Some(5)]);
-        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)])?;
+        let b = Int32Array::from(vec![Some(3), None, Some(14), Some(7)]);
+        let c = Int32Array::from(vec![Some(0), Some(-3), Some(777), None]);
+        let batch = RecordBatch::try_new(
+            Arc::new(schema),
+            vec![Arc::new(a), Arc::new(b), Arc::new(c)],
+        )?;
         Ok(batch)
     }
 
@@ -1306,18 +1318,8 @@ mod tests {
             lit(2i32),
             &batch.schema(),
         )?;
-        let then = binary(
-            col("a", &schema)?,
-            Operator::Plus,
-            lit(1i32),
-            &batch.schema(),
-        )?;
-        let else_expr = binary(
-            col("a", &schema)?,
-            Operator::Minus,
-            lit(1i32),
-            &batch.schema(),
-        )?;
+        let then = col("b", &schema)?;
+        let else_expr = col("c", &schema)?;
         let expr = CaseExpr::try_new(None, vec![(when, then)], Some(else_expr))?;
         assert!(matches!(
             expr.eval_method,
@@ -1329,7 +1331,7 @@ mod tests {
             .expect("Failed to convert to array");
         let result = as_int32_array(&result).expect("failed to downcast to Int32Array");
 
-        let expected = &Int32Array::from(vec![Some(2), Some(1), None, Some(4)]);
+        let expected = &Int32Array::from(vec![Some(3), None, Some(777), None]);
 
         assert_eq!(expected, result);
         Ok(())

datafusion/sqllogictest/test_files/case.slt

@@ -467,7 +467,18 @@ FROM t;
 ----
 [{foo: blarg}]
 
+query II
+SELECT v, CASE WHEN v != 0 THEN 10/v ELSE 42 END FROM (VALUES (0), (1), (2)) t(v)
+----
+0 42
+1 10
+2 5
 
+query II
+SELECT v, CASE WHEN v < 0 THEN 10/0 ELSE 1 END FROM (VALUES (1), (2)) t(v)
+----
+1 1
+2 1
 
 statement ok
 drop table t