[SPARK-32030][SPARK-32127][SQL] Support unlimited MATCHED and NOT MATCHED clauses in MERGE INTO

sql/catalyst/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBase.g4

@@ -437,8 +437,7 @@ dmlStatementNoWith
         USING (source=multipartIdentifier |
           '(' sourceQuery=query')') sourceAlias=tableAlias
         ON mergeCondition=booleanExpression
-        matchedClause*
-        notMatchedClause*                                                          #mergeIntoTable
+        matchedOrNotMatchedClause*                                                 #mergeIntoTable
     ;
 
 queryOrganization
@@ -526,6 +525,11 @@ setClause
     : SET assignmentList
     ;
 
+matchedOrNotMatchedClause
+    : matchedClause
+    | notMatchedClause
+    ;
+
 matchedClause
     : WHEN MATCHED (AND matchedCond=booleanExpression)? THEN matchedAction
     ;

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala

@@ -411,12 +411,9 @@ class AstBuilder(conf: SQLConf) extends SqlBaseBaseVisitor[AnyRef] with Logging
 
     val mergeCondition = expression(ctx.mergeCondition)
 
-    val matchedClauses = ctx.matchedClause()
-    if (matchedClauses.size() > 2) {
-      throw new ParseException("There should be at most 2 'WHEN MATCHED' clauses.",
-        matchedClauses.get(2))
-    }
-    val matchedActions = matchedClauses.asScala.map {
+    val whenClauses = ctx.matchedOrNotMatchedClause()
+    val matchedClauses = whenClauses.asScala.map(_.matchedClause()).filter(_ != null)
+    val matchedActions = matchedClauses.map {
       clause => {
         if (clause.matchedAction().DELETE() != null) {
           DeleteAction(Option(clause.matchedCond).map(expression))
@@ -435,12 +432,8 @@ class AstBuilder(conf: SQLConf) extends SqlBaseBaseVisitor[AnyRef] with Logging
         }
       }
     }
-    val notMatchedClauses = ctx.notMatchedClause()
-    if (notMatchedClauses.size() > 1) {
-      throw new ParseException("There should be at most 1 'WHEN NOT MATCHED' clause.",
-        notMatchedClauses.get(1))
-    }
-    val notMatchedActions = notMatchedClauses.asScala.map {
+    val notMatchedClauses = whenClauses.asScala.map(_.notMatchedClause()).filter(_ != null)
+    val notMatchedActions = notMatchedClauses.map {
       clause => {
         if (clause.notMatchedAction().INSERT() != null) {
           val condition = Option(clause.notMatchedCond).map(expression)
@@ -468,13 +461,25 @@ class AstBuilder(conf: SQLConf) extends SqlBaseBaseVisitor[AnyRef] with Logging
       throw new ParseException("There must be at least one WHEN clause in a MERGE statement", ctx)
     }
     // children being empty means that the condition is not set
-    if (matchedActions.length == 2 && matchedActions.head.children.isEmpty) {
-      throw new ParseException("When there are 2 MATCHED clauses in a MERGE statement, " +
-        "the first MATCHED clause must have a condition", ctx)
-    }
-    if (matchedActions.groupBy(_.getClass).mapValues(_.size).exists(_._2 > 1)) {
+    val matchedActionSize = matchedActions.length
+    if (matchedActionSize >= 2 && !matchedActions.init.forall(_.condition.nonEmpty)) {
+      throw new ParseException(
+        s"When there are $matchedActionSize MATCHED clauses in a MERGE statement, " +
+            s"${if (matchedActionSize == 2) {
+              "the first MATCHED clause must have a condition"
+            } else {
+              s"the first ${matchedActionSize - 1} MATCHED clauses must have conditions"
+            }}", ctx)
+    }
+    val notMatchedActionSize = notMatchedActions.length
+    if (notMatchedActionSize >= 2 && !notMatchedActions.init.forall(_.condition.nonEmpty)) {
       throw new ParseException(
-        "UPDATE and DELETE can appear at most once in MATCHED clauses in a MERGE statement", ctx)
+        s"When there are $notMatchedActionSize NOT MATCHED clauses in a MERGE statement, " +
+            s"${if (notMatchedActionSize == 2) {
+              "the first NOT MATCHED clause must have a condition"
+            } else {
+              s"the first ${notMatchedActionSize - 1} NOT MATCHED clauses must have conditions"
+            }}", ctx)
     }
 
     MergeIntoTable(

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/v2Commands.scala

@@ -347,23 +347,23 @@ case class MergeIntoTable(
 }
 
 sealed abstract class MergeAction(
-    condition: Option[Expression]) extends Expression with Unevaluable {
+    val condition: Option[Expression]) extends Expression with Unevaluable {
   override def foldable: Boolean = false
   override def nullable: Boolean = false
   override def dataType: DataType = throw new UnresolvedException(this, "nullable")
   override def children: Seq[Expression] = condition.toSeq
 }
 
-case class DeleteAction(condition: Option[Expression]) extends MergeAction(condition)
+case class DeleteAction(override val condition: Option[Expression]) extends MergeAction(condition)
 
 case class UpdateAction(
-    condition: Option[Expression],
+    override val condition: Option[Expression],
     assignments: Seq[Assignment]) extends MergeAction(condition) {
   override def children: Seq[Expression] = condition.toSeq ++ assignments
 }
 
 case class InsertAction(
-    condition: Option[Expression],
+    override val condition: Option[Expression],
     assignments: Seq[Assignment]) extends MergeAction(condition) {
   override def children: Seq[Expression] = condition.toSeq ++ assignments
 }

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/parser/DDLParserSuite.scala

@@ -1134,58 +1134,70 @@ class DDLParserSuite extends AnalysisTest {
     }
   }
 
-  test("merge into table: at most two matched clauses") {
-    val exc = intercept[ParseException] {
-      parsePlan(
-        """
-          |MERGE INTO testcat1.ns1.ns2.tbl AS target
-          |USING testcat2.ns1.ns2.tbl AS source
-          |ON target.col1 = source.col1
-          |WHEN MATCHED AND (target.col2='delete') THEN DELETE
-          |WHEN MATCHED AND (target.col2='update1') THEN UPDATE SET target.col2 = source.col2
-          |WHEN MATCHED AND (target.col2='update2') THEN UPDATE SET target.col2 = source.col2
-          |WHEN NOT MATCHED AND (target.col2='insert')
-          |THEN INSERT (target.col1, target.col2) values (source.col1, source.col2)
-        """.stripMargin)
-    }
-
-    assert(exc.getMessage.contains("There should be at most 2 'WHEN MATCHED' clauses."))
+  test("merge into table: multi matched and not matched clauses") {
+    parseCompare(
+      """
+        |MERGE INTO testcat1.ns1.ns2.tbl AS target
+        |USING testcat2.ns1.ns2.tbl AS source
+        |ON target.col1 = source.col1
+        |WHEN MATCHED AND (target.col2='delete') THEN DELETE
+        |WHEN MATCHED AND (target.col2='update to 1') THEN UPDATE SET target.col2 = 1
+        |WHEN MATCHED AND (target.col2='update to 2') THEN UPDATE SET target.col2 = 2
+        |WHEN NOT MATCHED AND (target.col2='insert 1')
+        |THEN INSERT (target.col1, target.col2) values (source.col1, 1)
+        |WHEN NOT MATCHED AND (target.col2='insert 2')
+        |THEN INSERT (target.col1, target.col2) values (source.col1, 2)
+      """.stripMargin,
+      MergeIntoTable(
+        SubqueryAlias("target", UnresolvedRelation(Seq("testcat1", "ns1", "ns2", "tbl"))),
+        SubqueryAlias("source", UnresolvedRelation(Seq("testcat2", "ns1", "ns2", "tbl"))),
+        EqualTo(UnresolvedAttribute("target.col1"), UnresolvedAttribute("source.col1")),
+        Seq(DeleteAction(Some(EqualTo(UnresolvedAttribute("target.col2"), Literal("delete")))),
+          UpdateAction(Some(EqualTo(UnresolvedAttribute("target.col2"), Literal("update to 1"))),
+            Seq(Assignment(UnresolvedAttribute("target.col2"), Literal(1)))),
+          UpdateAction(Some(EqualTo(UnresolvedAttribute("target.col2"), Literal("update to 2"))),
+            Seq(Assignment(UnresolvedAttribute("target.col2"), Literal(2))))),
+        Seq(InsertAction(Some(EqualTo(UnresolvedAttribute("target.col2"), Literal("insert 1"))),
+          Seq(Assignment(UnresolvedAttribute("target.col1"), UnresolvedAttribute("source.col1")),
+            Assignment(UnresolvedAttribute("target.col2"), Literal(1)))),
+          InsertAction(Some(EqualTo(UnresolvedAttribute("target.col2"), Literal("insert 2"))),
+            Seq(Assignment(UnresolvedAttribute("target.col1"), UnresolvedAttribute("source.col1")),
+              Assignment(UnresolvedAttribute("target.col2"), Literal(2)))))))
   }
 
-  test("merge into table: at most one not matched clause") {
+  test("merge into table: the first matched clause must have a condition if there's a second") {
     val exc = intercept[ParseException] {
       parsePlan(
         """
           |MERGE INTO testcat1.ns1.ns2.tbl AS target
           |USING testcat2.ns1.ns2.tbl AS source
           |ON target.col1 = source.col1
-          |WHEN MATCHED AND (target.col2='delete') THEN DELETE
-          |WHEN MATCHED AND (target.col2='update1') THEN UPDATE SET target.col2 = source.col2
-          |WHEN NOT MATCHED AND (target.col2='insert1')
-          |THEN INSERT (target.col1, target.col2) values (source.col1, source.col2)
-          |WHEN NOT MATCHED AND (target.col2='insert2')
+          |WHEN MATCHED THEN DELETE
+          |WHEN MATCHED THEN UPDATE SET target.col2 = source.col2
+          |WHEN NOT MATCHED AND (target.col2='insert')
           |THEN INSERT (target.col1, target.col2) values (source.col1, source.col2)
         """.stripMargin)
     }
 
-    assert(exc.getMessage.contains("There should be at most 1 'WHEN NOT MATCHED' clause."))
+    assert(exc.getMessage.contains("the first MATCHED clause must have a condition"))
   }
 
-  test("merge into table: the first matched clause must have a condition if there's a second") {
+  test("merge into table: the first 2 matched clause must have conditions if there's a third") {
     val exc = intercept[ParseException] {
       parsePlan(
         """
           |MERGE INTO testcat1.ns1.ns2.tbl AS target
           |USING testcat2.ns1.ns2.tbl AS source
           |ON target.col1 = source.col1
+          |WHEN MATCHED AND (source.col2 == '1') THEN UPDATE SET target.col2 = 1
+          |WHEN MATCHED THEN UPDATE SET target.col2 = 2
           |WHEN MATCHED THEN DELETE
-          |WHEN MATCHED THEN UPDATE SET target.col2 = source.col2
           |WHEN NOT MATCHED AND (target.col2='insert')
           |THEN INSERT (target.col1, target.col2) values (source.col1, source.col2)
         """.stripMargin)
     }
 
-    assert(exc.getMessage.contains("the first MATCHED clause must have a condition"))
+    assert(exc.getMessage.contains("the first 2 MATCHED clauses must have conditions"))
   }
 
   test("merge into table: there must be a when (not) matched condition") {
@@ -1201,26 +1213,6 @@ class DDLParserSuite extends AnalysisTest {
     assert(exc.getMessage.contains("There must be at least one WHEN clause in a MERGE statement"))
   }
 
-  test("merge into table: there can be only a single use DELETE or UPDATE") {
-    Seq("UPDATE SET *", "DELETE").foreach { op =>
-      val exc = intercept[ParseException] {
-        parsePlan(
-          s"""
-             |MERGE INTO testcat1.ns1.ns2.tbl AS target
-             |USING testcat2.ns1.ns2.tbl AS source
-             |ON target.col1 = source.col1
-             |WHEN MATCHED AND (target.col2='delete') THEN $op
-             |WHEN MATCHED THEN $op
-             |WHEN NOT MATCHED AND (target.col2='insert')
-             |THEN INSERT (target.col1, target.col2) values (source.col1, source.col2)
-           """.stripMargin)
-      }
-
-      assert(exc.getMessage.contains(
-        "UPDATE and DELETE can appear at most once in MATCHED clauses"))
-    }
-  }
-
   test("show tables") {
     comparePlans(
       parsePlan("SHOW TABLES"),