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..9bab78a70296b 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 @@ -1805,17 +1805,24 @@ class Analyzer( m case _ => + // Defer outer-reference and variable resolution until schema evolution has produced + // the final target schema so newly added target columns retain precedence. + val canResolveLastResort = !m.schemaEvolutionEnabled || + (m.schemaEvolutionReady && m.pendingSchemaChanges.isEmpty) + def findAttrInTarget(name: String): Option[Attribute] = { targetTable.output.find(targetAttr => conf.resolver(name, targetAttr.name)) } val newMatchedActions = m.matchedActions.map { case DeleteAction(deleteCondition) => val resolvedDeleteCondition = deleteCondition.map( - resolveExpressionByPlanChildren(_, m)) + resolveExpressionByPlanChildren( + _, m, includeLastResort = canResolveLastResort)) DeleteAction(resolvedDeleteCondition) case UpdateAction(updateCondition, assignments, fromStar) => val resolvedUpdateCondition = updateCondition.map( - resolveExpressionByPlanChildren(_, m)) + resolveExpressionByPlanChildren( + _, m, includeLastResort = canResolveLastResort)) UpdateAction( resolvedUpdateCondition, // The update value can access columns from both target and source tables. @@ -1838,7 +1845,9 @@ class Analyzer( } } UpdateAction( - updateCondition.map(resolveExpressionByPlanChildren(_, m)), + updateCondition.map( + resolveExpressionByPlanChildren( + _, m, includeLastResort = canResolveLastResort)), // For UPDATE *, the value must be from source table. resolveAssignments(assignments, m, MergeResolvePolicy.SOURCE, throws), fromStar = true) @@ -1849,7 +1858,8 @@ 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 = canResolveLastResort)) InsertAction( resolvedInsertCondition, resolveAssignments(assignments, m, MergeResolvePolicy.SOURCE, throws)) @@ -1857,7 +1867,8 @@ 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 = canResolveLastResort)) // 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 +1892,13 @@ class Analyzer( val newNotMatchedBySourceActions = m.notMatchedBySourceActions.map { case DeleteAction(deleteCondition) => val resolvedDeleteCondition = deleteCondition.map( - resolveExpressionByPlanOutput(_, targetTable)) + resolveExpressionByPlanOutput( + _, targetTable, includeLastResort = canResolveLastResort)) DeleteAction(resolvedDeleteCondition) case UpdateAction(updateCondition, assignments, fromStar) => val resolvedUpdateCondition = updateCondition.map( - resolveExpressionByPlanOutput(_, targetTable)) + resolveExpressionByPlanOutput( + _, targetTable, includeLastResort = canResolveLastResort)) UpdateAction( resolvedUpdateCondition, // The update value can access columns from the target table only. @@ -1894,7 +1907,9 @@ class Analyzer( case o => o } - val resolvedMergeCondition = resolveExpressionByPlanChildren(m.mergeCondition, m) + val resolvedMergeCondition = + resolveExpressionByPlanChildren( + m.mergeCondition, m, includeLastResort = canResolveLastResort) 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..9fe46a9ab7718 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 @@ -25,13 +25,16 @@ import org.apache.spark.sql.errors.QueryCompilationErrors /** * A virtual rule to resolve [[UnresolvedAttribute]] in [[UpdateTable]]. It's only used by the real - * rule `ResolveReferences`. The column resolution order for [[UpdateTable]] is: + * rule `ResolveReferences`. Assignments and the condition share the first two steps below, then + * diverge: step 3 applies to assignments only, and steps 4 and 5 to the condition only. * 1. Resolves the column to `AttributeReference` with the output of the child plan. This * includes metadata columns as well. * 2. Resolves the column to a literal function which is allowed to be invoked without braces, e.g. * `SELECT col, current_date FROM t`. * 3. Resolves the column to the default value expression, if the column is the assignment value * and the corresponding assignment key is a top-level column. + * 4. Resolves the column to an outer reference. + * 5. Resolves the column to a SQL variable, which is always tried after outer references. */ class ResolveReferencesInUpdate(val catalogManager: CatalogManager) extends SQLConfHelper with ColumnResolutionHelper { @@ -67,7 +70,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/MergeIntoSchemaEvolutionExtraSourceColumnTests.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/MergeIntoSchemaEvolutionExtraSourceColumnTests.scala index ac9c5dbf76b72..4312e4ffe71de 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/MergeIntoSchemaEvolutionExtraSourceColumnTests.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/MergeIntoSchemaEvolutionExtraSourceColumnTests.scala @@ -120,6 +120,41 @@ trait MergeIntoSchemaEvolutionExtraSourceColumnTests extends MergeIntoSchemaEvol "`active` cannot be resolved" ) + test("schema evolution - evolved target column takes precedence over a SQL variable") { + withTable(tableNameAsString) { + withTempView("source") { + createAndInitTable( + "pk INT NOT NULL, salary INT, dep STRING", + """{ "pk": 1, "salary": 100, "dep": "hr" } + |{ "pk": 2, "salary": 200, "dep": "software" } + |{ "pk": 3, "salary": 300, "dep": "hr" }""".stripMargin) + Seq((1, 150, "hr", true)).toDF("pk", "salary", "dep", "active") + .createOrReplaceTempView("source") + + withSessionVariable("active", "action_enabled") { + sql("DECLARE VARIABLE active BOOLEAN DEFAULT false") + sql("DECLARE VARIABLE action_enabled BOOLEAN DEFAULT true") + executeMerge( + withSchemaEvolution = true, + targetTableName = tableNameAsString, + sourceViewName = "source", + cond = "t.pk = s.pk", + clauses = Seq( + update(set = "salary = s.salary, active = s.active"), + updateNotMatched( + set = "salary = salary + 1", + condition = "active IS NULL AND action_enabled AND salary = 200"), + deleteNotMatched( + condition = "active IS NULL AND action_enabled AND salary = 300"))) + + checkAnswer( + sql(s"SELECT * FROM $tableNameAsString"), + Row(1, 150, "hr", true) :: Row(2, 201, "software", null) :: Nil) + } + } + } + } + testEvolution("source has extra column with set all columns")( targetData = Seq( (1, 100, "hr"), 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..93338ba043d9b 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,152 @@ 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") { + sql("DECLARE VARIABLE target_dep STRING DEFAULT '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") { + sql("DECLARE VARIABLE salary_threshold INT DEFAULT 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 matched delete 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("salary_threshold") { + sql("DECLARE VARIABLE salary_threshold INT DEFAULT 150") + sql( + s"""MERGE INTO $tableNameAsString t + |USING source s + |ON t.pk = s.pk + |WHEN MATCHED AND t.salary < salary_threshold THEN DELETE + |""".stripMargin) + + checkAnswer( + sql(s"SELECT * FROM $tableNameAsString"), + Row(2, 200, "software") :: Nil) + } + } + } + + test("merge with a SQL variable in matched update-star 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, 999, "finance"), (2, 888, "sales")) + .toDF("pk", "salary", "dep").createOrReplaceTempView("source") + + withSessionVariable("pk_threshold") { + sql("DECLARE VARIABLE pk_threshold INT DEFAULT 2") + sql( + s"""MERGE INTO $tableNameAsString t + |USING source s + |ON t.pk = s.pk + |WHEN MATCHED AND s.pk < pk_threshold THEN UPDATE SET * + |""".stripMargin) + + checkAnswer( + sql(s"SELECT * FROM $tableNameAsString"), + Row(1, 999, "finance") :: Row(2, 200, "software") :: 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") { + sql("DECLARE VARIABLE pk_threshold INT DEFAULT 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") { + sql("DECLARE VARIABLE pk_threshold INT DEFAULT 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/UpdateTableSuiteBase.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/UpdateTableSuiteBase.scala index c122713c1873f..3269081227e57 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,39 @@ 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") { + sql("DECLARE VARIABLE target_dep STRING DEFAULT '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) + } + }