Skip to content

Commit 0d64069

Browse files
committed
Improve performance of CASE WHEN x THEN y ELSE NULL expressions
1 parent 9de192a commit 0d64069

1 file changed

Lines changed: 43 additions & 26 deletions

File tree

  • datafusion/physical-expr/src/expressions

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

Lines changed: 43 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@ use crate::expressions::case::literal_lookup_table::LiteralLookupTable;
4242
use arrow::compute::kernels::merge::{MergeIndex, merge, merge_n};
4343
use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion};
4444
use datafusion_physical_expr_common::datum::compare_with_eq;
45+
use datafusion_physical_expr_common::utils::scatter;
4546
use itertools::Itertools;
4647
use std::fmt::{Debug, Formatter};
4748

@@ -659,7 +660,7 @@ impl CaseExpr {
659660
&& body.else_expr.as_ref().unwrap().as_any().is::<Literal>()
660661
{
661662
EvalMethod::ScalarOrScalar
662-
} else if body.when_then_expr.len() == 1 && body.else_expr.is_some() {
663+
} else if body.when_then_expr.len() == 1 {
663664
EvalMethod::ExpressionOrExpression(body.project()?)
664665
} else {
665666
EvalMethod::NoExpression(body.project()?)
@@ -961,32 +962,40 @@ impl CaseBody {
961962
let then_batch = filter_record_batch(batch, &when_filter)?;
962963
let then_value = self.when_then_expr[0].1.evaluate(&then_batch)?;
963964

964-
let else_selection = not(&when_value)?;
965-
let else_filter = create_filter(&else_selection, optimize_filter);
966-
let else_batch = filter_record_batch(batch, &else_filter)?;
967-
968-
// keep `else_expr`'s data type and return type consistent
969-
let e = self.else_expr.as_ref().unwrap();
970-
let return_type = self.data_type(&batch.schema())?;
971-
let else_expr = try_cast(Arc::clone(e), &batch.schema(), return_type.clone())
972-
.unwrap_or_else(|_| Arc::clone(e));
973-
974-
let else_value = else_expr.evaluate(&else_batch)?;
975-
976-
Ok(ColumnarValue::Array(match (then_value, else_value) {
977-
(ColumnarValue::Array(t), ColumnarValue::Array(e)) => {
978-
merge(&when_value, &t, &e)
979-
}
980-
(ColumnarValue::Scalar(t), ColumnarValue::Array(e)) => {
981-
merge(&when_value, &t.to_scalar()?, &e)
982-
}
983-
(ColumnarValue::Array(t), ColumnarValue::Scalar(e)) => {
984-
merge(&when_value, &t, &e.to_scalar()?)
965+
match &self.else_expr {
966+
None => {
967+
let then_array = then_value.to_array(when_value.true_count())?;
968+
scatter(&when_value, then_array.as_ref()).map(ColumnarValue::Array)
985969
}
986-
(ColumnarValue::Scalar(t), ColumnarValue::Scalar(e)) => {
987-
merge(&when_value, &t.to_scalar()?, &e.to_scalar()?)
970+
Some(else_expr) => {
971+
let else_selection = not(&when_value)?;
972+
let else_filter = create_filter(&else_selection, optimize_filter);
973+
let else_batch = filter_record_batch(batch, &else_filter)?;
974+
975+
// keep `else_expr`'s data type and return type consistent
976+
let return_type = self.data_type(&batch.schema())?;
977+
let else_expr =
978+
try_cast(Arc::clone(else_expr), &batch.schema(), return_type.clone())
979+
.unwrap_or_else(|_| Arc::clone(else_expr));
980+
981+
let else_value = else_expr.evaluate(&else_batch)?;
982+
983+
Ok(ColumnarValue::Array(match (then_value, else_value) {
984+
(ColumnarValue::Array(t), ColumnarValue::Array(e)) => {
985+
merge(&when_value, &t, &e)
986+
}
987+
(ColumnarValue::Scalar(t), ColumnarValue::Array(e)) => {
988+
merge(&when_value, &t.to_scalar()?, &e)
989+
}
990+
(ColumnarValue::Array(t), ColumnarValue::Scalar(e)) => {
991+
merge(&when_value, &t, &e.to_scalar()?)
992+
}
993+
(ColumnarValue::Scalar(t), ColumnarValue::Scalar(e)) => {
994+
merge(&when_value, &t.to_scalar()?, &e.to_scalar()?)
995+
}
996+
}?))
988997
}
989-
}?))
998+
}
990999
}
9911000
}
9921001

@@ -1137,7 +1146,15 @@ impl CaseExpr {
11371146
self.body.when_then_expr[0].1.evaluate(batch)
11381147
} else if true_count == 0 {
11391148
// All input rows are false/null, just call the 'else' expression
1140-
self.body.else_expr.as_ref().unwrap().evaluate(batch)
1149+
match &self.body.else_expr {
1150+
Some(else_expr) => else_expr.evaluate(batch),
1151+
None => {
1152+
let return_type = self.data_type(&batch.schema())?;
1153+
Ok(ColumnarValue::Scalar(ScalarValue::try_new_null(
1154+
&return_type,
1155+
)?))
1156+
}
1157+
}
11411158
} else if projected.projection.len() < batch.num_columns() {
11421159
// The case expressions do not use all the columns of the input batch.
11431160
// Project first to reduce time spent filtering.

0 commit comments

Comments
 (0)