davidm-db commented on code in PR #49006:
URL: https://github.com/apache/spark/pull/49006#discussion_r1871033655
##########
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:
I think we will need more than peeking at the top value, for example when
resolving variables we have to go through all the scopes to try to find the
variable.
Not sure for the frames at the moment, but I think it's fine to keep it like
this and improve later if possible.
--
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]
