Blog post about CASE optimization

content/blog/2025-11-11-datafusion_case.md

@@ -0,0 +1,377 @@
+---
+layout: post
+title: Optimizing CASE Expression Evaluation
+date: 2025-11-11
+author: Pepijn Van Eeckhoudt
+categories: [features]
+---
+<!--
+{% comment %}
+Licensed to the Apache Software Foundation (ASF) under one or more
+contributor license agreements.  See the NOTICE file distributed with
+this work for additional information regarding copyright ownership.
+The ASF licenses this file to you under the Apache License, Version 2.0
+(the "License"); you may not use this file except in compliance with
+the License.  You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+{% endcomment %}
+-->
+
+[TOC]
+
+<style>
+figure {
+  margin: 20px 0;
+}
+
+figure img {
+  display: block;
+  max-width: 80%;
+  margin: auto;
+}
+
+figcaption {
+  font-style: italic;
+  color: #555;
+  font-size: 0.9em;
+  max-width: 80%;
+  margin: auto;
+  text-align: center;
+}
+</style>
+
+# Optimizing CASE Expression Evaluation in DataFusion
+
+SQL's `CASE` expression is one of the few constructs the language provides to perform conditional logic.
+Its deceptively simple syntax hides significant implementation complexity.
+Over the past few weeks, we've landed a series of improvements to DataFusion's `CASE` expression evaluator that reduce both CPU time and memory allocations.
+This post walks through the original implementation, its performance bottlenecks, and how we addressed them step by step.
+Finally we'll also take a look at some future improvements to `CASE` that are in the works.
+
+## Background: CASE Expression Evaluation
+
+SQL supports two forms of CASE expressions:
+
+1. **Simple**: `CASE expr WHEN value1 THEN result1 WHEN value2 THEN result2 ... END`
+2. **Searched**: `CASE WHEN condition1 THEN result1 WHEN condition2 THEN result2 ... END`
+
+The simple form evaluates an expression once for each input row and then tests that value against the constants (or expressions) in each `WHEN` clause using equality comparisons.
+Think of it as a limited Rust `match` expression.
+
+Here's a simple example:
+
+```sql
+CASE status
+    WHEN 'pending' THEN 1
+    WHEN 'active' THEN 2
+    WHEN 'complete' THEN 3
+    ELSE 0
+END
+```
+
+In this `CASE` expression, `status` is evaluated once per row, and then its value is tested for equality with the values `'pending'`, `'active'`, and `'complete'` in that order.
+The `THEN` expression value for the first matching `WHEN` expression is returned per row.
+
+The searched `CASE` form is a more flexible variant.
+It evaluates completely independent boolean expressions for each branch.
+
+This allows you to test different columns with different operators per branch as can be seen in the following example:
+
+```sql
+CASE
+    WHEN age > 65 THEN 'senior'
+    WHEN childCount != 0 THEN 'parent'
+    WHEN age < 21 THEN 'minor'
+    ELSE 'adult'
+END
+```
+
+In both forms, branches are evaluated sequentially with short-circuit semantics: for each row, once a `WHEN` condition matches, the corresponding `THEN` expression is evaluated. Any further branches are not evaluated for that row.
+This lazy evaluation model is critical for correctness.
+It let's you safely write `CASE` expressions like
+
+```sql
+CASE
+    WHEN denominator == 0 THEN NULL
+    ELSE nominator / denominator
+END
+```
+
+that are guaranteed to not trigger divide-by-zero errors.
+
+## `CASE` Evaluation in DataFusion 50.0.0
+
+For the rest of this post we'll be looking at 'searched case' evaluation.
+'Simple case' uses a distinct, but very similar implementation.
+The same set of improvements has been applied to both.
+
+The baseline implementation in DataFusion 50.0.0 evaluated `CASE` using a straightforward approach:
+
+1. Start with an output array `out` with the same length as the input batch, filled with nulls. Additionally, create a bit vector `remainder` with the same length and each value set to `true`.
+2. For each `WHEN`/`THEN` branch:
+  - Evaluate the `WHEN` condition for remaining unmatched rows using `PhysicalExpr::evaluate_selection`, passing in the input batch and the `remainder` mask
+  - If any rows matched, evaluate the `THEN` expression for those rows using `PhysicalExpr::evaluate_selection`
+  - Merge the results into the `out` using the `zip` kernel
+  - Update the `remainder` mask to exclude matched rows
+3. If there's an `ELSE` clause, evaluate it for any remaining unmatched rows and merge using `zip`
+
+Here's a simplified version of the original loop:
+
+```rust
+let mut current_value = new_null_array(&return_type, batch.num_rows());
+let mut remainder = BooleanArray::from(vec![true; batch.num_rows()]);
+
+for (when_expr, then_expr) in &self.when_then_expr {
+    let when_value = when_expr.evaluate_selection(batch, &remainder)?
+        .into_array(batch.num_rows())?;
+    let when_value = and(&when_value, &remainder)?;
+
+    if when_value.true_count() == 0 {
+        continue;
+    }
+
+    let then_value = then_expr.evaluate_selection(batch, &when_value)?;
+    current_value = zip(&when_value, &then_value, &current_value)?;
+    remainder = and_not(&remainder, &when_value)?;
+}
+```
+
+While correct, this implementation had several performance issues mostly related to the usage of `evaluate_selection`.
+To understand why, we need to dig a little deeper into the implementation of that function.
+Here's a simplified version of it that captures the relevant parts:
+
+```rust
+pub trait PhysicalExpr {
+    fn evaluate_selection(
+        &self,
+        batch: &RecordBatch,
+        selection: &BooleanArray,
+    ) -> Result<ColumnarValue> {
+        let filtered_batch = filter_record_batch(batch, selection)?;
+        let filtered_result = self.evaluate(&filtered_batch)?;
+        scatter(selection, filtered_result)
+    }
+}
+```
+
+The `evaluate_selection` method first filters the input batch to only include rows that match the `selection` mask.
+It then calls the regular `evaluate` method using the filtered batch as input.
+Finally, to return a result array with the same number of rows as `batch`, the `scatter` function is called.
+This function produces a new array padded with `null` values for any rows that didn't match the `selection` mask.
+
+So how does the simple evaluation strategy and use of `evaluate_selection` cause performance overhead?
+
+### Problem 1: No Early Exit
+
+The case evaluation loop always iterated through all branches, even when every row had already been matched.
+In queries where early branches match many rows, this meant unnecessary work was done for remaining rows.
+
+### Problem 2: Repeated Filtering, Scattering, and Merging
+
+Each iteration performed operations that are very well-optimized but still not cost free to execute:
+- **Filtering**: `PhysicalExpr::evaluate_selection` filters the entire `RecordBatch` (all columns) for each branch. For the `WHEN` expression, this was done even if the selection mask was entirely empty.
+- **Scattering**: `PhysicalExpr::evaluate_selection` scatters the filtered result back to the original `RecordBatch` length.
+- **Merging**: The `zip` kernel is called once per branch to merge partial results into the output array
+
+Each of these steps allocates new arrays and shuffles a lot of data around. 
+
+### Problem 3: Filtering Unused Columns
+
+The `PhysicalExpr::evaluate_selection` method filters the entire record batch, including columns that the current branch's `WHEN` and `THEN` expressions don't reference.
+For wide tables (many columns) with narrow expressions (few column references), this is wasteful.
+
+Suppose we have a table with 26 columns named `a` through `z`.
+For a simple CASE expression like:
+
+```sql
+CASE
+  WHEN a > 1000 THEN 'large'
+  WHEN a >= 0 THEN 'positive'
+  ELSE 'negative'
+END
+```
+
+the implementation would filter all columns 26 columns even though only a single column is needed for the entire `CASE` expression evaluation.
+Again this involves a non-negligible amount of allocation and data copying.
+
+## Performance Improvements
+
+### Optimization 1: Short-Circuit Early Exit (commit 7c215ed)
+
+The first optimization added early exit logic to the evaluation loop:
+
+```rust
+let mut remainder_count = batch.num_rows();
+
+for (when_expr, then_expr) in &self.when_then_expr {
+    if remainder_count == 0 {
+        break;  // All rows matched, exit early
+    }
+
+    // ... evaluate branch ...
+
+    let when_match_count = when_value.true_count();
+    remainder_count -= when_match_count;
+}
+```
+
+Additionally, we avoid evaluating the `ELSE` clause when no rows remain:
+
+```rust
+if let Some(else_expr) = &self.else_expr {
+    remainder = or(&base_nulls, &remainder)?;
+    if remainder.true_count() > 0 {
+        // ... evaluate else ...
+    }
+}
+```
+
+**Impact**: For queries where early branches match all rows, this eliminates unnecessary branch evaluations and `ELSE` clause processing.
+
+### Optimization 2: Optimized Result Merging (commit e9431fc)
+
+The second optimization fundamentally restructured how partial results are merged.
+Instead of using `zip()` after each branch to merge results into an output array, we now:
+
+1. Maintain the subset of rows still needing evaluation across loop iterations
+2. Filter the batch progressively as rows are matched
+3. Build an index structure that tracks which branch produced each row's result
+4. Perform a single merge operation at the end
+
+The key insight is that we can defer all merging until the end of the evaluation loop by tracking result provenance.
+When a branch matches a number of rows, instead of immediately merging with `zip()`, we:
+1. Store the partial result array
+2. Mark the cells corresponding to each row in an indices array as needing to take one value from the partial result array
+
+In the example below, three `WHEN/THEN` branches produced results.
+The first branch produced the result `A` for 2, the second produced `B` for row 1, and the third produced `C` and `D` for rows 4 and 5.
+The final result array is obtained by running the arrays through the merge operation.
+
+```aiignore
+┌───────────┐  ┌─────────┐                             ┌─────────┐
+│┌─────────┐│  │   None  │                             │   NULL  │
+││    A    ││  ├─────────┤                             ├─────────┤
+│└─────────┘│  │    1    │                             │    B    │
+│┌─────────┐│  ├─────────┤                             ├─────────┤
+││    B    ││  │    0    │   merge_n(values, indices)  │    A    │
+│└─────────┘│  ├─────────┤  ─────────────────────────▶ ├─────────┤
+│┌─────────┐│  │   None  │                             │   NULL  │
+││    C    ││  ├─────────┤                             ├─────────┤
+│├─────────┤│  │    2    │                             │    C    │
+││    D    ││  ├─────────┤                             ├─────────┤
+│└─────────┘│  │    2    │                             │    D    │
+└───────────┘  └─────────┘                             └─────────┘
+   arrays        indices                                  result
+```
+
+The main benefits of this merge operation are that the `scatter` step is eleminated enitrely, and instead of requiring
+a `zip` per branch only a single `merge_n` is done.
+
+Besides more efficient merging, we also maintain a progressively filtered `remainder_batch`.
+
+```rust
+let mut remainder_batch = Cow::Borrowed(batch);
+let mut remainder_rows = UInt32Array::from_iter_values(0..batch.num_rows());
+
+for branch in branches {
+    // Evaluate on progressively smaller batch
+    let when_value = evaluate_on(&remainder_batch)?;
+
+    // Filter remainder_batch for next iteration
+    let next_filter = create_filter(&not(&when_value)?);
+    remainder_batch = Cow::Owned(filter_record_batch(&remainder_batch, &next_filter)?);
+    remainder_rows = filter_array(&remainder_rows, &next_filter)?;
+}
+```
+
+As rows are matched, the `remainder_batch` shrinks, making subsequent filter operations a bit faster.
+
+**Impact**: This eliminates N-1 merge operations and makes filtering progressively cheaper as branches match rows.
+
+This new operation was initially developed specifically for DataFusion's `CASE` evaluation, but in the meantime has been generalized and moved into the `arrow-rs` crate as [`arrow_select::merge::merge_n`](https://docs.rs/arrow-select/57.1.0/arrow_select/merge/fn.merge_n.html).
+
+### Optimization 3: Column Projection (commit c3e49fb)
+
+The third optimization addresses the "filtering unused columns" problem through projection.
+Before evaluating a CASE expression, we:
+
+1. Analyze all WHEN/THEN/ELSE expressions to find referenced columns
+2. Build a projection vector containing only those column indices
+3. Derive new versions of the expressions with updated column references
+
+For example, if the original CASE references columns at indices `[1, 5, 8]` in a 20-column batch:
+- Project the batch to a 3-column batch with those columns
+- Rewrite expressions from `col@1, col@5, col@8` to `col@0, col@1, col@2`
+- Evaluate using the projected batch
+
+This is encapsulated in a `ProjectedCaseBody` structure:
+
+```rust
+struct ProjectedCaseBody {
+    projection: Vec<usize>,     // [1, 5, 8]
+    body: CaseBody,             // Expressions with rewritten column indices
+}
+```
+
+The projection logic is only applied when it would be beneficial (i.e., when the number of used columns is lower than the total number of columns in the batch).
+
+**Impact**: For wide tables with narrow CASE expressions, this dramatically reduces filtering overhead by removing copying of unused columns.
+
+### Optimization 4: Eliminating Scatter in Two-Branch Case (commit 32d2618)
+
+The final optimization targets a common pattern: `CASE WHEN condition THEN expr1 ELSE expr2 END`.
+
+Previously, this used the specialized `ExpressionOrExpression` fast path, but still used `evaluate_selection()` which produces scattered results:
+
+```rust
+// Old approach: produces scattered array with batch.num_rows() elements
+let then_value = then_expr.evaluate_selection(batch, &when_value)?
+    .into_array(batch.num_rows())?;
+```
+
+The new approach filters the batch first, then evaluates:
+
+```rust
+// New approach: filter to only matching rows
+let when_filter = create_filter(&when_value);
+let then_batch = filter_record_batch(batch, &when_filter)?;
+let then_value = then_expr.evaluate(&then_batch)?;  // Produces compact array
+
+let else_filter = create_filter(&not(&when_value)?);
+let else_batch = filter_record_batch(batch, &else_filter)?;
+let else_value = else_expr.evaluate(&else_batch)?;
+```
+
+This produces two compact arrays (one for THEN values, one for ELSE values) which are then merged with a custom merge function that doesn't require pre-scattered inputs:
+
+```rust
+fn merge(mask: &BooleanArray, truthy: ColumnarValue, falsy: ColumnarValue) -> ArrayRef {
+    // Interleave truthy and falsy values directly using mask
+    // without requiring pre-alignment
+}
+```
+
+**Impact**: This eliminates unnecessary scatter operations and memory allocations for one of the most common CASE expression patterns.
+
+Just like `merge_n` this operation has been moved into `arrow-rs` as [`arrow_select::merge::merge`](https://docs.rs/arrow-select/57.1.0/arrow_select/merge/fn.merge.html).
+
+## Summary
+
+Through four targeted optimizations, we've transformed CASE expression evaluation from a simple but inefficient implementation to a highly optimized one that:
+
+1. **Exits early** when all rows are matched
+2. **Defers merging** until the end with a single interleave operation
+3. **Projects columns** to avoid filtering unused data
+4. **Eliminates scatter** operations in common patterns
+
+These improvements compound: a CASE expression on a wide table with multiple branches and early matches benefits from all four optimizations simultaneously.
+The result is significantly reduced CPU time and memory allocation in one of SQL's most frequently used constructs.
+
+All of these changes maintain full compatibility with DataFusion's existing semantics and are covered by the existing test suite, demonstrating that performance optimization doesn't have to compromise correctness or require API changes.