Only optimize filter in case of multiple columns

Author: pepijnve

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

@@ -247,10 +247,12 @@ fn is_cheap_and_infallible(expr: &Arc<dyn PhysicalExpr>) -> bool {
 }
 
 /// Creates a [FilterPredicate] from a boolean array.
-fn create_filter(predicate: &BooleanArray) -> FilterPredicate {
+fn create_filter(predicate: &BooleanArray, optimize: bool) -> FilterPredicate {
     let mut filter_builder = FilterBuilder::new(predicate);
-    // Always optimize the filter since we use them multiple times.
-    filter_builder = filter_builder.optimize();
+    if optimize {
+        // Always optimize the filter since we use them multiple times.
+        filter_builder = filter_builder.optimize();
+    }
     filter_builder.build()
 }
 
@@ -846,7 +848,7 @@ impl CaseBody {
                     result_builder.add_branch_result(&remainder_rows, nulls_value)?;
                 } else {
                     // Filter out the null rows and evaluate the else expression for those
-                    let nulls_filter = create_filter(&not(&base_not_nulls)?);
+                    let nulls_filter = create_filter(&not(&base_not_nulls)?, true);
                     let nulls_batch =
                         filter_record_batch(&remainder_batch, &nulls_filter)?;
                     let nulls_rows = filter_array(&remainder_rows, &nulls_filter)?;
@@ -861,7 +863,7 @@ impl CaseBody {
             }
 
             // Remove the null rows from the remainder batch
-            let not_null_filter = create_filter(&base_not_nulls);
+            let not_null_filter = create_filter(&base_not_nulls, true);
             remainder_batch =
                 Cow::Owned(filter_record_batch(&remainder_batch, &not_null_filter)?);
             remainder_rows = filter_array(&remainder_rows, &not_null_filter)?;
@@ -907,7 +909,7 @@ impl CaseBody {
             // for the current branch
             // Still no need to call `prep_null_mask_filter` since `create_filter` will already do
             // this unconditionally.
-            let then_filter = create_filter(&when_value);
+            let then_filter = create_filter(&when_value, true);
             let then_batch = filter_record_batch(&remainder_batch, &then_filter)?;
             let then_rows = filter_array(&remainder_rows, &then_filter)?;
 
@@ -930,7 +932,7 @@ impl CaseBody {
                     not(&prep_null_mask_filter(&when_value))
                 }
             }?;
-            let next_filter = create_filter(&next_selection);
+            let next_filter = create_filter(&next_selection, true);
             remainder_batch =
                 Cow::Owned(filter_record_batch(&remainder_batch, &next_filter)?);
             remainder_rows = filter_array(&remainder_rows, &next_filter)?;
@@ -996,7 +998,7 @@ impl CaseBody {
             // for the current branch
             // Still no need to call `prep_null_mask_filter` since `create_filter` will already do
             // this unconditionally.
-            let then_filter = create_filter(when_value);
+            let then_filter = create_filter(when_value, true);
             let then_batch = filter_record_batch(&remainder_batch, &then_filter)?;
             let then_rows = filter_array(&remainder_rows, &then_filter)?;
 
@@ -1019,7 +1021,7 @@ impl CaseBody {
                     not(&prep_null_mask_filter(when_value))
                 }
             }?;
-            let next_filter = create_filter(&next_selection);
+            let next_filter = create_filter(&next_selection, true);
             remainder_batch =
                 Cow::Owned(filter_record_batch(&remainder_batch, &next_filter)?);
             remainder_rows = filter_array(&remainder_rows, &next_filter)?;
@@ -1044,20 +1046,29 @@ impl CaseBody {
         when_value: &BooleanArray,
         return_type: &DataType,
     ) -> Result<ColumnarValue> {
-        let when_filter = create_filter(&when_value);
+        let when_value = match when_value.null_count() {
+            0 => Cow::Borrowed(when_value),
+            _ => {
+                // `prep_null_mask_filter` is required to ensure null is treated as false
+                Cow::Owned(prep_null_mask_filter(when_value))
+            }
+        };
+
+        let optimize_filter = batch.num_columns() > 1;
+        let when_filter = create_filter(&when_value, optimize_filter);
         let then_batch = filter_record_batch(batch, &when_filter)?;
         let then_value = self.when_then_expr[0].1.evaluate(&then_batch)?;
 
         let else_selection = not(&when_value)?;
-        let else_filter = create_filter(&else_selection);
+        let else_filter = create_filter(&else_selection, optimize_filter);
         let else_batch = filter_record_batch(batch, &else_filter)?;
         let e = self.else_expr.as_ref().unwrap();
         // keep `else_expr`'s data type and return type consistent
         let else_expr = try_cast(Arc::clone(e), &batch.schema(), return_type.clone())
             .unwrap_or_else(|_| Arc::clone(e));
         let else_value = else_expr.evaluate(&else_batch)?;
 
-        Ok(ColumnarValue::Array(zip_unaligned(
+        Ok(ColumnarValue::Array(merge(
             &when_value,
             then_value,
             else_value,
@@ -1215,14 +1226,6 @@ impl CaseExpr {
             return self.body.else_expr.as_ref().unwrap().evaluate(batch);
         }
 
-        let when_value = match when_value.null_count() {
-            0 => Cow::Borrowed(when_value),
-            _ => {
-                // `prep_null_mask_filter` is required to ensure null is treated as false
-                Cow::Owned(prep_null_mask_filter(when_value))
-            }
-        };
-
         if projected.projection.len() < batch.num_columns() {
             let projected_batch = batch.project(&projected.projection)?;
             projected
@@ -2236,7 +2239,7 @@ mod tests {
             PartialResultIndex::try_new(2).unwrap(),
         ];
 
-        let merged = merge(&[a1, a2, a3], &indices).unwrap();
+        let merged = merge_n(&[a1, a2, a3], &indices).unwrap();
         let merged = merged.as_string::<i32>();
 
         assert_eq!(merged.len(), indices.len());