Avoid expanding scalars

Author: pepijnve

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

@@ -804,6 +804,26 @@ impl CaseExpr {
 }
 
 impl CaseBody {
+    fn data_type(&self, input_schema: &Schema) -> Result<DataType> {
+        // since all then results have the same data type, we can choose any one as the
+        // return data type except for the null.
+        let mut data_type = DataType::Null;
+        for i in 0..self.when_then_expr.len() {
+            data_type = self.when_then_expr[i].1.data_type(input_schema)?;
+            if !data_type.equals_datatype(&DataType::Null) {
+                break;
+            }
+        }
+        // if all then results are null, we use data type of else expr instead if possible.
+        if data_type.equals_datatype(&DataType::Null) {
+            if let Some(e) = &self.else_expr {
+                data_type = e.data_type(input_schema)?;
+            }
+        }
+
+        Ok(data_type)
+    }
+
     /// See [CaseExpr::case_when_with_expr].
     fn case_when_with_expr(
         &self,
@@ -1044,7 +1064,6 @@ impl CaseBody {
         &self,
         batch: &RecordBatch,
         when_value: &BooleanArray,
-        return_type: &DataType,
     ) -> Result<ColumnarValue> {
         let when_value = match when_value.null_count() {
             0 => Cow::Borrowed(when_value),
@@ -1055,17 +1074,21 @@ impl CaseBody {
         };
 
         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, 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 e = self.else_expr.as_ref().unwrap();
+        let return_type = self.data_type(&batch.schema())?;
         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(merge(
@@ -1204,35 +1227,34 @@ impl CaseExpr {
         batch: &RecordBatch,
         projected: &ProjectedCaseBody,
     ) -> Result<ColumnarValue> {
-        let return_type = self.data_type(&batch.schema())?;
-
         // evaluate when condition on batch
         let when_value = self.body.when_then_expr[0].0.evaluate(batch)?;
-        let when_value = when_value.into_array(batch.num_rows())?;
+        // `num_rows == 1` is intentional to avoid expanding scalars.
+        // If the `when_value` is effectively a scalar, the 'all true' and 'all false' checks
+        // below will avoid incorrectly using the scalar as a merge/zip mask.
+        let when_value = when_value.into_array(1)?;
         let when_value = as_boolean_array(&when_value).map_err(|e| {
             DataFusionError::Context(
                 "WHEN expression did not return a BooleanArray".to_string(),
                 Box::new(e),
             )
         })?;
 
-        // For the true and false/null selection vectors, bypass filtering and merging
-        // results. This avoids materializing the array for the other branch which we will discard
-        // entirely anyway.
         let true_count = when_value.true_count();
-        if true_count == batch.num_rows() {
-            return self.body.when_then_expr[0].1.evaluate(batch);
+        if true_count == when_value.len() {
+            // All input rows are true, just call the 'then' expression
+            self.body.when_then_expr[0].1.evaluate(batch)
         } else if true_count == 0 {
-            return self.body.else_expr.as_ref().unwrap().evaluate(batch);
-        }
-
-        if projected.projection.len() < batch.num_columns() {
+            // All input rows are false/null, just call the 'else' expression
+            self.body.else_expr.as_ref().unwrap().evaluate(batch)
+        } else if projected.projection.len() < batch.num_columns() {
+            // The case expressions do not use all the columns of the input batch.
+            // Project first to reduce time spent filtering.
             let projected_batch = batch.project(&projected.projection)?;
-            projected
-                .body
-                .expr_or_expr(&projected_batch, &when_value, &return_type)
+            projected.body.expr_or_expr(&projected_batch, when_value)
         } else {
-            self.body.expr_or_expr(batch, &when_value, &return_type)
+            // All columns are used in the case expressions, so there is no need to project.
+            self.body.expr_or_expr(batch, when_value)
         }
     }
 }
@@ -1244,23 +1266,7 @@ impl PhysicalExpr for CaseExpr {
     }
 
     fn data_type(&self, input_schema: &Schema) -> Result<DataType> {
-        // since all then results have the same data type, we can choose any one as the
-        // return data type except for the null.
-        let mut data_type = DataType::Null;
-        for i in 0..self.body.when_then_expr.len() {
-            data_type = self.body.when_then_expr[i].1.data_type(input_schema)?;
-            if !data_type.equals_datatype(&DataType::Null) {
-                break;
-            }
-        }
-        // if all then results are null, we use data type of else expr instead if possible.
-        if data_type.equals_datatype(&DataType::Null) {
-            if let Some(e) = &self.body.else_expr {
-                data_type = e.data_type(input_schema)?;
-            }
-        }
-
-        Ok(data_type)
+        self.body.data_type(input_schema)
     }
 
     fn nullable(&self, input_schema: &Schema) -> Result<bool> {