Simplify error handling in case.rs (#13990) (#14033)

* Simplify error handling in case.rs (#13990) * Fix issues causing GitHub checks to fail * Update datafusion/physical-expr/src/expressions/case.rs Co-authored-by: Andrew Lamb <[email protected]> --------- Co-authored-by: Sergey Zhukov <[email protected]> Co-authored-by: Andrew Lamb <[email protected]>
apache · Jan 9, 2025 · ac9584e · ac9584e
1 parent 61afb0d
commit ac9584e
Showing 1 changed file with 15 additions and 21 deletions.
diff --git a/datafusion/physical-expr/src/expressions/case.rs b/datafusion/physical-expr/src/expressions/case.rs
@@ -27,7 +27,9 @@ use arrow::compute::kernels::zip::zip;
 use arrow::compute::{and, and_not, is_null, not, nullif, or, prep_null_mask_filter};
 use arrow::datatypes::{DataType, Schema};
 use datafusion_common::cast::as_boolean_array;
-use datafusion_common::{exec_err, internal_err, DataFusionError, Result, ScalarValue};
+use datafusion_common::{
+    exec_err, internal_datafusion_err, internal_err, DataFusionError, Result, ScalarValue,
+};
 use datafusion_expr::ColumnarValue;
 
 use super::{Column, Literal};
@@ -249,10 +251,9 @@ impl CaseExpr {
             remainder = and_not(&remainder, &when_match)?;
         }
 
-        if let Some(e) = &self.else_expr {
+        if let Some(e) = self.else_expr() {
             // keep `else_expr`'s data type and return type consistent
-            let expr = try_cast(Arc::clone(e), &batch.schema(), return_type.clone())
-                .unwrap_or_else(|_| Arc::clone(e));
+            let expr = try_cast(Arc::clone(e), &batch.schema(), return_type.clone())?;
             // null and unmatched tuples should be assigned else value
             remainder = or(&base_nulls, &remainder)?;
             let else_ = expr
@@ -282,11 +283,8 @@ impl CaseExpr {
                 .0
                 .evaluate_selection(batch, &remainder)?;
             let when_value = when_value.into_array(batch.num_rows())?;
-            let when_value = as_boolean_array(&when_value).map_err(|e| {
-                DataFusionError::Context(
-                    "WHEN expression did not return a BooleanArray".to_string(),
-                    Box::new(e),
-                )
+            let when_value = as_boolean_array(&when_value).map_err(|_| {
+                internal_datafusion_err!("WHEN expression did not return a BooleanArray")
             })?;
             // Treat 'NULL' as false value
             let when_value = match when_value.null_count() {
@@ -322,10 +320,9 @@ impl CaseExpr {
             remainder = and_not(&remainder, &when_value)?;
         }
 
-        if let Some(e) = &self.else_expr {
+        if let Some(e) = self.else_expr() {
             // keep `else_expr`'s data type and return type consistent
-            let expr = try_cast(Arc::clone(e), &batch.schema(), return_type.clone())
-                .unwrap_or_else(|_| Arc::clone(e));
+            let expr = try_cast(Arc::clone(e), &batch.schema(), return_type.clone())?;
             let else_ = expr
                 .evaluate_selection(batch, &remainder)?
                 .into_array(batch.num_rows())?;
@@ -376,11 +373,8 @@ impl CaseExpr {
         // evaluate when expression
         let when_value = self.when_then_expr[0].0.evaluate(batch)?;
         let when_value = when_value.into_array(batch.num_rows())?;
-        let when_value = as_boolean_array(&when_value).map_err(|e| {
-            DataFusionError::Context(
-                "WHEN expression did not return a BooleanArray".to_string(),
-                Box::new(e),
-            )
+        let when_value = as_boolean_array(&when_value).map_err(|_| {
+            internal_datafusion_err!("WHEN expression did not return a BooleanArray")
         })?;
 
         // Treat 'NULL' as false value
@@ -393,12 +387,12 @@ impl CaseExpr {
         let then_value = self.when_then_expr[0].1.evaluate(batch)?;
         let then_value = Scalar::new(then_value.into_array(1)?);
 
+        let Some(e) = self.else_expr() else {
+            return internal_err!("expression did not evaluate to an array");
+        };
         // keep `else_expr`'s data type and return type consistent
-        let e = self.else_expr.as_ref().unwrap();
-        let expr = try_cast(Arc::clone(e), &batch.schema(), return_type)
-            .unwrap_or_else(|_| Arc::clone(e));
+        let expr = try_cast(Arc::clone(e), &batch.schema(), return_type)?;
         let else_ = Scalar::new(expr.evaluate(batch)?.into_array(1)?);
-
         Ok(ColumnarValue::Array(zip(&when_value, &then_value, &else_)?))
     }