From e79dd148ad9f79b502d9305a22dd5a5fcf9ec8ea Mon Sep 17 00:00:00 2001 From: duripeng Date: Mon, 11 Apr 2022 05:17:46 +0800 Subject: [PATCH 1/3] case when support literal --- datafusion/core/tests/sql/expr.rs | 34 +++++++++++++++++++ .../physical-expr/src/expressions/case.rs | 11 +++++- 2 files changed, 44 insertions(+), 1 deletion(-) diff --git a/datafusion/core/tests/sql/expr.rs b/datafusion/core/tests/sql/expr.rs index 565af320394ea..765d495830d6e 100644 --- a/datafusion/core/tests/sql/expr.rs +++ b/datafusion/core/tests/sql/expr.rs @@ -109,6 +109,40 @@ async fn case_when_else_with_base_expr() -> Result<()> { Ok(()) } +#[tokio::test] +async fn case_when_else_with_null_literal() -> Result<()> { + let ctx = create_case_context()?; + let sql = "SELECT \ + CASE WHEN c1 = 'a' THEN 1 \ + WHEN NULL THEN 2 \ + ELSE 999 END \ + FROM t1"; + let actual = execute_to_batches(&ctx, sql).await; + let expected = vec![ + "+----------------------------------------------------------------------------------------------+", + "| CASE WHEN #t1.c1 = Utf8(\"a\") THEN Int64(1) WHEN Utf8(NULL) THEN Int64(2) ELSE Int64(999) END |", + "+----------------------------------------------------------------------------------------------+", + "| 1 |", + "| 999 |", + "| 999 |", + "| 999 |", + "+----------------------------------------------------------------------------------------------+", + ]; + assert_batches_eq!(expected, &actual); + + let sql = "SELECT CASE WHEN NULL THEN 'foo' ELSE 'bar' END"; + let actual = execute_to_batches(&ctx, sql).await; + let expected = vec![ + "+------------------------------------------------------------+", + "| CASE WHEN Utf8(NULL) THEN Utf8(\"foo\") ELSE Utf8(\"bar\") END |", + "+------------------------------------------------------------+", + "| bar |", + "+------------------------------------------------------------+", + ]; + assert_batches_eq!(expected, &actual); + Ok(()) +} + #[tokio::test] async fn query_not() -> Result<()> { let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Boolean, true)])); diff --git a/datafusion/physical-expr/src/expressions/case.rs b/datafusion/physical-expr/src/expressions/case.rs index df763ec9a6f60..5c4d9c42e8a91 100644 --- a/datafusion/physical-expr/src/expressions/case.rs +++ b/datafusion/physical-expr/src/expressions/case.rs @@ -24,7 +24,7 @@ use arrow::compute::kernels::zip::zip; use arrow::compute::{and, eq_dyn, is_null, not, or, or_kleene}; use arrow::datatypes::{DataType, Schema}; use arrow::record_batch::RecordBatch; -use datafusion_common::{DataFusionError, Result}; +use datafusion_common::{DataFusionError, Result, ScalarValue}; use datafusion_expr::ColumnarValue; type WhenThen = (Arc, Arc); @@ -178,6 +178,15 @@ impl CaseExpr { let when_value = self.when_then_expr[i] .0 .evaluate_selection(batch, &remainder)?; + let when_value = match when_value { + ColumnarValue::Scalar(value) => + if value.is_null() { + ColumnarValue::Scalar(ScalarValue::Boolean(Some(false))) + } else { + ColumnarValue::Scalar(value) + } + _ => when_value + }; let when_value = when_value.into_array(batch.num_rows()); let when_value = when_value .as_ref() From 178b4aa2b857d4f1bca7a66dcefcb0d004628a13 Mon Sep 17 00:00:00 2001 From: duripeng Date: Mon, 11 Apr 2022 05:24:45 +0800 Subject: [PATCH 2/3] fmt fix --- datafusion/physical-expr/src/expressions/case.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/datafusion/physical-expr/src/expressions/case.rs b/datafusion/physical-expr/src/expressions/case.rs index 5c4d9c42e8a91..9eac3103f0e15 100644 --- a/datafusion/physical-expr/src/expressions/case.rs +++ b/datafusion/physical-expr/src/expressions/case.rs @@ -178,14 +178,16 @@ impl CaseExpr { let when_value = self.when_then_expr[i] .0 .evaluate_selection(batch, &remainder)?; + // Treat 'NULL' as false value let when_value = match when_value { - ColumnarValue::Scalar(value) => + ColumnarValue::Scalar(value) => { if value.is_null() { ColumnarValue::Scalar(ScalarValue::Boolean(Some(false))) } else { ColumnarValue::Scalar(value) } - _ => when_value + } + _ => when_value, }; let when_value = when_value.into_array(batch.num_rows()); let when_value = when_value From 7092618ad806ba388d12025baba7c10bffe396b8 Mon Sep 17 00:00:00 2001 From: duripeng Date: Mon, 11 Apr 2022 14:53:40 +0800 Subject: [PATCH 3/3] code clean --- datafusion/core/tests/sql/expr.rs | 2 +- datafusion/physical-expr/src/expressions/case.rs | 10 +++------- 2 files changed, 4 insertions(+), 8 deletions(-) diff --git a/datafusion/core/tests/sql/expr.rs b/datafusion/core/tests/sql/expr.rs index 765d495830d6e..79bf9e2c4a18d 100644 --- a/datafusion/core/tests/sql/expr.rs +++ b/datafusion/core/tests/sql/expr.rs @@ -110,7 +110,7 @@ async fn case_when_else_with_base_expr() -> Result<()> { } #[tokio::test] -async fn case_when_else_with_null_literal() -> Result<()> { +async fn case_when_else_with_null_contant() -> Result<()> { let ctx = create_case_context()?; let sql = "SELECT \ CASE WHEN c1 = 'a' THEN 1 \ diff --git a/datafusion/physical-expr/src/expressions/case.rs b/datafusion/physical-expr/src/expressions/case.rs index 9eac3103f0e15..e7db10d1772a0 100644 --- a/datafusion/physical-expr/src/expressions/case.rs +++ b/datafusion/physical-expr/src/expressions/case.rs @@ -24,7 +24,7 @@ use arrow::compute::kernels::zip::zip; use arrow::compute::{and, eq_dyn, is_null, not, or, or_kleene}; use arrow::datatypes::{DataType, Schema}; use arrow::record_batch::RecordBatch; -use datafusion_common::{DataFusionError, Result, ScalarValue}; +use datafusion_common::{DataFusionError, Result}; use datafusion_expr::ColumnarValue; type WhenThen = (Arc, Arc); @@ -180,12 +180,8 @@ impl CaseExpr { .evaluate_selection(batch, &remainder)?; // Treat 'NULL' as false value let when_value = match when_value { - ColumnarValue::Scalar(value) => { - if value.is_null() { - ColumnarValue::Scalar(ScalarValue::Boolean(Some(false))) - } else { - ColumnarValue::Scalar(value) - } + ColumnarValue::Scalar(value) if value.is_null() => { + continue; } _ => when_value, };