dusantism-db commented on code in PR #50026:
URL: https://github.com/apache/spark/pull/50026#discussion_r1965190136
##########
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:
If the last statement in the FOR uses one of the local variables, we have to
call drop variables after that statement has been executed. Before this PR we
achieved this with a third state, `ForState.VariableCleanup`. In this state we
passed back `DropVariableStatements` to the iterator.
```
case ForState.VariableCleanup =>
dropVariablesExec.getTreeIterator.next()
```
As we no longer use `DropVariable`, and drop variables with `exitScope`, we
don't have anything to return here. This is why NOOP is needed. To simplify the
code i removed the third state and appended the NOOP to the exec body, but the
principle is the same to something like this:
```
case ForState.VariableCleanup =>
exitScope()
NoOpStatementExec()
```
--
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]
