Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 3 commits into from
Jan 9, 2025
Merged

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

Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
4a39d4c
Simplify error handling in case.rs (#13990)
Jan 7, 2025
53d8876
Fix issues causing GitHub checks to fail
Jan 8, 2025
702128f
Update datafusion/physical-expr/src/expressions/case.rs
cj-zhukov Jan 9, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 19 additions & 24 deletions datafusion/physical-expr/src/expressions/case.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -242,10 +244,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())?;
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this change is an improvement -- previously the error was ignored. In this new formulation it is returned.

// null and unmatched tuples should be assigned else value
remainder = or(&base_nulls, &remainder)?;
let else_ = expr
Expand Down Expand Up @@ -275,11 +276,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() {
Expand Down Expand Up @@ -315,10 +313,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())?;
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

likewise here it is a good change not to silently ignore the error -- that would likely just make any future error more confusing

let else_ = expr
.evaluate_selection(batch, &remainder)?
.into_array(batch.num_rows())?;
Expand Down Expand Up @@ -369,11 +366,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(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Another potential way to to keep the context would be something like

.map_err(|e| e.context("WHEN expression did not return a BooleanArray"))

I think what you have here is good though as it makes it clear this is an internal / non expected error

"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
Expand All @@ -386,13 +380,14 @@ impl CaseExpr {
let then_value = self.when_then_expr[0].1.evaluate(batch)?;
let then_value = Scalar::new(then_value.into_array(1)?);

// 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 else_ = Scalar::new(expr.evaluate(batch)?.into_array(1)?);

Ok(ColumnarValue::Array(zip(&when_value, &then_value, &else_)?))
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)?;
let else_ = Scalar::new(expr.evaluate(batch)?.into_array(1)?);
Ok(ColumnarValue::Array(zip(&when_value, &then_value, &else_)?))
} else {
internal_err!("expression did not evaluate to an array")
}
cj-zhukov marked this conversation as resolved.
Show resolved Hide resolved
}
}

Expand Down
Loading