Skip to content

Commit

Permalink
fix case_column_or_null with nullable when conditions
Browse files Browse the repository at this point in the history
  • Loading branch information
zhangli20 committed Dec 23, 2024
1 parent 8fd792f commit 86b4a41
Show file tree
Hide file tree
Showing 2 changed files with 32 additions and 3 deletions.
31 changes: 30 additions & 1 deletion datafusion/physical-expr/src/expressions/case.rs
Original file line number Diff line number Diff line change
Expand Up @@ -346,7 +346,10 @@ impl CaseExpr {
.downcast_ref::<BooleanArray>()
.expect("predicate should evaluate to a boolean array");
// invert the bitmask
let bit_mask = not(bit_mask)?;
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)?))
Expand Down Expand Up @@ -885,6 +888,32 @@ mod tests {
Ok(())
}

#[test]
fn test_when_null_and_some_cond_else_null() -> Result<()> {
let batch = case_test_batch()?;
let schema = batch.schema();

let when = binary(
Arc::new(Literal::new(ScalarValue::Boolean(None))),
Operator::And,
binary(col("a", &schema)?, Operator::Eq, lit("foo"), &schema)?,
&schema,
)?;
let then = col("a", &schema)?;

// SELECT CASE WHEN (NULL AND a = 'foo') THEN a ELSE NULL END
let expr = Arc::new(CaseExpr::try_new(None, vec![(when, then)], None)?);
let result = expr
.evaluate(&batch)?
.into_array(batch.num_rows())
.expect("Failed to convert to array");
let result = as_string_array(&result);

// all result values should be null
assert_eq!(result.logical_null_count(), batch.num_rows());
Ok(())
}

fn case_test_batch() -> Result<RecordBatch> {
let schema = Schema::new(vec![Field::new("a", DataType::Utf8, true)]);
let a = StringArray::from(vec![Some("foo"), Some("baz"), None, Some("bar")]);
Expand Down
4 changes: 2 additions & 2 deletions datafusion/sqllogictest/test_files/case.slt
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,7 @@ NULL
6
NULL
NULL
7
NULL

# column or implicit null
query I
Expand All @@ -61,7 +61,7 @@ NULL
6
NULL
NULL
7
NULL

# scalar or scalar (string)
query T
Expand Down

0 comments on commit 86b4a41

Please sign in to comment.