Skip to content

[SPARK-33798][SQL] Add new rule to push down the foldable expressions through CaseWhen/If #30790

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wangyum wants to merge 12 commits into from
Closed

[SPARK-33798][SQL] Add new rule to push down the foldable expressions through CaseWhen/If #30790

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
19b0a83
Simplify EqualTo(CaseWhen/If, Literal) always false
wangyum Dec 16, 2020
831bf66
Address comments
wangyum Dec 16, 2020
859893d
Null value should be handled by NullPropagation.
wangyum Dec 16, 2020
db584ef
Handle Null values
wangyum Dec 16, 2020
55f4528
address comments
wangyum Dec 16, 2020
0a48048
address comments
wangyum Dec 17, 2020
f9f622f
Push down EqualTo through CaseWhen/If
wangyum Dec 17, 2020
21ef0c3
Push down EqualTo through CaseWhen/If
wangyum Dec 17, 2020
d5135ec
Fix
wangyum Dec 17, 2020
8ccc3c1
fix
wangyum Dec 17, 2020
45b56fc
fix
wangyum Dec 18, 2020
cfac0e8
fix
wangyum Dec 18, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
fix
  • Loading branch information
@wangyum
wangyum committed Dec 18, 2020
commit 45b56fcc7b08a89893082b7b9b6c6e2f8e4747f7
5 changes: 5 additions & 0 deletions sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Expression.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -636,6 +636,11 @@ abstract class BinaryExpression extends Expression {
}


object BinaryExpression {
def unapply(e: BinaryExpression): Option[(Expression, Expression)] = Some((e.left, e.right))
}


