davidm-db commented on code in PR #49006:
URL: https://github.com/apache/spark/pull/49006#discussion_r1868408939
##########
sql/core/src/main/scala/org/apache/spark/sql/scripting/SqlScriptingExecution.scala:
##########
@@ -35,27 +35,49 @@ class SqlScriptingExecution(
session: SparkSession,
args: Map[String, Expression]) extends Iterator[DataFrame] {
- // Build the execution plan for the script.
- private val executionPlan: Iterator[CompoundStatementExec] =
- SqlScriptingInterpreter(session).buildExecutionPlan(sqlScript, args)
+ private val interpreter = SqlScriptingInterpreter(session)
- private var current = getNextResult
+ // Frames to keep what is being executed.
+ private val context: SqlScriptingExecutionContext = {
+ val ctx = new SqlScriptingExecutionContext()
+ interpreter.buildExecutionPlan(sqlScript, args, ctx)
+ ctx
+ }
+
+ private var current: Option[DataFrame] = None
+ private var resultConsumed: Boolean = true
- override def hasNext: Boolean = current.isDefined
+ override def hasNext: Boolean = {
+ // If the previous result was not consumed, return true if current element
exists.
+ if (!resultConsumed) {
+ return current.isDefined
+ }
+
+ // If the previous result was consumed, get the next result and return
true if it exists.
+ current = getNextResult
+ resultConsumed = false
+ current.isDefined
+ }
override def next(): DataFrame = {
if (!hasNext) throw SparkException.internalError("No more elements to
iterate through.")
- val nextDataFrame = current.get
- current = getNextResult
- nextDataFrame
+ resultConsumed = true
+ current.get
+ }
+
+ /** Helper method to iterate get next statements from the first available
frame. */
+ private def getNextStatement: Option[CompoundStatementExec] = {
+ while (context.frames.nonEmpty && !context.frames.last.hasNext) {
+ context.frames.remove(context.frames.size - 1)
+ }
+ if (context.frames.nonEmpty && context.frames.last.hasNext) {
Review Comment:
same here, just `context.frames.nonEmpty` should be enough
--
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]
