fix: handle scalar predicates in CASE expressions to prevent internal errors for InfallibleExprOrNull eval method

Author: avkirilishin

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

@@ -344,7 +344,16 @@ impl CaseExpr {
     fn case_column_or_null(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
         let when_expr = &self.when_then_expr[0].0;
         let then_expr = &self.when_then_expr[0].1;
-        if let ColumnarValue::Array(bit_mask) = when_expr.evaluate(batch)? {
+
+        let when_expr_value = when_expr.evaluate(batch)?;
+        let when_expr_value = match when_expr_value {
+            ColumnarValue::Scalar(_) => {
+                ColumnarValue::Array(when_expr_value.into_array(batch.num_rows())?)
+            }
+            other => other,
+        };
+
+        if let ColumnarValue::Array(bit_mask) = when_expr_value {
             let bit_mask = bit_mask
                 .as_any()
                 .downcast_ref::<BooleanArray>()
@@ -896,6 +905,53 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn case_with_scalar_predicate() -> Result<()> {
+        let batch = case_test_batch_nulls()?;
+        let schema = batch.schema();
+
+        // SELECT CASE WHEN TRUE THEN load4 END
+        let when = lit(true);
+        let then = col("load4", &schema)?;
+        let expr = generate_case_when_with_type_coercion(
+            None,
+            vec![(when, then)],
+            None,
+            schema.as_ref(),
+        )?;
+
+        // many rows
+        let result = expr
+            .evaluate(&batch)?
+            .into_array(batch.num_rows())
+            .expect("Failed to convert to array");
+        let result =
+            as_float64_array(&result).expect("failed to downcast to Float64Array");
+        let expected = &Float64Array::from(vec![
+            Some(1.77),
+            None,
+            None,
+            Some(1.78),
+            None,
+            Some(1.77),
+        ]);
+        assert_eq!(expected, result);
+
+        // one row
+        let expected = Float64Array::from(vec![Some(1.1)]);
+        let batch =
+            RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(expected.clone())])?;
+        let result = expr
+            .evaluate(&batch)?
+            .into_array(batch.num_rows())
+            .expect("Failed to convert to array");
+        let result =
+            as_float64_array(&result).expect("failed to downcast to Float64Array");
+        assert_eq!(&expected, result);
+
+        Ok(())
+    }
+
     #[test]
     fn case_expr_matches_and_nulls() -> Result<()> {
         let batch = case_test_batch_nulls()?;