Skip to content

Commit 386e464

Browse files
dusantism-dbcloud-fan
authored andcommitted
[SPARK-52870][SQL] Properly quote variable names in FOR statement
### What changes were proposed in this pull request? This PR fixes the bugs presented in https://issues.apache.org/jira/browse/SPARK-52870 . Before this change, we passed FOR statement referenced column names as multipart identifiers to `SetVariable`. This fails when the columns are not formatted as proper multipart identifiers. This PR changes the logic to always quote the column names, similar to how backticks are used from the SQL side. ### Why are the changes needed? Bugfix. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Unit tests added. ### Was this patch authored or co-authored using generative AI tooling? No. Closes apache#51560 from dusantism-db/for-quote-variable-names. Authored-by: Dušan Tišma <[email protected]> Signed-off-by: Wenchen Fan <[email protected]>
1 parent f8c2671 commit 386e464

File tree

2 files changed

+73
-1
lines changed

2 files changed

+73
-1
lines changed

sql/core/src/main/scala/org/apache/spark/sql/scripting/SqlScriptingExecutionNode.scala

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1187,7 +1187,7 @@ class ForStatementExec(
11871187
OneRowRelation()
11881188
)
11891189
val setIdentifierToCurrentRow =
1190-
SetVariable(Seq(UnresolvedAttribute(varName)), projectNamedStruct)
1190+
SetVariable(Seq(UnresolvedAttribute.quoted(varName)), projectNamedStruct)
11911191
new SingleStatementExec(
11921192
setIdentifierToCurrentRow,
11931193
Origin(),

sql/core/src/test/scala/org/apache/spark/sql/scripting/SqlScriptingExecutionSuite.scala

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2938,4 +2938,76 @@ class SqlScriptingExecutionSuite extends QueryTest with SharedSparkSession {
29382938
)
29392939
verifySqlScriptResult(sqlScript, expected = expected)
29402940
}
2941+
2942+
test("Integer literal column in FOR query") {
2943+
val sqlScript1 =
2944+
"""
2945+
|BEGIN
2946+
| FOR SELECT 1 DO
2947+
| SELECT 2;
2948+
| END FOR;
2949+
|END
2950+
|""".stripMargin
2951+
verifySqlScriptResult(sqlScript1, Seq(Seq(Row(2))))
2952+
2953+
val sqlScript2 =
2954+
"""
2955+
|BEGIN
2956+
| FOR x AS SELECT 1 DO
2957+
| SELECT x.`1`;
2958+
| END FOR;
2959+
|END
2960+
|""".stripMargin
2961+
verifySqlScriptResult(sqlScript2, Seq(Seq(Row(1))))
2962+
}
2963+
2964+
test("Column with space in FOR query") {
2965+
val sqlScript1 =
2966+
"""
2967+
|BEGIN
2968+
| FOR SELECT 1 AS `Space Column` DO
2969+
| SELECT `Space Column`;
2970+
| END FOR;
2971+
|END
2972+
|""".stripMargin
2973+
val expected = Seq(Seq(Row(1)))
2974+
verifySqlScriptResult(sqlScript1, expected)
2975+
2976+
val sqlScript2 =
2977+
"""
2978+
|BEGIN
2979+
| FOR x AS SELECT 1 AS `Space Column` DO
2980+
| SELECT x.`Space Column`;
2981+
| END FOR;
2982+
|END
2983+
|""".stripMargin
2984+
verifySqlScriptResult(sqlScript2, expected)
2985+
}
2986+
2987+
test("FOR query referencing table with space") {
2988+
withTable("test_tbl") {
2989+
sql("CREATE TABLE test_tbl (`Space Column` INT) USING parquet")
2990+
sql("INSERT INTO test_tbl VALUES (1)")
2991+
val sqlScript1 =
2992+
"""
2993+
|BEGIN
2994+
| FOR SELECT * FROM test_tbl DO
2995+
| SELECT `Space Column`;
2996+
| END FOR;
2997+
|END
2998+
|""".stripMargin
2999+
val expected = Seq(Seq(Row(1)))
3000+
verifySqlScriptResult(sqlScript1, expected)
3001+
3002+
val sqlScript2 =
3003+
"""
3004+
|BEGIN
3005+
| FOR x AS SELECT * FROM test_tbl DO
3006+
| SELECT x.`Space Column`;
3007+
| END FOR;
3008+
|END
3009+
|""".stripMargin
3010+
verifySqlScriptResult(sqlScript2, expected)
3011+
}
3012+
}
29413013
}

0 commit comments

Comments
 (0)