davidm-db commented on code in PR #48794:
URL: https://github.com/apache/spark/pull/48794#discussion_r1850845635
##########
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"
+ ))
+ }
+
+ test("for statement - nested - iterate outer loop") {
Review Comment:
nit: just a minor thought for test readability - since we are using named
parameters everywhere, maybe put `body` as the last parameter, so it's a bit
cleaner when multiple FOR statements are nested.
--
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]
