davidm-db commented on code in PR #48794:
URL: https://github.com/apache/spark/pull/48794#discussion_r1850838477
##########
sql/core/src/test/scala/org/apache/spark/sql/scripting/SqlScriptingExecutionNodeSuite.scala:
##########
@@ -686,4 +705,368 @@ class SqlScriptingExecutionNodeSuite extends
SparkFunSuite with SharedSparkSessi
val statements = iter.map(extractStatementValue).toSeq
assert(statements === Seq("body1", "lbl"))
}
+
+ test("for statement - enters body once") {
+ val iter = new CompoundBodyExec(Seq(
+ new ForStatementExec(
+ query = MockQuery(1, "intCol", "query1"),
+ variableName = Some("x"),
+ body = new CompoundBodyExec(Seq(TestLeafStatement("body"))),
+ label = Some("for1"),
+ session = spark
+ )
+ )).getTreeIterator
+ val statements = iter.map(extractStatementValue).toSeq
+ assert(statements === Seq(
+ "body",
+ "SingleStatementExec", // drop local var
+ "SingleStatementExec" // drop local var
+ ))
+ }
+
+ test("for statement - enters body with multiple statements multiple times") {
+ val iter = new CompoundBodyExec(Seq(
+ new ForStatementExec(
+ query = MockQuery(2, "intCol", "query1"),
+ variableName = Some("x"),
+ body = new CompoundBodyExec(Seq(
+ TestLeafStatement("statement1"),
+ TestLeafStatement("statement2"))),
+ label = Some("for1"),
+ session = spark
+ )
+ )).getTreeIterator
+ val statements = iter.map(extractStatementValue).toSeq
+ assert(statements === Seq(
+ "statement1",
+ "statement2",
+ "statement1",
+ "statement2",
+ "SingleStatementExec", // drop local var
+ "SingleStatementExec", // drop local var
+ ))
+ }
+
+ test("for statement - empty result") {
+ val iter = new CompoundBodyExec(Seq(
+ new ForStatementExec(
+ query = MockQuery(0, "intCol", "query1"),
+ variableName = Some("x"),
+ body = new CompoundBodyExec(Seq(TestLeafStatement("body1"))),
+ label = Some("for1"),
+ session = spark
+ )
+ )).getTreeIterator
+ val statements = iter.map(extractStatementValue).toSeq
+ assert(statements === Seq.empty[String])
+ }
+
+ test("for statement - nested") {
+ val iter = new CompoundBodyExec(Seq(
+ new ForStatementExec(
+ query = MockQuery(2, "intCol", "query1"),
+ variableName = Some("x"),
+ body = new CompoundBodyExec(Seq(
+ new ForStatementExec(
+ query = MockQuery(2, "intCol1", "query2"),
+ variableName = Some("y"),
+ body = new CompoundBodyExec(Seq(TestLeafStatement("body"))),
+ label = Some("for2"),
+ session = spark
+ )
+ )),
+ label = Some("for1"),
+ session = spark
+ )
+ )).getTreeIterator
+ val statements = iter.map(extractStatementValue).toSeq
+ assert(statements === Seq(
+ "body",
+ "body",
+ "SingleStatementExec", // drop inner local var
+ "SingleStatementExec", // drop inner local var
+ "body",
+ "body",
+ "SingleStatementExec", // drop inner local var
+ "SingleStatementExec", // drop inner local var
+ "SingleStatementExec", // drop outer local var
+ "SingleStatementExec", // drop outer local var
+ ))
+ }
+
+ test("for statement no variable - enters body once") {
+ val iter = new CompoundBodyExec(Seq(
+ new ForStatementExec(
+ query = MockQuery(1, "intCol", "query1"),
+ variableName = None,
+ body = new CompoundBodyExec(Seq(TestLeafStatement("body"))),
+ label = Some("for1"),
+ session = spark
+ )
+ )).getTreeIterator
+ val statements = iter.map(extractStatementValue).toSeq
+ assert(statements === Seq(
+ "body",
+ "SingleStatementExec", // drop local var
+ ))
+ }
+
+ test("for statement no variable - enters body with multiple statements
multiple times") {
+ val iter = new CompoundBodyExec(Seq(
+ new ForStatementExec(
+ query = MockQuery(2, "intCol", "query1"),
+ variableName = None,
+ body = new CompoundBodyExec(Seq(
+ TestLeafStatement("statement1"),
+ TestLeafStatement("statement2"))),
+ label = Some("for1"),
+ session = spark
+ )
+ )).getTreeIterator
+ val statements = iter.map(extractStatementValue).toSeq
+ assert(statements === Seq(
+ "statement1", "statement2", "statement1", "statement2",
+ "SingleStatementExec", // drop local var
+ ))
+ }
+
+ test("for statement no variable - empty result") {
+ val iter = new CompoundBodyExec(Seq(
+ new ForStatementExec(
+ query = MockQuery(0, "intCol", "query1"),
+ variableName = None,
+ body = new CompoundBodyExec(Seq(TestLeafStatement("body1"))),
+ label = Some("for1"),
+ session = spark
+ )
+ )).getTreeIterator
+ val statements = iter.map(extractStatementValue).toSeq
+ assert(statements === Seq.empty[String])
+ }
+
+ test("for statement no variable - nested") {
+ val iter = new CompoundBodyExec(Seq(
+ new ForStatementExec(
+ query = MockQuery(2, "intCol", "query1"),
+ variableName = None,
+ body = new CompoundBodyExec(Seq(
+ new ForStatementExec(
+ query = MockQuery(2, "intCol1", "query2"),
+ variableName = None,
+ body = new CompoundBodyExec(Seq(TestLeafStatement("body"))),
+ label = Some("for2"),
+ session = spark
+ )
+ )),
+ label = Some("for1"),
+ session = spark
+ )
+ )).getTreeIterator
+ val statements = iter.map(extractStatementValue).toSeq
+ assert(statements === Seq(
+ "body", "body",
+ "SingleStatementExec", // drop inner local var
+ "body", "body",
+ "SingleStatementExec", // drop inner local var
+ "SingleStatementExec", // drop outer local var
+ ))
+ }
+
+ test("for statement - iterate") {
+ val iter = new CompoundBodyExec(Seq(
+ new ForStatementExec(
+ query = MockQuery(2, "intCol", "query1"),
+ variableName = Some("x"),
+ body = new CompoundBodyExec(Seq(
+ TestLeafStatement("statement1"),
+ new IterateStatementExec("lbl1"),
+ TestLeafStatement("statement2"))),
+ label = Some("lbl1"),
+ session = spark
+ )
+ )).getTreeIterator
+ val statements = iter.map(extractStatementValue).toSeq
+ assert(statements === Seq(
+ "statement1",
+ "lbl1",
+ "statement1",
+ "lbl1",
+ "SingleStatementExec", // drop local var
+ "SingleStatementExec", // drop local var
+ ))
+ }
+
+ test("for statement - leave") {
+ val iter = new CompoundBodyExec(Seq(
+ new ForStatementExec(
+ query = MockQuery(2, "intCol", "query1"),
+ variableName = Some("x"),
+ body = new CompoundBodyExec(Seq(
+ TestLeafStatement("statement1"),
+ new LeaveStatementExec("lbl1"),
+ TestLeafStatement("statement2"))),
+ label = Some("lbl1"),
+ session = spark
+ )
+ )).getTreeIterator
+ val statements = iter.map(extractStatementValue).toSeq
+ assert(statements === Seq(
+ "statement1",
+ "lbl1"
+ ))
Review Comment:
we don't have drop var statements here due to the fact that they are dropped
in `handleLeaveStatement`?
this is the thing we talked about that will be properly resolved once the
proper execution and scopes are introduced?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