/**
* A [[BinaryExpression]] that is an operator, with two properties:
*
Expand Down
22 changes: 17 additions & 5 deletions sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/expressions.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ import scala.collection.immutable.HashSet
import scala.collection.mutable.{ArrayBuffer, Stack}

import org.apache.spark.sql.catalyst.analysis._
import org.apache.spark.sql.catalyst.expressions._
import org.apache.spark.sql.catalyst.expressions.{BinaryExpression, _}
import org.apache.spark.sql.catalyst.expressions.Literal.{FalseLiteral, TrueLiteral}
import org.apache.spark.sql.catalyst.expressions.aggregate._
import org.apache.spark.sql.catalyst.expressions.objects.AssertNotNull
Expand Down Expand Up @@ -534,17 +534,29 @@ object SimplifyConditionals extends Rule[LogicalPlan] with PredicateHelper {
object PushFoldableIntoBranches extends Rule[LogicalPlan] with PredicateHelper {
Copy link
Member

@dongjoon-hyun dongjoon-hyun Dec 18, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice idea.

def apply(plan: LogicalPlan): LogicalPlan = plan transform {
Copy link
Member

@maropu maropu Dec 17, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: we cannot use transformAllExpressions here?

Copy link
Member Author

@wangyum wangyum Dec 17, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This copied from

spark/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/expressions.scala

Line 474 in 8ccc3c1

case q: LogicalPlan => q transformExpressionsUp {

Copy link
Member

@maropu maropu Dec 18, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah, I see. I think its okay to simply use transformAllExpressions here.

case q: LogicalPlan => q transformExpressionsUp {
case b @ BinaryComparison(i @ If(_, trueValue, falseValue), right) if i.deterministic &&
right.foldable && trueValue.foldable && falseValue.foldable =>
case b @ BinaryExpression(i @ If(_, trueValue, falseValue), right) if i.deterministic &&
right.foldable && (trueValue.foldable || falseValue.foldable) =>
Copy link
Member

@HyukjinKwon HyukjinKwon Dec 18, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sorry if I missed something but can you explain why we need this foldable condition?

Shouldn't bin_op((if cond: a else: b), c) be same as if cond: bin_op(a, c) else: bin_op(b, c) as long as the results are deterministic?

Copy link
Member

@HyukjinKwon HyukjinKwon Dec 18, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ohh, okay but we just want to push down foldables only that can have benefits. gotya.

Copy link
Member

@viirya viirya Dec 18, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is good idea when both trueValue and falseValue are foldable like the example in the description.

But if only trueValue is foldable? For example, from If(_, 'b', col1) = 'a', we get If(_, 'b' = 'a', col1 = 'a'), are some benefits here too?

Copy link
Member Author

@wangyum wangyum Dec 18, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The differences in generated code:
image

Copy link
Member

@viirya viirya Dec 18, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hmm, I'm not sure if the difference means obvious benefits. Looks like it should be very close but I think I'm okay with it.

i.copy(
trueValue = b.makeCopy(Array(trueValue, right)),
falseValue = b.makeCopy(Array(falseValue, right)))

case b @ BinaryComparison(c @ CaseWhen(branches, elseValue), right) if c.deterministic &&
right.foldable && (branches.map(_._2) ++ elseValue).forall(_.foldable) =>
case b @ BinaryExpression(left, i @ If(_, trueValue, falseValue)) if i.deterministic &&
Copy link
Contributor

@cloud-fan cloud-fan Dec 18, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we care about deterministic here? This optimization doesn't change the execution order of the branch conditions and/or branch values.

left.foldable && (trueValue.foldable || falseValue.foldable) =>
Copy link
Contributor

@cloud-fan cloud-fan Dec 18, 2020 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

let's create a method to count foldables, with comments to explain why

// be conservative here: it's only a guaranteed win if all but at most only one branch
// end up being not foldable.
private def atMostOneUnfoldable(exprs: Seq[Expression]): Boolean = ...

then here left.foldable && atMostOneUnfoldable(Seq(trueValue, falseValue))

Copy link
Member Author

@wangyum wangyum Dec 18, 2020 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we need to push down this case, only one branch and non foldable:

assertEquivalent(EqualTo(CaseWhen(Seq((a, b)), None), Literal(1)), CaseWhen(Seq((a, EqualTo(b, Literal(1)))), None))

Copy link
Member Author

@wangyum wangyum Dec 18, 2020 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How about?

  // To be conservative here: it's only a guaranteed win if all but at most only one branch
  // end up being not foldable.
  private def atMostOneUnfoldable(exprs: Seq[Expression]): Boolean = {
    val (foldables, others) = exprs.partition(_.foldable)
    foldables.nonEmpty && others.length < 2
  }

i.copy(
trueValue = b.makeCopy(Array(left, trueValue)),
falseValue = b.makeCopy(Array(left, falseValue)))

case b @ BinaryExpression(c @ CaseWhen(branches, elseValue), right) if c.deterministic &&
right.foldable && (branches.map(_._2) ++ elseValue).exists(_.foldable) =>
c.copy(
branches.map(e => e.copy(_2 = b.makeCopy(Array(e._2, right)))),
elseValue.map(e => b.makeCopy(Array(e, right))))

case b @ BinaryExpression(left, c @ CaseWhen(branches, elseValue)) if c.deterministic &&
left.foldable && (branches.map(_._2) ++ elseValue).exists(_.foldable) =>
c.copy(
branches.map(e => e.copy(_2 = b.makeCopy(Array(left, e._2)))),
elseValue.map(e => b.makeCopy(Array(left, e))))
}
}
}
Expand Down
93 changes: 82 additions & 11 deletions ...rc/test/scala/org/apache/spark/sql/catalyst/optimizer/PushFoldableIntoBranchesSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,8 @@

package org.apache.spark.sql.catalyst.optimizer

import java.sql.Date

import org.apache.spark.sql.catalyst.analysis.UnresolvedAttribute
import org.apache.spark.sql.catalyst.dsl.expressions._
import org.apache.spark.sql.catalyst.dsl.plans._
Expand Down Expand Up @@ -49,20 +51,18 @@ class PushFoldableIntoBranchesSuite
comparePlans(actual, correctAnswer)
}

private val normalBranch = (NonFoldableLiteral(true), Literal(10))

test("SPARK-33798: Push down EqualTo through If") {
Copy link
Contributor

@cloud-fan cloud-fan Dec 18, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: when creating a new test suite, we don't need to add the JIRA ID as prefix.

assertEquivalent(EqualTo(ifExp, Literal(4)), FalseLiteral)
assertEquivalent(EqualTo(ifExp, Literal(3)), If(a, FalseLiteral, TrueLiteral))
Copy link
Contributor

@cloud-fan cloud-fan Dec 18, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If(a, FalseLiteral, TrueLiteral) can be turned into !a. We can optimize it in followups.

Copy link
Member Author

@wangyum wangyum Dec 19, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes. we can optimize it:
https://github.com/apache/spark/pull/30849/files#diff-d43484d56a4d9991066b5c00d12ec2465c75131e055fc02ee7fb6dfd45b5006fR478-R479

assertEquivalent(EqualTo(ifExp, Literal("4")), FalseLiteral)
Copy link
Contributor

@cloud-fan cloud-fan Dec 18, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

let's not test invalid expressions, EqualTo can't compare int and string.

assertEquivalent(EqualTo(ifExp, Literal("3")), If(a, FalseLiteral, TrueLiteral))

// Do not simplify if it contains non foldable expressions.
// Partially push down if it contains non foldable expressions.
assertEquivalent(
EqualTo(If(a, b, Literal(2)), Literal(2)),
EqualTo(If(a, b, Literal(2)), Literal(2)))
If(a, EqualTo(b, Literal(2)), TrueLiteral))

// Do not simplify if it contains non-deterministic expressions.
// Do not push down if it contains non-deterministic expressions.
val nonDeterministic = If(LessThan(Rand(1), Literal(0.5)), Literal(1), Literal(1))
assert(!nonDeterministic.deterministic)
assertEquivalent(EqualTo(nonDeterministic, Literal(-1)), EqualTo(nonDeterministic, Literal(-1)))
Expand Down Expand Up @@ -90,6 +90,29 @@ class PushFoldableIntoBranchesSuite
assertEquivalent(LessThanOrEqual(ifExp, Literal(4)), TrueLiteral)
}

test("SPARK-33798: Push down other BinaryOperator through If") {
assertEquivalent(Add(ifExp, Literal(4)), If(a, Literal(6), Literal(7)))
assertEquivalent(Subtract(ifExp, Literal(4)), If(a, Literal(-2), Literal(-1)))
assertEquivalent(Multiply(ifExp, Literal(4)), If(a, Literal(8), Literal(12)))
assertEquivalent(Pmod(ifExp, Literal(4)), If(a, Literal(2), Literal(3)))
assertEquivalent(Remainder(ifExp, Literal(4)), If(a, Literal(2), Literal(3)))
assertEquivalent(Divide(If(a, Literal(2.0), Literal(3.0)), Literal(1.0)),
If(a, Literal(2.0), Literal(3.0)))
assertEquivalent(And(If(a, FalseLiteral, TrueLiteral), TrueLiteral),
If(a, FalseLiteral, TrueLiteral))
assertEquivalent(Or(If(a, FalseLiteral, TrueLiteral), TrueLiteral), TrueLiteral)
}

test("SPARK-33798: Push down other BinaryExpression through If") {
assertEquivalent(BRound(If(a, Literal(1.23), Literal(1.24)), Literal(1)), Literal(1.2))
assertEquivalent(StartsWith(If(a, Literal("ab"), Literal("ac")), Literal("a")), TrueLiteral)
assertEquivalent(FindInSet(If(a, Literal("ab"), Literal("ac")), Literal("a")), Literal(0))
assertEquivalent(
AddMonths(If(a, Literal(Date.valueOf("2020-01-01")), Literal(Date.valueOf("2021-01-01"))),
Literal(1)),
If(a, Literal(Date.valueOf("2020-02-01")), Literal(Date.valueOf("2021-02-01"))))
}

test("SPARK-33798: Push down EqualTo through CaseWhen") {
assertEquivalent(EqualTo(caseWhen, Literal(4)), FalseLiteral)
assertEquivalent(EqualTo(caseWhen, Literal(3)),
Expand All @@ -105,13 +128,12 @@ class PushFoldableIntoBranchesSuite
And(EqualTo(caseWhen, Literal(5)), EqualTo(caseWhen, Literal(6))),
FalseLiteral)

// Do not simplify if it contains non foldable expressions.
assertEquivalent(EqualTo(caseWhen, NonFoldableLiteral(true)),
EqualTo(caseWhen, NonFoldableLiteral(true)))
val nonFoldable = CaseWhen(Seq(normalBranch, (a, b)), None)
assertEquivalent(EqualTo(nonFoldable, Literal(1)), EqualTo(nonFoldable, Literal(1)))
// Partially push down if it contains non foldable expressions.
val nonFoldable = CaseWhen(Seq((NonFoldableLiteral(true), Literal(10)), (a, b)), None)
assertEquivalent(EqualTo(nonFoldable, Literal(1)),
CaseWhen(Seq((NonFoldableLiteral(true), FalseLiteral), (a, EqualTo(b, Literal(1)))), None))

// Do not simplify if it contains non-deterministic expressions.
// Do not push down if it contains non-deterministic expressions.
val nonDeterministic = CaseWhen(Seq((LessThan(Rand(1), Literal(0.5)), Literal(1))), Some(b))
assert(!nonDeterministic.deterministic)
assertEquivalent(EqualTo(nonDeterministic, Literal(-1)), EqualTo(nonDeterministic, Literal(-1)))
Expand Down Expand Up @@ -143,4 +165,53 @@ class PushFoldableIntoBranchesSuite
assertEquivalent(LessThan(caseWhen, Literal(4)), TrueLiteral)
assertEquivalent(LessThanOrEqual(caseWhen, Literal(4)), TrueLiteral)
}

test("SPARK-33798: Push down other BinaryOperator through CaseWhen") {
assertEquivalent(Add(caseWhen, Literal(4)),
CaseWhen(Seq((a, Literal(5)), (c, Literal(6))), Some(Literal(7))))
assertEquivalent(Subtract(caseWhen, Literal(4)),
CaseWhen(Seq((a, Literal(-3)), (c, Literal(-2))), Some(Literal(-1))))
assertEquivalent(Multiply(caseWhen, Literal(4)),
CaseWhen(Seq((a, Literal(4)), (c, Literal(8))), Some(Literal(12))))
assertEquivalent(Pmod(caseWhen, Literal(4)),
CaseWhen(Seq((a, Literal(1)), (c, Literal(2))), Some(Literal(3))))
assertEquivalent(Remainder(caseWhen, Literal(4)),
CaseWhen(Seq((a, Literal(1)), (c, Literal(2))), Some(Literal(3))))
assertEquivalent(Divide(CaseWhen(Seq((a, Literal(1.0)), (c, Literal(2.0))), Some(Literal(3.0))),
Literal(1.0)),
CaseWhen(Seq((a, Literal(1.0)), (c, Literal(2.0))), Some(Literal(3.0))))
assertEquivalent(And(CaseWhen(Seq((a, FalseLiteral), (c, TrueLiteral)), Some(TrueLiteral)),
TrueLiteral),
CaseWhen(Seq((a, FalseLiteral), (c, TrueLiteral)), Some(TrueLiteral)))
assertEquivalent(Or(CaseWhen(Seq((a, FalseLiteral), (c, TrueLiteral)), Some(TrueLiteral)),
TrueLiteral), TrueLiteral)
}

test("SPARK-33798: Push down other BinaryExpression through CaseWhen") {
assertEquivalent(
BRound(CaseWhen(Seq((a, Literal(1.23)), (c, Literal(1.24))), Some(Literal(1.25))),
Literal(1)),
Literal(1.2))
assertEquivalent(
StartsWith(CaseWhen(Seq((a, Literal("ab")), (c, Literal("ac"))), Some(Literal("ad"))),
Literal("a")),
TrueLiteral)
assertEquivalent(
FindInSet(CaseWhen(Seq((a, Literal("ab")), (c, Literal("ac"))), Some(Literal("ad"))),
Literal("a")),
Literal(0))
assertEquivalent(
AddMonths(CaseWhen(Seq((a, Literal(Date.valueOf("2020-01-01"))),
(c, Literal(Date.valueOf("2021-01-01")))),
Some(Literal(Date.valueOf("2022-01-01")))),
Literal(1)),
CaseWhen(Seq((a, Literal(Date.valueOf("2020-02-01"))),
(c, Literal(Date.valueOf("2021-02-01")))),
Some(Literal(Date.valueOf("2022-02-01")))))
}

test("SPARK-33798: Push down BinaryExpression through If/CaseWhen backwards") {
assertEquivalent(EqualTo(Literal(4), ifExp), FalseLiteral)
assertEquivalent(EqualTo(Literal(4), caseWhen), FalseLiteral)
}
}