[SPARK-58966][SQL] Resolve SQL variables in UPDATE and MERGE INTO conditions

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala

@@ -1811,11 +1811,11 @@ class Analyzer(
             val newMatchedActions = m.matchedActions.map {
               case DeleteAction(deleteCondition) =>
                 val resolvedDeleteCondition = deleteCondition.map(
-                  resolveExpressionByPlanChildren(_, m))
+                  resolveExpressionByPlanChildren(_, m, includeLastResort = true))
                 DeleteAction(resolvedDeleteCondition)
               case UpdateAction(updateCondition, assignments, fromStar) =>
                 val resolvedUpdateCondition = updateCondition.map(
-                  resolveExpressionByPlanChildren(_, m))
+                  resolveExpressionByPlanChildren(_, m, includeLastResort = true))
                 UpdateAction(
                   resolvedUpdateCondition,
                   // The update value can access columns from both target and source tables.
@@ -1838,7 +1838,8 @@ class Analyzer(
                   }
                 }
                 UpdateAction(
-                  updateCondition.map(resolveExpressionByPlanChildren(_, m)),
+                  updateCondition.map(
+                    resolveExpressionByPlanChildren(_, m, includeLastResort = true)),
                   // For UPDATE *, the value must be from source table.
                   resolveAssignments(assignments, m, MergeResolvePolicy.SOURCE, throws),
                   fromStar = true)
@@ -1849,15 +1850,15 @@ class Analyzer(
                 // The insert action is used when not matched, so its condition and value can only
                 // access columns from the source table.
                 val resolvedInsertCondition = insertCondition.map(
-                  resolveExpressionByPlanOutput(_, m.sourceTable))
+                  resolveExpressionByPlanOutput(_, m.sourceTable, includeLastResort = true))
                 InsertAction(
                   resolvedInsertCondition,
                   resolveAssignments(assignments, m, MergeResolvePolicy.SOURCE, throws))
               case InsertStarAction(insertCondition) =>
                 // The insert action is used when not matched, so its condition and value can only
                 // access columns from the source table.
                 val resolvedInsertCondition = insertCondition.map(
-                  resolveExpressionByPlanOutput(_, m.sourceTable))
+                  resolveExpressionByPlanOutput(_, m.sourceTable, includeLastResort = true))
                 // Expand star to top level source columns.  If source has less columns than target,
                 // assignments will be added by ResolveRowLevelCommandAssignments later.
                 val assignments = if (m.schemaEvolutionEnabled) {
@@ -1881,11 +1882,11 @@ class Analyzer(
             val newNotMatchedBySourceActions = m.notMatchedBySourceActions.map {
               case DeleteAction(deleteCondition) =>
                 val resolvedDeleteCondition = deleteCondition.map(
-                  resolveExpressionByPlanOutput(_, targetTable))
+                  resolveExpressionByPlanOutput(_, targetTable, includeLastResort = true))
                 DeleteAction(resolvedDeleteCondition)
               case UpdateAction(updateCondition, assignments, fromStar) =>
                 val resolvedUpdateCondition = updateCondition.map(
-                  resolveExpressionByPlanOutput(_, targetTable))
+                  resolveExpressionByPlanOutput(_, targetTable, includeLastResort = true))
                 UpdateAction(
                   resolvedUpdateCondition,
                   // The update value can access columns from the target table only.
@@ -1894,7 +1895,8 @@ class Analyzer(
               case o => o
             }
 
-            val resolvedMergeCondition = resolveExpressionByPlanChildren(m.mergeCondition, m)
+            val resolvedMergeCondition =
+              resolveExpressionByPlanChildren(m.mergeCondition, m, includeLastResort = true)
             m.copy(mergeCondition = resolvedMergeCondition,
               matchedActions = newMatchedActions,
               notMatchedActions = newNotMatchedActions,

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveReferencesInUpdate.scala

@@ -67,7 +67,8 @@ class ResolveReferencesInUpdate(val catalogManager: CatalogManager)
 
     val newUpdate = u.copy(
       assignments = newAssignments,
-      condition = u.condition.map(resolveExpressionByPlanChildren(_, u)))
+      condition = u.condition.map(
+        resolveExpressionByPlanChildren(_, u, includeLastResort = true)))
     newUpdate.copyTagsFrom(u)
     newUpdate
   }

sql/core/src/test/scala/org/apache/spark/sql/connector/MergeIntoTableSuiteBase.scala

@@ -3052,4 +3052,99 @@ abstract class MergeIntoTableSuiteBase extends RowLevelOperationSuiteBase
       case None => fail(s"$metricName metric not found")
     }
   }
