[SPARK-48530][SQL] Support for local variables in SQL Scripting

common/utils/src/main/resources/error/error-conditions.json

@@ -3355,6 +3355,16 @@
         "message" : [
           "Variable <varName> can only be declared at the beginning of the compound."
         ]
+      },
+      "QUALIFIED_LOCAL_VARIABLE" : {
+        "message" : [
+          "The variable <varName> must be declared without a qualifier, as qualifiers are not allowed for local variable declarations."
+        ]
+      },
+      "REPLACE_LOCAL_VARIABLE" : {
+        "message": [
+          "The variable <varName> does not support DECLARE OR REPLACE, as local variables cannot be replaced."
+        ]
       }
     },
     "sqlState" : "42K0M"
@@ -3501,6 +3511,12 @@
     ],
     "sqlState" : "42K0L"
   },
+  "LABEL_NAME_FORBIDDEN" : {
+    "message" : [
+      "The label name <label> is forbidden."
+    ],
+    "sqlState" : "42K0L"
+  },
   "LOAD_DATA_PATH_NOT_EXISTS" : {
     "message" : [
       "LOAD DATA input path does not exist: <path>."

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ColumnResolutionHelper.scala

@@ -251,7 +251,9 @@ trait ColumnResolutionHelper extends Logging with DataTypeErrorsBase {
     }
   }
 
-  def lookupVariable(nameParts: Seq[String]): Option[VariableReference] = {
+  def lookupVariable(
+      nameParts: Seq[String],
+      sessionVariablesOnly: Boolean = false): Option[VariableReference] = {
     // The temp variables live in `SYSTEM.SESSION`, and the name can be qualified or not.
     def maybeTempVariableName(nameParts: Seq[String]): Boolean = {
       nameParts.length == 1 || {
@@ -266,22 +268,40 @@ trait ColumnResolutionHelper extends Logging with DataTypeErrorsBase {
       }
     }
 
-    if (maybeTempVariableName(nameParts)) {
-      val variableName = if (conf.caseSensitiveAnalysis) {
-        nameParts.last
-      } else {
-        nameParts.last.toLowerCase(Locale.ROOT)
-      }
-      catalogManager.tempVariableManager.get(variableName).map { varDef =>
+    val namePartsCaseAdjusted = if (conf.caseSensitiveAnalysis) {
+      nameParts
+    } else {
+      nameParts.map(_.toLowerCase(Locale.ROOT))
+    }
+
+    catalogManager.scriptingLocalVariableManager
+      // If sessionOnly is set to true lookup only session variables.
+      .filterNot(_ => sessionVariablesOnly)
+      // If variable name is qualified with system.session.<varName> treat it as a session variable.
+      .filterNot(_ => nameParts.take(2).map(_.toLowerCase(Locale.ROOT)) == Seq("system", "session"))
+      // Local variable must be in format <varName> or <label>.<varName>
+      .filter(_ => namePartsCaseAdjusted.nonEmpty && namePartsCaseAdjusted.length <= 2)
+      .flatMap(_.get(namePartsCaseAdjusted))
+      .map { varDef =>
         VariableReference(
           nameParts,
           FakeSystemCatalog,
-          Identifier.of(Array(CatalogManager.SESSION_NAMESPACE), variableName),
+          Identifier.of(Array(varDef.identifier.namespace().last), namePartsCaseAdjusted.last),
           varDef)
       }
-    } else {
-      None
-    }
+      .orElse(
+        Option.when(maybeTempVariableName(nameParts)) {
+          catalogManager.tempVariableManager
+            .get(namePartsCaseAdjusted)
+            .map { varDef =>
+              VariableReference(
+                nameParts,
+                FakeSystemCatalog,
+                Identifier.of(Array(CatalogManager.SESSION_NAMESPACE), namePartsCaseAdjusted.last),
+                varDef)
+            }
+        }.flatten
+      )
   }
 
   // Resolves `UnresolvedAttribute` to its value.

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveCatalogs.scala

@@ -19,9 +19,11 @@ package org.apache.spark.sql.catalyst.analysis
 
 import scala.jdk.CollectionConverters._
 
+import org.apache.spark.sql.AnalysisException
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.connector.catalog.{CatalogManager, CatalogPlugin, Identifier, LookupCatalog, SupportsNamespaces}
+import org.apache.spark.sql.errors.DataTypeErrors.toSQLId
 import org.apache.spark.sql.errors.QueryCompilationErrors
 import org.apache.spark.util.ArrayImplicits._
 
@@ -35,10 +37,28 @@ class ResolveCatalogs(val catalogManager: CatalogManager)
     // We only support temp variables for now and the system catalog is not properly implemented
     // yet. We need to resolve `UnresolvedIdentifier` for variable commands specially.
     case c @ CreateVariable(UnresolvedIdentifier(nameParts, _), _, _) =>
-      val resolved = resolveVariableName(nameParts)
+      // From scripts we can only create local variables, which must be unqualified,
+      // and must not be DECLARE OR REPLACE.
+      if (catalogManager.scriptingLocalVariableManager.isDefined) {
+        // TODO [SPARK-50785]: Uncomment this when For Statement starts properly using local vars.
+//        if (c.replace) {
+//          throw new AnalysisException(
+//            "INVALID_VARIABLE_DECLARATION.REPLACE_LOCAL_VARIABLE",
+//            Map("varName" -> toSQLId(nameParts))
+//          )
+//        }
+
+        if (nameParts.length != 1) {
+          throw new AnalysisException(
+            "INVALID_VARIABLE_DECLARATION.QUALIFIED_LOCAL_VARIABLE",
+            Map("varName" -> toSQLId(nameParts)))
+        }
+      }
+
+      val resolved = resolveCreateVariableName(nameParts)
       c.copy(name = resolved)
     case d @ DropVariable(UnresolvedIdentifier(nameParts, _), _) =>
-      val resolved = resolveVariableName(nameParts)
+      val resolved = resolveDropVariableName(nameParts)
       d.copy(name = resolved)
 
     case UnresolvedIdentifier(nameParts, allowTemp) =>
@@ -73,28 +93,39 @@ class ResolveCatalogs(val catalogManager: CatalogManager)
     }
   }
 
-  private def resolveVariableName(nameParts: Seq[String]): ResolvedIdentifier = {
-    def ident: Identifier = Identifier.of(Array(CatalogManager.SESSION_NAMESPACE), nameParts.last)
-    if (nameParts.length == 1) {
+  private def resolveCreateVariableName(nameParts: Seq[String]): ResolvedIdentifier = {
+    val ident = catalogManager.scriptingLocalVariableManager
+      .getOrElse(catalogManager.tempVariableManager)
+      .createIdentifier(nameParts.last)
+
+    resolveVariableName(nameParts, ident)
+  }
+
+  private def resolveDropVariableName(nameParts: Seq[String]): ResolvedIdentifier = {
+    // Only session variables can be dropped, so catalogManager.scriptingLocalVariableManager
+    // is not checked in the case of DropVariable.
+    val ident = catalogManager.tempVariableManager.createIdentifier(nameParts.last)
+    resolveVariableName(nameParts, ident)
+  }
+
+  private def resolveVariableName(
+      nameParts: Seq[String],
+      ident: Identifier): ResolvedIdentifier = nameParts.length match {
+    case 1 => ResolvedIdentifier(FakeSystemCatalog, ident)
+
+    // On declare variable, local variables support only unqualified names.
+    // On drop variable, local variables are not supported at all.
+    case 2 if nameParts.head.equalsIgnoreCase(CatalogManager.SESSION_NAMESPACE) =>
       ResolvedIdentifier(FakeSystemCatalog, ident)
-    } else if (nameParts.length == 2) {
-      if (nameParts.head.equalsIgnoreCase(CatalogManager.SESSION_NAMESPACE)) {
-        ResolvedIdentifier(FakeSystemCatalog, ident)
-      } else {
-        throw QueryCompilationErrors.unresolvedVariableError(
-          nameParts, Seq(CatalogManager.SYSTEM_CATALOG_NAME, CatalogManager.SESSION_NAMESPACE))
-      }
-    } else if (nameParts.length == 3) {
-      if (nameParts(0).equalsIgnoreCase(CatalogManager.SYSTEM_CATALOG_NAME) &&
-        nameParts(1).equalsIgnoreCase(CatalogManager.SESSION_NAMESPACE)) {
-        ResolvedIdentifier(FakeSystemCatalog, ident)
-      } else {
-        throw QueryCompilationErrors.unresolvedVariableError(
-          nameParts, Seq(CatalogManager.SYSTEM_CATALOG_NAME, CatalogManager.SESSION_NAMESPACE))
-      }
-    } else {
+
+    // When there are 3 nameParts the variable must be a fully qualified session variable
+    // i.e. "system.session.<varName>"
+    case 3 if nameParts(0).equalsIgnoreCase(CatalogManager.SYSTEM_CATALOG_NAME) &&
+      nameParts(1).equalsIgnoreCase(CatalogManager.SESSION_NAMESPACE) =>
+      ResolvedIdentifier(FakeSystemCatalog, ident)
+
+    case _ =>
       throw QueryCompilationErrors.unresolvedVariableError(
-        nameParts, Seq(CatalogManager.SYSTEM_CATALOG_NAME, CatalogManager.SESSION_NAMESPACE))
-    }
+        nameParts, Seq(CatalogManager.SYSTEM_CATALOG_NAME, ident.namespace().head))
   }
 }

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveSetVariable.scala

@@ -53,11 +53,12 @@ class ResolveSetVariable(val catalogManager: CatalogManager) extends Rule[Logica
       // Names are normalized when the variables are created.
       // No need for case insensitive comparison here.
       // TODO: we need to group by the qualified variable name once other catalogs support it.
-      val dups = resolvedVars.groupBy(_.identifier.name).filter(kv => kv._2.length > 1)
+      val dups = resolvedVars.groupBy(_.identifier).filter(kv => kv._2.length > 1)
       if (dups.nonEmpty) {
         throw new AnalysisException(
           errorClass = "DUPLICATE_ASSIGNMENTS",
-          messageParameters = Map("nameList" -> dups.keys.map(toSQLId).mkString(", ")))
+          messageParameters = Map("nameList" ->
+            dups.keys.map(key => toSQLId(key.name())).mkString(", ")))
       }
 
       setVariable.copy(targetVariables = resolvedVars)

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/executeImmediate.scala

@@ -185,7 +185,7 @@ class SubstituteExecuteImmediate(val catalogManager: CatalogManager)
   }
 
   private def getVariableReference(expr: Expression, nameParts: Seq[String]): VariableReference = {
-    lookupVariable(nameParts) match {
+    lookupVariable(nameParts, sessionVariablesOnly = true) match {
       case Some(variable) => variable
       case _ =>
         throw QueryCompilationErrors

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/ResolverGuard.scala

@@ -209,7 +209,8 @@ class ResolverGuard(catalogManager: CatalogManager) extends SQLConfHelper {
     // Case-sensitive inference is not supported for Hive table schema.
     conf.caseSensitiveInferenceMode == HiveCaseSensitiveInferenceMode.NEVER_INFER
 
-  private def checkVariables() = catalogManager.tempVariableManager.isEmpty
+  private def checkVariables() = catalogManager.tempVariableManager.isEmpty &&
+      catalogManager.scriptingLocalVariableManager.forall(_.isEmpty)
 }
 
 object ResolverGuard {

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/catalog/TempVariableManager.scala → sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/catalog/VariableManager.scala

@@ -23,8 +23,39 @@ import scala.collection.mutable
 
 import org.apache.spark.sql.AnalysisException
 import org.apache.spark.sql.catalyst.expressions.Literal
+import org.apache.spark.sql.connector.catalog.{CatalogManager, Identifier}
 import org.apache.spark.sql.connector.catalog.CatalogManager.{SESSION_NAMESPACE, SYSTEM_CATALOG_NAME}
 import org.apache.spark.sql.errors.DataTypeErrorsBase
+import org.apache.spark.sql.errors.QueryCompilationErrors.unresolvedVariableError
+
+trait VariableManager {
+  def create(
+    identifier: Identifier,
+    defaultValueSQL: String,
+    initValue: Literal,
+    overrideIfExists: Boolean): Unit
+
+  def set(
+    nameParts: Seq[String],
+    defaultValueSQL: String,
+    initValue: Literal,
+    identifier: Identifier): Unit
+
+  def get(nameParts: Seq[String]): Option[VariableDefinition]
+
+  def remove(name: String): Boolean
+
+  def createIdentifier(name: String): Identifier
+
+  def clear(): Unit
+
+  def isEmpty: Boolean
+}
+
+case class VariableDefinition(
+  identifier: Identifier,
+  defaultValueSQL: String,
+  currentValue: Literal)
 
 /**
  * A thread-safe manager for temporary SQL variables (that live in the schema `SYSTEM.SESSION`),
@@ -33,40 +64,55 @@ import org.apache.spark.sql.errors.DataTypeErrorsBase
  * Note that, the variable name is always case-sensitive here, callers are responsible to format the
  * variable name w.r.t. case-sensitive config.
  */
-class TempVariableManager extends DataTypeErrorsBase {
+class TempVariableManager extends VariableManager with DataTypeErrorsBase {
 
   @GuardedBy("this")
   private val variables = new mutable.HashMap[String, VariableDefinition]
 
-  def create(
-      name: String,
+  override def create(
+      identifier: Identifier,
       defaultValueSQL: String,
       initValue: Literal,
       overrideIfExists: Boolean): Unit = synchronized {
+    val name = identifier.name()
     if (!overrideIfExists && variables.contains(name)) {
       throw new AnalysisException(
         errorClass = "VARIABLE_ALREADY_EXISTS",
         messageParameters = Map(
           "variableName" -> toSQLId(Seq(SYSTEM_CATALOG_NAME, SESSION_NAMESPACE, name))))
     }
-    variables.put(name, VariableDefinition(defaultValueSQL, initValue))
+    variables.put(name, VariableDefinition(identifier, defaultValueSQL, initValue))
   }
 
-  def get(name: String): Option[VariableDefinition] = synchronized {
-    variables.get(name)
+  override def set(
+      nameParts: Seq[String],
+      defaultValueSQL: String,
+      initValue: Literal,
+      identifier: Identifier): Unit = synchronized {
+    val name = nameParts.last
+    // Sanity check as this is already checked in ResolveSetVariable.
+    if (!variables.contains(name)) {
+      throw unresolvedVariableError(nameParts, Seq("SYSTEM", "SESSION"))
+    }
+    variables.put(name, VariableDefinition(identifier, defaultValueSQL, initValue))
   }
 
-  def remove(name: String): Boolean = synchronized {
+  override def get(nameParts: Seq[String]): Option[VariableDefinition] = synchronized {
+    variables.get(nameParts.last)
+  }
+
+  override def remove(name: String): Boolean = synchronized {
     variables.remove(name).isDefined
   }
 
-  def clear(): Unit = synchronized {
+  override def createIdentifier(name: String): Identifier =
+    Identifier.of(Array(CatalogManager.SESSION_NAMESPACE), name)
+
+  override def clear(): Unit = synchronized {
     variables.clear()
   }
 
-  def isEmpty: Boolean = synchronized {
+  override def isEmpty: Boolean = synchronized {
     variables.isEmpty
   }
 }
-
-case class VariableDefinition(defaultValueSQL: String, currentValue: Literal)

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/ParserUtils.scala

@@ -19,7 +19,8 @@ package org.apache.spark.sql.catalyst.parser
 import java.util
 import java.util.Locale
 
-import scala.collection.mutable.Set
+import scala.collection.immutable
+import scala.collection.mutable
 
 import org.antlr.v4.runtime.{ParserRuleContext, Token}
 import org.antlr.v4.runtime.misc.Interval
@@ -142,7 +143,7 @@ object ParserUtils extends SparkParserUtils {
 
 class SqlScriptingLabelContext {
   /** Set to keep track of labels seen so far */
-  private val seenLabels = Set[String]()
+  private val seenLabels = mutable.Set[String]()
 
   /**
    * Check if the beginLabelCtx and endLabelCtx match.
@@ -203,6 +204,11 @@ class SqlScriptingLabelContext {
       // Do not add the label to the seenLabels set if it is not defined.
       java.util.UUID.randomUUID.toString.toLowerCase(Locale.ROOT)
     }
+    if (SqlScriptingLabelContext.forbiddenLabelNames.contains(labelText.toLowerCase(Locale.ROOT))) {
+      withOrigin(beginLabelCtx.get) {
+        throw SqlScriptingErrors.labelNameForbidden(CurrentOrigin.get, labelText)
+      }
+    }
     labelText
   }
 
@@ -216,3 +222,7 @@ class SqlScriptingLabelContext {
     }
   }
 }
+
+object SqlScriptingLabelContext {
+  val forbiddenLabelNames: immutable.Set[String] = immutable.Set("system", "session")
+}

sql/catalyst/src/main/scala/org/apache/spark/sql/connector/catalog/CatalogManager.scala

@@ -21,7 +21,7 @@ import scala.collection.mutable
 
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.catalyst.SQLConfHelper
-import org.apache.spark.sql.catalyst.catalog.{SessionCatalog, TempVariableManager}
+import org.apache.spark.sql.catalyst.catalog.{SessionCatalog, TempVariableManager, VariableManager}
 import org.apache.spark.sql.catalyst.util.StringUtils
 import org.apache.spark.sql.errors.QueryCompilationErrors
 import org.apache.spark.sql.internal.SQLConf
@@ -49,6 +49,9 @@ class CatalogManager(
   // TODO: create a real SYSTEM catalog to host `TempVariableManager` under the SESSION namespace.
   val tempVariableManager: TempVariableManager = new TempVariableManager
 
+  // This field will be populated and cleaned up by SqlScriptingExecution.
+  var scriptingLocalVariableManager: Option[VariableManager] = None
+
   def catalog(name: String): CatalogPlugin = synchronized {
     if (name.equalsIgnoreCase(SESSION_CATALOG_NAME)) {
       v2SessionCatalog