feat: Optimize CASE expression for "column or null" use case #11534

andygrove · 2024-07-18T14:43:03Z

Which issue does this PR close?

Rationale for this change

Some TPC-DS queries include expressions such as case when (d_day_name='Sunday') then ss_sales_price else null end. The current implementation of CaseExpr is very expensive for an expression which is basically IF(expr, column, null). For this case we can take a fast path of just changing the null bitmask on the column.

case_when: column or null
                        time:   [842.29 ns 844.91 ns 847.49 ns]
                        change: [-92.101% -92.074% -92.046%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 10 outliers among 100 measurements (10.00%)
  1 (1.00%) low severe
  7 (7.00%) low mild
  2 (2.00%) high mild

What changes are included in this PR?

The existing evaluate method already had conditional logic for choosing between two execution methods based on context. This PR adds a third that specializes for the "column or null" use case.

The benchmarks were also improved.

Are these changes tested?

New test added for this optimization.

Are there any user-facing changes?

No, just faster queries in some cases.

viirya · 2024-07-18T14:55:22Z

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

+    /// END
+    CaseWithExpression,
+    /// This is a specialization for a specific use case where we can take a fast path
+    /// CASE WHEN condition THEN column [ELSE NULL] END


Actually I'm wondering why this special format of case when is not optimized to if/else during query optimization.

I don't think DataFusion has an if/else expression. We have one in Comet that could be upstreamed.

viirya · 2024-07-18T14:57:47Z

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

+                EvalMethod::CaseWithExpression
+            } else if when_then_expr.len() == 1
+                && when_then_expr[0].1.as_any().is::<Column>()
+                && else_expr.is_none()


If else value is a null literal, is else_expr still none? Or is it a null literal?

From SQL planning, it will be None, but it makes sense to add code here to normalize this for other use cases (other query engines that delegate to DataFusion). I have added this.

viirya · 2024-07-18T15:00:48Z

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

