fix

Author: wangyum

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

@@ -511,9 +511,13 @@ object SimplifyConditionals extends Rule[LogicalPlan] with PredicateHelper {
           e.copy(branches = branches.take(i).map(branch => (branch._1, elseValue)))
         }
 
-      case EqualTo(CaseWhen(branches, _), right)
-          if branches.count(_._2.semanticEquals(right)) == 1 =>
-        branches.filter(_._2.semanticEquals(right)).head._1
+      case EqualTo(CaseWhen(branches, elseValue), right)
+          if right.foldable && branches.forall(_._2.foldable) =>
+        (branches.filter(_._2.equals(right)).map(_._1) ++
+          elseValue.map(e => EqualTo(e, right))).reduceLeftOption(Or) match {
+          case Some(value) => value
+          case None => FalseLiteral
+        }
     }
   }
 }

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/SimplifyConditionalSuite.scala

@@ -202,23 +202,29 @@ class SimplifyConditionalSuite extends PlanTest with ExpressionEvalHelper with P
 
   test("SPARK-33315: simplify CaseWhen with EqualTo") {
     val e1 = EqualTo(UnresolvedAttribute("a"), Literal(100))
-    val e2 = GreaterThan(UnresolvedAttribute("b"), Literal(1000))
-    val e3 = IsNotNull(UnresolvedAttribute("c"))
+    val e3 = EqualTo(UnresolvedAttribute("c"), Literal(true))
     val caseWhen = CaseWhen(
-      Seq(normalBranch, (e1, Literal(1)), (e2, Literal(2)), (e3, Literal(3))), None)
-    assertEquivalent(EqualTo(caseWhen, Literal(1)), e1)
-    assertEquivalent(EqualTo(caseWhen, Literal(3)), e3)
+      Seq(normalBranch, (e1, Literal(1)), (e3, Literal(2))), Some(UnresolvedAttribute("b")))
 
-    assertEquivalent(
-      And(EqualTo(caseWhen, Literal(1)), EqualTo(caseWhen, Literal(2))),
-      And(e1, e2))
-    assertEquivalent(
-      Or(EqualTo(caseWhen, Literal(1)), EqualTo(caseWhen, Literal(2))),
-      Or(e1, e2))
+    assertEquivalent(EqualTo(caseWhen, Literal(1)),
+      Or(e1, EqualTo(UnresolvedAttribute("b"), Literal(1))))
+    assertEquivalent(EqualTo(caseWhen, Literal(3)),
+      EqualTo(UnresolvedAttribute("b"), Literal(3)))
+    assertEquivalent(EqualTo(caseWhen, Literal(4)),
+      EqualTo(UnresolvedAttribute("b"), Literal(4)))
+
+    assertEquivalent(And(EqualTo(caseWhen, Literal(1)), EqualTo(caseWhen, Literal(2))),
+      And(Or(e1, EqualTo(UnresolvedAttribute("b"), Literal(1))),
+        Or(e3, EqualTo(UnresolvedAttribute("b"), Literal(2)))))
 
-    assertEquivalent(EqualTo(caseWhen, Literal(4)), EqualTo(caseWhen, Literal(4)))
     assertEquivalent(
-      Or(EqualTo(caseWhen, Literal(3)), EqualTo(caseWhen, Literal(4))),
-      Or(e3, EqualTo(caseWhen, Literal(4))))
+      EqualTo(CaseWhen(Seq(normalBranch, (e1, Literal(1)), (e3, Literal(2))), None), Literal(3)),
+      FalseLiteral)
+
+    // Do not simplify if it contains non foldable expressions.
+    assertEquivalent(EqualTo(caseWhen, NonFoldableLiteral(true)),
+      EqualTo(caseWhen, NonFoldableLiteral(true)))
+    val nonFoldable = CaseWhen(Seq(normalBranch, (e1, UnresolvedAttribute("b"))), None)
+    assertEquivalent(EqualTo(nonFoldable, Literal(1)), EqualTo(nonFoldable, Literal(1)))
   }
 }