More work reduction in case of early exit

Author: pepijnve

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

@@ -269,14 +269,14 @@ impl CaseExpr {
             )
         };
 
-        let (mut remainder_batch, mut base_value, mut current_value) = initial_values;
+        let (mut remainder, mut base_value, mut current_value) = initial_values;
 
         for i in 0..self.when_then_expr.len() {
             // Evaluate the 'when' predicate for the remainder batch
             // This results in a boolean array with the same length as the remaining number of rows
             let when_expression = &self.when_then_expr[i].0;
-            let when_value = when_expression.evaluate(&remainder_batch)?;
-            let when_value = when_value.to_array(remainder_batch.num_rows())?;
+            let when_value = when_expression.evaluate(&remainder)?;
+            let when_value = when_value.to_array(remainder.num_rows())?;
             let when_value = compare_with_eq(
                 &base_value,
                 &when_value,
@@ -294,43 +294,45 @@ impl CaseExpr {
                 continue;
             }
 
-            // Make sure 'NULL' is treated as false
-            let when_value = match when_value.null_count() {
-                0 => Cow::Borrowed(when_value),
-                _ => Cow::Owned(prep_null_mask_filter(when_value)),
-            };
-
-            // Filter the remainder batch based on the 'when' value
-            // This results in a batch containing only the rows that need to be evaluated
-            // for the current branch
-            let then_batch = filter_record_batch(&remainder_batch, &when_value)?;
-
             // Evaluate the then expression for the matching rows
             let then_expression = &self.when_then_expr[i].1;
-            let then_value = then_expression.evaluate(&then_batch)?;
 
-            // Expand the 'when' match array using the 'remainder scatter' array
-            // This results in a truthy boolean array than we can use to merge the
-            // 'then' values with the `current_value` array.
-            let then_merge = scatter(&remainder_scatter, when_value.as_ref())?;
-            let then_merge = then_merge.as_boolean();
+            if when_match_count == remainder.num_rows() {
+                // If the 'when' predicate matched all remaining row, there's nothing left to do so
+                // we can return early
+                let then_value = then_expression.evaluate(remainder.as_ref())?;
+                current_value =
+                    merge_result(&current_value, then_value, &remainder_scatter)?;
+                return Ok(ColumnarValue::Array(current_value));
+            } else {
+                // Make sure 'NULL' is treated as false
+                let when_value = match when_value.null_count() {
+                    0 => Cow::Borrowed(when_value),
+                    _ => Cow::Owned(prep_null_mask_filter(when_value)),
+                };
+
+                // Filter the remainder batch based on the 'when' value
+                // This results in a batch containing only the rows that need to be evaluated
+                // for the current branch
+                let then_batch = filter_record_batch(&remainder, &when_value)?;
+                let then_value = then_expression.evaluate(&then_batch)?;
 
-            // Merge the 'then' values with the `current_value` array
-            current_value = merge_result(&current_value, then_value, then_merge)?;
+                // Expand the 'when' match array using the 'remainder scatter' array
+                // This results in a truthy boolean array than we can use to merge the
+                // 'then' values with the `current_value` array.
+                let then_merge = scatter(&remainder_scatter, when_value.as_ref())?;
+                let then_merge = then_merge.as_boolean();
 
-            // If the 'when' predicate matched all remaining row, there's nothing left to do so
-            // we can return early
-            if remainder_batch.num_rows() == when_match_count {
-                return Ok(ColumnarValue::Array(current_value));
-            }
+                // Merge the 'then' values with the `current_value` array
+                current_value = merge_result(&current_value, then_value, then_merge)?;
 
-            // Clear the positions in 'remainder scatter' for which we just evaluated a value
-            remainder_scatter = and_not(&remainder_scatter, then_merge)?;
+                // Clear the positions in 'remainder scatter' for which we just evaluated a value
+                remainder_scatter = and_not(&remainder_scatter, then_merge)?;
+            };
 
             // Finally, prepare the remainder batch for the next branch
-            let next_selection = not(&when_value)?;
-            remainder_batch =
-                Cow::Owned(filter_record_batch(&remainder_batch, &next_selection)?);
+            let next_selection = not(when_value)?;
+            remainder = Cow::Owned(filter_record_batch(&remainder, &next_selection)?);
             base_value = filter(base_value.as_ref(), &next_selection)?;
         }
 
@@ -339,7 +341,7 @@ impl CaseExpr {
         if let Some(e) = self.else_expr() {
             // keep `else_expr`'s data type and return type consistent
             let expr = try_cast(Arc::clone(e), &batch.schema(), return_type.clone())?;
-            let else_value = expr.evaluate(&remainder_batch)?;
+            let else_value = expr.evaluate(&remainder)?;
             current_value = merge_result(&current_value, else_value, &remainder_scatter)?;
         }
 
@@ -360,14 +362,14 @@ impl CaseExpr {
         let mut current_value = new_null_array(&return_type, batch.num_rows());
 
         let mut remainder_scatter = BooleanArray::from(vec![true; batch.num_rows()]);
-        let mut remainder_batch = Cow::Borrowed(batch);
+        let mut remainder = Cow::Borrowed(batch);
 
         for i in 0..self.when_then_expr.len() {
             // Evaluate the 'when' predicate for the remainder batch
             // This results in a boolean array with the same length as the remaining number of rows
             let when_predicate = &self.when_then_expr[i].0;
-            let when_value = when_predicate.evaluate(&remainder_batch)?;
-            let when_value = when_value.to_array(remainder_batch.num_rows())?;
+            let when_value = when_predicate.evaluate(&remainder)?;
+            let when_value = when_value.to_array(remainder.num_rows())?;
             let when_value = as_boolean_array(&when_value).map_err(|_| {
                 internal_datafusion_err!("WHEN expression did not return a BooleanArray")
             })?;
@@ -378,51 +380,53 @@ impl CaseExpr {
                 continue;
             }
 
-            // Make sure 'NULL' is treated as false
-            let when_value = match when_value.null_count() {
-                0 => Cow::Borrowed(when_value),
-                _ => Cow::Owned(prep_null_mask_filter(when_value)),
-            };
-
-            // Filter the remainder batch based on the 'when' value
-            // This results in a batch containing only the rows that need to be evaluated
-            // for the current branch
-            let then_batch = filter_record_batch(&remainder_batch, &when_value)?;
-
             // Evaluate the then expression for the matching rows
             let then_expression = &self.when_then_expr[i].1;
-            let then_value = then_expression.evaluate(&then_batch)?;
 
-            // Expand the 'when' match array using the 'remainder scatter' array
-            // This results in a truthy boolean array than we can use to merge the
-            // 'then' values with the `current_value` array.
-            let then_merge = scatter(&remainder_scatter, when_value.as_ref())?;
-            let then_merge = then_merge.as_boolean();
+            if when_match_count == remainder.num_rows() {
+                // If the 'when' predicate matched all remaining row, there's nothing left to do so
+                // we can return early
+                let then_value = then_expression.evaluate(remainder.as_ref())?;
+                current_value =
+                    merge_result(&current_value, then_value, &remainder_scatter)?;
+                return Ok(ColumnarValue::Array(current_value));
+            } else {
+                // Make sure 'NULL' is treated as false
+                let when_value = match when_value.null_count() {
+                    0 => Cow::Borrowed(when_value),
+                    _ => Cow::Owned(prep_null_mask_filter(when_value)),
+                };
+
+                // Filter the remainder batch based on the 'when' value
+                // This results in a batch containing only the rows that need to be evaluated
+                // for the current branch
+                let then_batch = filter_record_batch(&remainder, &when_value)?;
+                let then_value = then_expression.evaluate(&then_batch)?;
 
-            // Merge the 'then' values with the `current_value` array
-            current_value = merge_result(&current_value, then_value, then_merge)?;
+                // Expand the 'when' match array using the 'remainder scatter' array
+                // This results in a truthy boolean array than we can use to merge the
+                // 'then' values with the `current_value` array.
+                let then_merge = scatter(&remainder_scatter, when_value.as_ref())?;
+                let then_merge = then_merge.as_boolean();
 
-            // If the 'when' predicate matched all remaining row, there's nothing left to do so
-            // we can return early
-            if remainder_batch.num_rows() == when_match_count {
-                return Ok(ColumnarValue::Array(current_value));
-            }
+                // Merge the 'then' values with the `current_value` array
+                current_value = merge_result(&current_value, then_value, then_merge)?;
 
-            // Clear the positions in 'remainder scatter' for which we just evaluated a value
-            remainder_scatter = and_not(&remainder_scatter, then_merge)?;
+                // Clear the positions in 'remainder scatter' for which we just evaluated a value
+                remainder_scatter = and_not(&remainder_scatter, then_merge)?;
+            };
 
             // Finally, prepare the remainder batch for the next branch
-            let next_selection = not(&when_value)?;
-            remainder_batch =
-                Cow::Owned(filter_record_batch(&remainder_batch, &next_selection)?);
+            let next_selection = not(when_value)?;
+            remainder = Cow::Owned(filter_record_batch(&remainder, &next_selection)?);
         }
 
         // If we reached this point, some rows were left unmatched.
         // Check if those need to be evaluated using the 'else' expression.
         if let Some(e) = self.else_expr() {
             // keep `else_expr`'s data type and return type consistent
             let expr = try_cast(Arc::clone(e), &batch.schema(), return_type.clone())?;
-            let else_value = expr.evaluate(&remainder_batch)?;
+            let else_value = expr.evaluate(&remainder)?;
             current_value = merge_result(&current_value, else_value, &remainder_scatter)?;
         }