[SPARK-35564][SQL] Support subexpression elimination for conditionally evaluated expressions

[Kimahriman]

What changes were proposed in this pull request?

I am proposing to add support for conditionally evaluated expressions during subexpression elimination. Currently, only expressions that will definitely be always at least twice are candidates for subexpression elimination. This PR updates that logic so that expressions that are always evaluated at least once and conditionally evaluated at least once are also candidates for subexpression elimination. This helps optimize a common case during data normalization and cleaning and want to null out values that don't match a certain pattern, where you have something like:

transformed = F.regexp_replace(F.lower(F.trim('my_column')))
df.withColumn('normalized_value', F.when(F.length(transformed) > 0, transformed))

or

df.withColumn('normalized_value', F.when(transformed.rlike(<some regex>), transformed))

In these cases, transformed will always be fully calculated twice, because it might only be needed once. I am proposing creating a subexpression for transformed in this case.

In practice I've seen a decrease in runtime and codegen size of 10-30% in our production pipelines that heavily make use of this type of logic.

The only potential downside is creating extra subexpressions, and therefore function calls, more than necessary. This should only be an issue for certain edge cases where your conditional overwhelming evaluates to false. And then the only overhead is running your conditional logic potentially in a separate function rather than inlined in the codegen. I added a config to control this behavior if that is actually a real concern to anyone, but I'd be happy to just remove the config.

I also updated some of the existing logic for common expressions in coalesce and when that are actually better handled by the new logic, since you are only guaranteed to have the first value of a Coalesce evaluated, as well as the first conditional of a CaseWhen expression.

Why are the changes needed?

To increase the performance of conditional expressions.

Does this PR introduce any user-facing change?

No, just performance improvements.

How was this patch tested?

New and updated UT.

[cloud-fan]

OK to test

[SparkQA]

Kubernetes integration test starting
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/44671/

[SparkQA]

Kubernetes integration test status failure
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/44671/

[viirya]

The only potential downside is creating extra subexpressions, and therefore function calls, more than necessary. This should only be an issue for certain edge cases where your conditional overwhelming evaluates to false. And then the only overhead is running your conditional logic potentially in a separate function rather than inlined in the codegen. I added a config to control this behavior if that is actually a real concern to anyone, but I'd be happy to just remove the config.

I don't think the downside is edge case. On the contrary, I rather think it is common case than the use-case proposed here.

After this, any common expression shared between conditionally evaluated expression and a normal expression will be subexpression. I have a concern that gen-ed code will be overwhelmed with such subexpressions.

At least we need a config for this and I don't think it should be enabled by default.

[Kimahriman]

After this, any common expression shared between conditionally evaluated expression and a normal expression will be subexpression. I have a concern that gen-ed code will be overwhelmed with such subexpressions.

What exactly is the overwhelming part? I figured smaller overall code size would be beneficial.

[viirya]

What exactly is the overwhelming part? I figured smaller overall code size would be beneficial.

It is not zero-cost. For example, too many subexpressions will possibly make non-split case to be split case.

[Kimahriman]

Could you elaborate on how that could happen? I don't know that much about the codegen process

[viirya]

Could you elaborate on how that could happen? I don't know that much about the codegen process

In short, during subexpressions codegen, if the total code length is more than a threshold, we choose to split it as functions to avoid reach the max size of a method.

[Kimahriman]

Oh you're specifically talking about the subexpressions being split into functions versus inlined, not the general splitting the whole codegen into functions?

[SparkQA]

Test build #140145 has finished for PR 32987 at commit d4d64c7.

• This patch passes all tests.
• This patch merges cleanly.
• This patch adds no public classes.

[Kimahriman]

My main assumption in creating this was that it's always faster to run an expression once in a function than twice inlined. If this creates a lot of extra subexpressions that pushes the code over the 1kb threshold for breaking into functions, then the alternative is that you are running a lot of duplicate inlined logic, so at the end of the day it all comes down to how often a subexpression created by this logic is only evaluated once.

The two extremes of performance impact I can think of would be:

• Worst case: Without this logic, you have subexpressions that are just small enough to remain inlined. You add one conditional that creates a new subexpression that pushes your code over the (default) 1kb limit. That conditional never evaluates to true, so your conditional subexpression is evaluated once in a function rather than inlined, and all your other subexpressions are evaluated with a function call instead of inlined as well. This is somewhat bound by the number of subexpressions that can be fit inline in the first place, plus the function calls of the one-time evaluated conditional subexpressions.
• Best case: Your existing subexpressions have already been broken out into functions before this change, or the new subexpression fits inline as well, and the conditional always evaluates to true, so you are running the conditional expression once instead of two or more times. This is essentially the existing logic where we create a subexpression for things that are always evaluated at least twice, so obviously a win here.

Realistically things are going to fall somewhere in the middle. Where the extra function calls outweigh the deduped expression execution, who knows. But the upside here is pretty large, and I would expect most Spark users would expect this to logically happen (don't run the same code twice). If we want to leave it with the setting defaulted to disabled I'm fine with that.

[cloud-fan]

I feel it's a bit more overcomplicated: now we have childen, commonChildren, conditionallyEvaluatedChildren.

[Kimahriman]

Yeah it's just all different cases that need to be handled. I can think about how to simplify or if #33142 would help simplify

[Kimahriman]

i.e. useCount and conditionalUseCount instead of separate map and all that or something, idk

[Kimahriman]

Let me know if it's less overcomplicated now...

[Kimahriman]

Updated based on the refactor. It's still a little rough and needs some cleaning, renaming things, and updating a lot of comments, but wanted to get initial feedback

[SparkQA]

Kubernetes integration test starting
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/45267/

[SparkQA]

Kubernetes integration test status success
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/45267/

[SparkQA]

Test build #140756 has finished for PR 32987 at commit 1111b9b.

• This patch passes all tests.
• This patch merges cleanly.
• This patch adds the following public classes (experimental):
• case class ExpressionStats(expr: Expression)(
• case class RecurseChildren(

[SparkQA]

Test build #141494 has started for PR 32987 at commit d956d22.

[SparkQA]

Test build #141497 has started for PR 32987 at commit 5d6e1ad.

[SparkQA]

Test build #141498 has started for PR 32987 at commit 80245a6.

[Kimahriman]

The not-first conditions are now handled as a conditional instead. Supports all the same existing behavior but additionally can create subexpressions for things only in one of the remaining conditions instead of all. For example, CaseWhen((a + b) / (c + d) > 1, 1, a + b > 1, 2, c + d > 1, 3), a + b and c + d will become subexpressions now where they wouldn't previously, though only with this config enabled if we need to keep the config

[Kimahriman]

Same as the CaseWhen above

[Kimahriman]

#33142 (comment) in other cases it's already accepted that the performance overhead of maybe only using a subexpression once is worth the trade-off of not having to potentially evaluate it twice, so this just expands the places that could happen. Personally I don't think it needs a config defaulting to turned off, but I'm fine leaving it in if necessary. It does effectively prevent all the existing cases of creating a subexpression for an expression that might only be evaluated once, like mentioned in the comment, if the config is turned off.

[SparkQA]

Kubernetes integration test starting
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/46014/

[cloud-fan]

shall we update the doc of this method? no equivalenceMap in this method now.

[Kimahriman]

Yep good call, haven't kept up with with some of the docs

[cloud-fan]

do we have an example this commonChildren?

[Kimahriman]

Updated the docs a little bit to clarify. Currently it's only If and CaseWhen expressions that commonChildren applies too, should I put one of those as an example in the doc?