[SPARK-33736][SQL] Handle MERGE in ReplaceNullWithFalseInPredicate (apache#908)

### What changes were proposed in this pull request?

This PR adds MERGE operations to `ReplaceNullWithFalseInPredicate`.

### Why are the changes needed?

These changes are needed to optimize conditions of MERGE operations and match the existing logic for UPDATE and DELETE.

### Does this PR introduce _any_ user-facing change?

No.

### How was this patch tested?

This PR comes with new tests.

Author: aokolnychyi

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

@@ -17,10 +17,9 @@
 
 package org.apache.spark.sql.catalyst.optimizer
 
-import org.apache.spark.sql.catalyst.expressions.{And, ArrayExists, ArrayFilter, CaseWhen, Expression, If}
-import org.apache.spark.sql.catalyst.expressions.{LambdaFunction, Literal, MapFilter, Or}
+import org.apache.spark.sql.catalyst.expressions.{And, ArrayExists, ArrayFilter, CaseWhen, Expression, If, LambdaFunction, Literal, MapFilter, Or}
 import org.apache.spark.sql.catalyst.expressions.Literal.FalseLiteral
-import org.apache.spark.sql.catalyst.plans.logical.{DeleteFromTable, Filter, Join, LogicalPlan, UpdateTable}
+import org.apache.spark.sql.catalyst.plans.logical.{DeleteAction, DeleteFromTable, Filter, InsertAction, Join, LogicalPlan, MergeAction, MergeIntoTable, UpdateAction, UpdateTable}
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.types.BooleanType
 import org.apache.spark.util.Utils
@@ -55,6 +54,11 @@ object ReplaceNullWithFalseInPredicate extends Rule[LogicalPlan] {
     case j @ Join(_, _, _, Some(cond), _) => j.copy(condition = Some(replaceNullWithFalse(cond)))
     case d @ DeleteFromTable(_, Some(cond)) => d.copy(condition = Some(replaceNullWithFalse(cond)))
     case u @ UpdateTable(_, _, Some(cond)) => u.copy(condition = Some(replaceNullWithFalse(cond)))
+    case m @ MergeIntoTable(_, _, cond, matchedActions, notMatchedActions) =>
+      m.copy(
+        mergeCondition = replaceNullWithFalse(cond),
+        matchedActions = replaceNullWithFalse(matchedActions),
+        notMatchedActions = replaceNullWithFalse(notMatchedActions))
     case p: LogicalPlan => p transformExpressions {
       case i @ If(pred, _, _) => i.copy(predicate = replaceNullWithFalse(pred))
       case cw @ CaseWhen(branches, _) =>
@@ -109,4 +113,13 @@ object ReplaceNullWithFalseInPredicate extends Rule[LogicalPlan] {
         e
       }
   }
+
+  private def replaceNullWithFalse(mergeActions: Seq[MergeAction]): Seq[MergeAction] = {
+    mergeActions.map {
+      case u @ UpdateAction(Some(cond), _) => u.copy(condition = Some(replaceNullWithFalse(cond)))
+      case d @ DeleteAction(Some(cond)) => d.copy(condition = Some(replaceNullWithFalse(cond)))
+      case i @ InsertAction(Some(cond), _) => i.copy(condition = Some(replaceNullWithFalse(cond)))
+      case other => other
+    }
+  }
 }

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

@@ -24,7 +24,7 @@ import org.apache.spark.sql.catalyst.dsl.plans._
 import org.apache.spark.sql.catalyst.expressions.{And, ArrayExists, ArrayFilter, ArrayTransform, CaseWhen, Expression, GreaterThan, If, LambdaFunction, Literal, MapFilter, NamedExpression, Or, UnresolvedNamedLambdaVariable}
 import org.apache.spark.sql.catalyst.expressions.Literal.{FalseLiteral, TrueLiteral}
 import org.apache.spark.sql.catalyst.plans.{Inner, PlanTest}
-import org.apache.spark.sql.catalyst.plans.logical.{DeleteFromTable, LocalRelation, LogicalPlan, UpdateTable}
+import org.apache.spark.sql.catalyst.plans.logical.{Assignment, DeleteAction, DeleteFromTable, InsertAction, LocalRelation, LogicalPlan, MergeIntoTable, UpdateAction, UpdateTable}
 import org.apache.spark.sql.catalyst.rules.RuleExecutor
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types.{BooleanType, IntegerType}
@@ -50,6 +50,7 @@ class ReplaceNullWithFalseInPredicateSuite extends PlanTest {
     testJoin(originalCond = Literal(null, BooleanType), expectedCond = FalseLiteral)
     testDelete(originalCond = Literal(null, BooleanType), expectedCond = FalseLiteral)
     testUpdate(originalCond = Literal(null, BooleanType), expectedCond = FalseLiteral)
+    testMerge(originalCond = Literal(null, BooleanType), expectedCond = FalseLiteral)
   }
 
   test("Not expected type - replaceNullWithFalse") {
@@ -68,6 +69,7 @@ class ReplaceNullWithFalseInPredicateSuite extends PlanTest {
     testJoin(originalCond, expectedCond = FalseLiteral)
     testDelete(originalCond, expectedCond = FalseLiteral)
     testUpdate(originalCond, expectedCond = FalseLiteral)
+    testMerge(originalCond, expectedCond = FalseLiteral)
   }
 
   test("replace nulls in nested expressions in branches of If") {
@@ -79,6 +81,7 @@ class ReplaceNullWithFalseInPredicateSuite extends PlanTest {
     testJoin(originalCond, expectedCond = FalseLiteral)
     testDelete(originalCond, expectedCond = FalseLiteral)
     testUpdate(originalCond, expectedCond = FalseLiteral)
+    testMerge(originalCond, expectedCond = FalseLiteral)
   }
 
   test("replace null in elseValue of CaseWhen") {
@@ -91,6 +94,7 @@ class ReplaceNullWithFalseInPredicateSuite extends PlanTest {
     testJoin(originalCond, expectedCond)
     testDelete(originalCond, expectedCond)
     testUpdate(originalCond, expectedCond)
+    testMerge(originalCond, expectedCond)
   }
 
   test("replace null in branch values of CaseWhen") {
@@ -102,6 +106,7 @@ class ReplaceNullWithFalseInPredicateSuite extends PlanTest {
     testJoin(originalCond, expectedCond = FalseLiteral)
     testDelete(originalCond, expectedCond = FalseLiteral)
     testUpdate(originalCond, expectedCond = FalseLiteral)
+    testMerge(originalCond, expectedCond = FalseLiteral)
   }
 
   test("replace null in branches of If inside CaseWhen") {
@@ -120,6 +125,7 @@ class ReplaceNullWithFalseInPredicateSuite extends PlanTest {
     testJoin(originalCond, expectedCond)
     testDelete(originalCond, expectedCond)
     testUpdate(originalCond, expectedCond)
+    testMerge(originalCond, expectedCond)
   }
 
   test("replace null in complex CaseWhen expressions") {
@@ -141,6 +147,7 @@ class ReplaceNullWithFalseInPredicateSuite extends PlanTest {
     testJoin(originalCond, expectedCond)
     testDelete(originalCond, expectedCond)
     testUpdate(originalCond, expectedCond)
+    testMerge(originalCond, expectedCond)
   }
 
   test("replace null in Or") {
@@ -150,6 +157,7 @@ class ReplaceNullWithFalseInPredicateSuite extends PlanTest {
     testJoin(originalCond, expectedCond)
     testDelete(originalCond, expectedCond)
     testUpdate(originalCond, expectedCond)
+    testMerge(originalCond, expectedCond)
   }
 
   test("replace null in And") {
@@ -158,6 +166,7 @@ class ReplaceNullWithFalseInPredicateSuite extends PlanTest {
     testJoin(originalCond, expectedCond = FalseLiteral)
     testDelete(originalCond, expectedCond = FalseLiteral)
     testUpdate(originalCond, expectedCond = FalseLiteral)
+    testMerge(originalCond, expectedCond = FalseLiteral)
   }
 
   test("replace nulls in nested And/Or expressions") {
@@ -168,6 +177,7 @@ class ReplaceNullWithFalseInPredicateSuite extends PlanTest {
     testJoin(originalCond, expectedCond = FalseLiteral)
     testDelete(originalCond, expectedCond = FalseLiteral)
     testUpdate(originalCond, expectedCond = FalseLiteral)
+    testMerge(originalCond, expectedCond = FalseLiteral)
   }
 
   test("replace null in And inside branches of If") {
@@ -179,6 +189,7 @@ class ReplaceNullWithFalseInPredicateSuite extends PlanTest {
     testJoin(originalCond, expectedCond = FalseLiteral)
     testDelete(originalCond, expectedCond = FalseLiteral)
     testUpdate(originalCond, expectedCond = FalseLiteral)
+    testMerge(originalCond, expectedCond = FalseLiteral)
   }
 
   test("replace null in branches of If inside And") {
@@ -192,6 +203,7 @@ class ReplaceNullWithFalseInPredicateSuite extends PlanTest {
     testJoin(originalCond, expectedCond = FalseLiteral)
     testDelete(originalCond, expectedCond = FalseLiteral)
     testUpdate(originalCond, expectedCond = FalseLiteral)
+    testMerge(originalCond, expectedCond = FalseLiteral)
   }
 
   test("replace null in branches of If inside another If") {
@@ -203,6 +215,7 @@ class ReplaceNullWithFalseInPredicateSuite extends PlanTest {
     testJoin(originalCond, expectedCond = FalseLiteral)
     testDelete(originalCond, expectedCond = FalseLiteral)
     testUpdate(originalCond, expectedCond = FalseLiteral)
+    testMerge(originalCond, expectedCond = FalseLiteral)
   }
 
   test("replace null in CaseWhen inside another CaseWhen") {
@@ -212,6 +225,7 @@ class ReplaceNullWithFalseInPredicateSuite extends PlanTest {
     testJoin(originalCond, expectedCond = FalseLiteral)
     testDelete(originalCond, expectedCond = FalseLiteral)
     testUpdate(originalCond, expectedCond = FalseLiteral)
+    testMerge(originalCond, expectedCond = FalseLiteral)
   }
 
   test("inability to replace null in non-boolean branches of If") {
@@ -226,6 +240,7 @@ class ReplaceNullWithFalseInPredicateSuite extends PlanTest {
     testJoin(originalCond = condition, expectedCond = condition)
     testDelete(originalCond = condition, expectedCond = condition)
     testUpdate(originalCond = condition, expectedCond = condition)
+    testMerge(originalCond = condition, expectedCond = condition)
   }
 
   test("inability to replace null in non-boolean values of CaseWhen") {
@@ -395,6 +410,21 @@ class ReplaceNullWithFalseInPredicateSuite extends PlanTest {
     test((rel, expr) => UpdateTable(rel, Seq.empty, Some(expr)), originalCond, expectedCond)
   }
 
+  private def testMerge(originalCond: Expression, expectedCond: Expression): Unit = {
+    val func = (rel: LogicalPlan, expr: Expression) => {
+      val assignments = Seq(
+        Assignment('i, 'i),
+        Assignment('b, 'b),
+        Assignment('a, 'a),
+        Assignment('m, 'm)
+      )
+      val matchedActions = UpdateAction(Some(expr), assignments) :: DeleteAction(Some(expr)) :: Nil
+      val notMatchedActions = InsertAction(Some(expr), assignments) :: Nil
+      MergeIntoTable(rel, rel, mergeCondition = expr, matchedActions, notMatchedActions)
+    }
+    test(func, originalCond, expectedCond)
+  }
+
   private def testHigherOrderFunc(
       argument: Expression,
       createExpr: (Expression, Expression) => Expression,