davidm-db commented on code in PR #50026:
URL: https://github.com/apache/spark/pull/50026#discussion_r1965209406
##########
sql/core/src/main/scala/org/apache/spark/sql/scripting/SqlScriptingExecutionNode.scala:
##########
@@ -823,35 +827,41 @@ class ForStatementExec(
override def hasNext: Boolean = !interrupted && (state match {
case ForState.VariableAssignment => cachedQueryResult().hasNext
- case ForState.Body => true
- case ForState.VariableCleanup =>
dropVariablesExec.getTreeIterator.hasNext
+ case ForState.Body => bodyWithVariables.getTreeIterator.hasNext
})
@scala.annotation.tailrec
override def next(): CompoundStatementExec = state match {
case ForState.VariableAssignment =>
- variablesMap = createVariablesMapFromRow(cachedQueryResult().next())
-
- if (!areVariablesDeclared) {
- // create and execute declare var statements
- variablesMap.keys.toSeq
- .map(colName => createDeclareVarExec(colName,
variablesMap(colName)))
- .foreach(declareVarExec =>
declareVarExec.buildDataFrame(session).collect())
- areVariablesDeclared = true
- }
-
- // create and execute set var statements
- variablesMap.keys.toSeq
- .map(colName => createSetVarExec(colName, variablesMap(colName)))
- .foreach(setVarExec =>
setVarExec.buildDataFrame(session).collect())
+ val row = cachedQueryResult().next()
+
+ val variableInitStatements = row.schema.names.toSeq
+ .map { colName => (colName,
createExpressionFromValue(row.getAs(colName))) }
+ .flatMap { case (colName, expr) => Seq(
+ createDeclareVarExec(colName, expr),
+ createSetVarExec(colName, expr)
+ ) }
+
+ bodyWithVariables = new CompoundBodyExec(
+ // NoOpStatementExec appended to end of body to prevent
+ // dropping variables before last statement is executed.
Review Comment:
oh, I get it, it's basically because we have this `if` at the end of the
`Body` case:
```
if (!bodyWithVariables.getTreeIterator.hasNext) {
bodyWithVariables.exitScope()
state = ForState.VariableAssignment
}
```
we are doing this before returning the statement.
maybe we should include this in this comment related to the Noop - something
simple like "we are calling the exitScope <there> before returning the last
statement".
--
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]
