Skip to content

Commit

Permalink
Add variable cleanup when executing leave statement
Browse files Browse the repository at this point in the history
  • Loading branch information
miland-db committed Aug 13, 2024
1 parent bf9f409 commit ff0f330
Showing 1 changed file with 18 additions and 4 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -18,11 +18,11 @@
package org.apache.spark.sql.scripting

import scala.collection.mutable

import org.apache.spark.{SparkException, SparkThrowable}
import org.apache.spark.internal.Logging
import org.apache.spark.sql.catalyst.parser.SingleStatement
import org.apache.spark.sql.{Dataset, Row, SparkSession}
import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
import org.apache.spark.sql.catalyst.plans.logical.{DropVariable, LogicalPlan}
import org.apache.spark.sql.catalyst.trees.{Origin, WithOrigin}
import org.apache.spark.sql.errors.SqlScriptingErrors
import org.apache.spark.sql.types.BooleanType
Expand Down Expand Up @@ -224,6 +224,18 @@ class CompoundBodyExec(
statement
}

/**
* Drop variables declared in this CompoundBody.
*/
private def cleanup(): Unit = {
// Filter out internal DropVariable statements and execute them.
statements.filter(
dropVar => dropVar.isInstanceOf[SingleStatementExec]
&& dropVar.asInstanceOf[SingleStatementExec].parsedPlan.isInstanceOf[DropVariable]
&& dropVar.isInternal)
.foreach(_.asInstanceOf[SingleStatementExec].execute(session))
}

/**
* Check if the leave statement was used, if it is not used stop iterating surrounding
* [[CompoundBodyExec]] and move iterator forward. If the label of the block matches the label of
Expand All @@ -233,8 +245,10 @@ class CompoundBodyExec(
*/
private def handleLeave(leave: LeaveStatementExec): LeaveStatementExec = {
if (!leave.used) {
// Hard stop the iteration of the current begin/end block
// Hard stop the iteration of the current begin/end block.
stopIteration = true
// Cleanup variables declared in the current block.
cleanup()
// If label of the block matches the label of the leave statement,
// mark the leave statement as used. label can be None in case of a
// CompoundBody inside loop or if/else structure. In such cases,
Expand Down Expand Up @@ -300,7 +314,7 @@ class CompoundBodyExec(
handleLeave(leave)
case leafStatement: LeafStatementExec =>
// This check is done to handle error in surrounding begin/end block
// if it was not handled in the nested block
// if it was not handled in the nested block.
handleError(leafStatement)
case nonLeafStatement: NonLeafStatementExec => nonLeafStatement
}
Expand Down

0 comments on commit ff0f330

Please sign in to comment.