miland-db commented on code in PR #49006:
URL: https://github.com/apache/spark/pull/49006#discussion_r1869760703
##########
sql/core/src/main/scala/org/apache/spark/sql/scripting/SqlScriptingExecution.scala:
##########
@@ -35,27 +35,50 @@ 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()
+ val executionPlan = interpreter.buildExecutionPlan(sqlScript, args, ctx)
+ ctx.frames.addOne(new SqlScriptingExecutionFrame(executionPlan))
+ 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)
Review Comment:
The idea is to have _stack like_ data structure that you can also traverse
without removing elements from it. In the next iterations we can consider other
structures if needed.
--
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]
