diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala index 22a887ae12ce8..c778d3d6f48eb 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala +++ b/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,7 +1850,7 @@ 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)) @@ -1857,7 +1858,7 @@ 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)) // 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, diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveReferencesInUpdate.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveReferencesInUpdate.scala index 92813a156988b..185c91a727f38 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveReferencesInUpdate.scala +++ b/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 } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/MergeIntoTableSuiteBase.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/MergeIntoTableSuiteBase.scala index ea15f13c225cf..303931e17dce4 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/MergeIntoTableSuiteBase.scala +++ b/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) + } + } + } + } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationSuiteBase.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationSuiteBase.scala index 21e27e9b6c1cb..bdc3577922e25 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationSuiteBase.scala +++ b/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) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/UpdateTableSuiteBase.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/UpdateTableSuiteBase.scala index c122713c1873f..2cd47760d5e81 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/UpdateTableSuiteBase.scala +++ b/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) + } + }