+
+  test("merge with a SQL variable in the merge condition") {
+    withTempView("source") {
+      createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+        """{ "pk": 1, "salary": 100, "dep": "hr" }
+          |{ "pk": 2, "salary": 200, "dep": "software" }
+          |""".stripMargin)
+      Seq(1, 2).toDF("pk").createOrReplaceTempView("source")
+
+      withSessionVariable("target_dep", "STRING", "'hr'") {
+        sql(
+          s"""MERGE INTO $tableNameAsString t
+             |USING source s
+             |ON t.pk = s.pk AND t.dep = target_dep
+             |WHEN MATCHED THEN UPDATE SET t.salary = 999
+             |""".stripMargin)
+
+        checkAnswer(
+          sql(s"SELECT * FROM $tableNameAsString"),
+          Row(1, 999, "hr") :: Row(2, 200, "software") :: Nil)
+      }
+    }
+  }
+
+  test("merge with a SQL variable in matched and not-matched-by-source conditions") {
+    withTempView("source") {
+      createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+        """{ "pk": 1, "salary": 100, "dep": "hr" }
+          |{ "pk": 2, "salary": 200, "dep": "software" }
+          |""".stripMargin)
+      Seq(1).toDF("pk").createOrReplaceTempView("source")
+
+      withSessionVariable("salary_threshold", "INT", "150") {
+        sql(
+          s"""MERGE INTO $tableNameAsString t
+             |USING source s
+             |ON t.pk = s.pk
+             |WHEN MATCHED AND t.salary < salary_threshold THEN UPDATE SET t.salary = 999
+             |WHEN NOT MATCHED BY SOURCE AND t.salary > salary_threshold THEN DELETE
+             |""".stripMargin)
+
+        checkAnswer(
+          sql(s"SELECT * FROM $tableNameAsString"),
+          Row(1, 999, "hr") :: Nil)
+      }
+    }
+  }
+
+  test("merge with a SQL variable in the not-matched insert condition") {
+    withTempView("source") {
+      createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+        """{ "pk": 1, "salary": 100, "dep": "hr" }
+          |""".stripMargin)
+      Seq((2, 200, "software"), (3, 300, "hr"))
+        .toDF("pk", "salary", "dep").createOrReplaceTempView("source")
+
+      withSessionVariable("pk_threshold", "INT", "3") {
+        sql(
+          s"""MERGE INTO $tableNameAsString t
+             |USING source s
+             |ON t.pk = s.pk
+             |WHEN NOT MATCHED AND s.pk < pk_threshold THEN
+             |  INSERT (pk, salary, dep) VALUES (s.pk, s.salary, s.dep)
+             |""".stripMargin)
+
+        checkAnswer(
+          sql(s"SELECT * FROM $tableNameAsString"),
+          Row(1, 100, "hr") :: Row(2, 200, "software") :: Nil)
+      }
+    }
+  }
+
+  test("merge with a SQL variable in the not-matched insert-star condition") {
+    withTempView("source") {
+      createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+        """{ "pk": 1, "salary": 100, "dep": "hr" }
+          |""".stripMargin)
+      Seq((2, 200, "software"), (3, 300, "hr"))
+        .toDF("pk", "salary", "dep").createOrReplaceTempView("source")
+
+      withSessionVariable("pk_threshold", "INT", "3") {
+        sql(
+          s"""MERGE INTO $tableNameAsString t
+             |USING source s
+             |ON t.pk = s.pk
+             |WHEN NOT MATCHED AND s.pk < pk_threshold THEN INSERT *
+             |""".stripMargin)
+
+        checkAnswer(
+          sql(s"SELECT * FROM $tableNameAsString"),
+          Row(1, 100, "hr") :: Row(2, 200, "software") :: Nil)
+      }
+    }
+  }
+
 }

sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationSuiteBase.scala

@@ -127,6 +127,15 @@ abstract class RowLevelOperationSuiteBase
     catalog.createTable(ident, tableInfo)
   }
 
+  /** Declares a session variable for the duration of `f`, dropping it afterwards. */
+  protected def withSessionVariable(
+      name: String,
+      dataType: String,
+      default: String)(f: => Unit): Unit = {
+    sql(s"DECLARE OR REPLACE VARIABLE $name $dataType DEFAULT $default")
+    try f finally sql(s"DROP TEMPORARY VARIABLE IF EXISTS $name")
+  }
+
   protected def createAndInitTable(schemaString: String, jsonData: String): Unit = {
     createTable(schemaString)
     append(schemaString, jsonData)

sql/core/src/test/scala/org/apache/spark/sql/connector/UpdateTableSuiteBase.scala

@@ -1370,4 +1370,38 @@ abstract class UpdateTableSuiteBase extends RowLevelOperationSuiteBase {
       sql(s"SELECT * FROM $tableNameAsString"),
       Row(1, -1, "hr") :: Row(2, 200, "software") :: Row(3, -1, "hr") :: Nil)
   }
+
+  test("update with a SQL variable in the condition") {
+    createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+      """{ "pk": 1, "salary": 100, "dep": "hr" }
+        |{ "pk": 2, "salary": 200, "dep": "software" }
+        |""".stripMargin)
+
+    withSessionVariable("target_dep", "STRING", "'hr'") {
+      sql(s"UPDATE $tableNameAsString SET salary = 999 WHERE dep = target_dep")
+
+      checkAnswer(
+        sql(s"SELECT * FROM $tableNameAsString"),
+        Row(1, 999, "hr") :: Row(2, 200, "software") :: Nil)
+    }
+  }
+
+  test("update with a SQL scripting local variable in the condition") {
+    createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+      """{ "pk": 1, "salary": 100, "dep": "hr" }
+        |{ "pk": 2, "salary": 200, "dep": "software" }
+        |""".stripMargin)
+
+    sql(
+      s"""BEGIN
+         |  DECLARE local_dep STRING DEFAULT 'hr';
+         |  UPDATE $tableNameAsString SET salary = 999 WHERE dep = local_dep;
+         |END
+         |""".stripMargin)
+
+    checkAnswer(
+      sql(s"SELECT * FROM $tableNameAsString"),
+      Row(1, 999, "hr") :: Row(2, 200, "software") :: Nil)
+  }
+
 }