@@ -42,6 +42,7 @@ use crate::expressions::case::literal_lookup_table::LiteralLookupTable;
4242use arrow:: compute:: kernels:: merge:: { MergeIndex , merge, merge_n} ;
4343use datafusion_common:: tree_node:: { Transformed , TreeNode , TreeNodeRecursion } ;
4444use datafusion_physical_expr_common:: datum:: compare_with_eq;
45+ use datafusion_physical_expr_common:: utils:: scatter;
4546use itertools:: Itertools ;
4647use 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