+    fn case_column_or_null(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
+        let when_expr = &self.when_then_expr[0].0;
+        let then_expr = &self.when_then_expr[0].1;
+        if let ColumnarValue::Array(bit_mask) = when_expr.evaluate(batch)? {


If it is evaluated to Scalar, maybe we still can work on it? I.e., we can convert the scalar value to an boolean array.

I am planning on adding a specialization for scalar values as well to avoid converting scalars to arrays

actually, nm, your question is different. Yes, we could implement special handling for CASE WHEN [true|false] THEN but I'm not sure that is a real-world use case that is worth optimizing

Oh, I forgot that in this optimization, the when expression could only be a Column.

alamb

Thanks @andygrove -- I think the performance improvement is pretty compelling and this PR does what it says.

I think it would be great to continue to optimize the CASE evaluation in follow on PRs. I bet if you filed some tickets others in the community would likely pitch in to help

alamb · 2024-07-18T20:35:02Z

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

+                ColumnarValue::Array(array) => {
+                    Ok(ColumnarValue::Array(nullif(&array, &bit_mask)?))
+                }
+                ColumnarValue::Scalar(_) => Err(DataFusionError::Execution(


this might be better as an internal error

You could also use the macros like

internal_err!("expression did not evaluate to an array")

alamb · 2024-07-18T20:39:34Z

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

        if when_then_expr.is_empty() {
            exec_err!("There must be at least one WHEN clause")
        } else {
+            let eval_method = if expr.is_some() {
+                EvalMethod::WithExpression
+            } else if when_then_expr.len() == 1


I was thinking about this optimization and I think it would be valid for any CASE expression that had a NULL ELSE, not just column

So in other words, I think you could remove the when_then_expr[0].1.as_any().is::<Column>() check and this would still work fine

Generally, it is only safe to use this approach for expressions that are infallible. One of the main reasons that regular CaseExpr is expensive is that we need to only evaluate the "true" expression on rows where the predicate has evaluated to true.

I do think that this could be extended beyond just Column expressions though.

I will add a comment in the code to explain this

I added more comments to explain the rationale and what is safe or not

alamb · 2024-07-18T20:42:11Z

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

+            }
+        } else {
+            Err(DataFusionError::Execution(
+                "predicate did not evaluate to an array".to_string(),


I agree with @viirya I think it would be fairly straightforward here to handle ColumnarValue::Scalar s well. I don't think we need to do it in this PR

Given you say in https://github.com/apache/datafusion/pull/11534/files#r1683081792

I am planning on adding a specialization for scalar values as well to avoid converting scalars to arrays

I think we could definitely do it in a follow on PR

alamb · 2024-07-18T20:43:27Z

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

@@ -998,6 +1080,53 @@ mod tests {
        Ok(())
    }

+    #[test]


I think it would be good to make sure we had a slt level test to cover this as well,

Maybe in
https://github.com/apache/datafusion/blob/382bf4f3c7a730828684b9e4ce01369b89717e19/datafusion/sqllogictest/test_files/expr.slt

Or we could start adding a file just for CASE if we are about to spend a bunch of time optimizing it 🤔

I willl add slt tests

I am planning on more CASE optimizations so will create a separate file

I added the slt test. This is my first time using slt and it is very cool

Dandandan · 2024-07-19T07:17:24Z

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

+                .downcast_ref::<BooleanArray>()
+                .expect("predicate should evaluate to a boolean array");
+            // invert the bitmask
+            let bit_mask = not(bit_mask)?;


As a further optimization I think the when predicate can be transformed to be not(condition) so it's possible to use bit_mask directly instead of not(bit_mask).
This makes it possible to optimize/simplify a condition like x=1 to x!=1 instead of not(x=1), saving the invert.

Thanks! I will look at this suggestion as a follow on. Great idea.

…11534)

Optimize CASE expression for column or null use case

f02dbc0

github-actions bot added the physical-expr Physical Expressions label Jul 18, 2024

cargo fmt

58be11a

github-actions bot added the core Core DataFusion crate label Jul 18, 2024

andygrove mentioned this pull request Jul 18, 2024

feat: Optimize for CASE WHEN .. THEN column ELSE null END apache/datafusion-comet#672

Closed

andygrove requested review from alamb and viirya July 18, 2024 14:49

viirya reviewed Jul 18, 2024

View reviewed changes

andygrove added 2 commits July 18, 2024 09:46

normalize nulls

0490765

clippy

5ed7c70

viirya approved these changes Jul 18, 2024

View reviewed changes

andygrove requested review from mingmwang, Dandandan and avantgardnerio July 18, 2024 16:42

alamb approved these changes Jul 18, 2024

View reviewed changes

Dandandan reviewed Jul 19, 2024

View reviewed changes

address feedback

c04235c

github-actions bot added the sqllogictest SQL Logic Tests (.slt) label Jul 19, 2024

andygrove merged commit 28fa74b into apache:main Jul 19, 2024
23 checks passed

andygrove added the performance Make DataFusion faster label Jul 19, 2024

Lordworms pushed a commit to Lordworms/arrow-datafusion that referenced this pull request Jul 23, 2024

feat: Optimize CASE expression for "column or null" use case (apache#…

dfcdd02

…11534)

wiedld pushed a commit to influxdata/arrow-datafusion that referenced this pull request Jul 31, 2024

feat: Optimize CASE expression for "column or null" use case (apache#…

a57a975

…11534)

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: Optimize CASE expression for "column or null" use case #11534

feat: Optimize CASE expression for "column or null" use case #11534

andygrove commented Jul 18, 2024 •

edited

Loading

viirya Jul 18, 2024

andygrove Jul 18, 2024

viirya Jul 18, 2024

andygrove Jul 18, 2024

viirya Jul 18, 2024

andygrove Jul 18, 2024

andygrove Jul 18, 2024

viirya Jul 18, 2024

alamb left a comment

alamb Jul 18, 2024

alamb Jul 18, 2024

andygrove Jul 18, 2024

andygrove Jul 18, 2024

andygrove Jul 19, 2024

alamb Jul 18, 2024

alamb Jul 18, 2024

andygrove Jul 18, 2024

andygrove Jul 18, 2024

andygrove Jul 19, 2024

Dandandan Jul 19, 2024 •

edited

Loading

andygrove Jul 19, 2024

@@ @@ -998,6 +1080,53 @@ mod tests { @@
                       Ok(())
                   }
+                  #[test]

feat: Optimize CASE expression for "column or null" use case #11534

feat: Optimize CASE expression for "column or null" use case #11534

Conversation

andygrove commented Jul 18, 2024 • edited Loading

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alamb left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Dandandan Jul 19, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

andygrove commented Jul 18, 2024 •

edited

Loading

Dandandan Jul 19, 2024 •

edited

Loading