Add capability to inherit SessionState (SQL conf, temp tables, registered functions) when forking a new SparkSession.

Author: kunalkhamar

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/catalog/SessionCatalog.scala

@@ -1180,4 +1180,25 @@ class SessionCatalog(
     }
   }
 
+  /**
+    * Get an identical copy of the `SessionCatalog`.
+    * The temporary tables and function registry are retained.
+    * The table relation cache will not be populated.
+    */
+  override def clone: SessionCatalog = {
+    val catalog = new SessionCatalog(
+      externalCatalog,
+      globalTempViewManager,
+      functionResourceLoader,
+      functionRegistry,
+      conf,
+      hadoopConf,
+      parser)
+
+    catalog.currentDb = currentDb
+    tempTables.foreach(kv => catalog.tempTables.put(kv._1, kv._2)) // copy over temporary tables
+
+    catalog
+  }
+
 }

sql/core/src/main/scala/org/apache/spark/sql/ExperimentalMethods.scala

@@ -17,6 +17,8 @@
 
 package org.apache.spark.sql
 
+import scala.collection.mutable.ListBuffer
+
 import org.apache.spark.annotation.{Experimental, InterfaceStability}
 import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
 import org.apache.spark.sql.catalyst.rules.Rule
@@ -46,4 +48,23 @@ class ExperimentalMethods private[sql]() {
 
   @volatile var extraOptimizations: Seq[Rule[LogicalPlan]] = Nil
 
+  /**
+    * Get an identical copy of this `ExperimentalMethods` instance.
+    * @note This is used when forking a `SparkSession`.
+    * `clone` is provided here instead of implementing equivalent functionality
+    * in `SparkSession.clone` since it needs to be updated
+    * as the class `ExperimentalMethods` is extended or  modified.
+    */
+  override def clone: ExperimentalMethods = {
+    def cloneSeq[T](seq: Seq[T]): Seq[T] = {
+      val newSeq = new ListBuffer[T]
+      newSeq ++= seq
+      newSeq
+    }
+
+    val result = new ExperimentalMethods
+    result.extraStrategies = cloneSeq(extraStrategies)
+    result.extraOptimizations = cloneSeq(extraOptimizations)
+    result
+  }
 }

sql/core/src/main/scala/org/apache/spark/sql/SparkSession.scala

@@ -72,7 +72,8 @@ import org.apache.spark.util.Utils
 @InterfaceStability.Stable
 class SparkSession private(
     @transient val sparkContext: SparkContext,
-    @transient private val existingSharedState: Option[SharedState])
+    @transient private val existingSharedState: Option[SharedState],
+    existingSessionState: Option[SessionState] = None)
   extends Serializable with Closeable with Logging { self =>
 
   private[sql] def this(sc: SparkContext) {
@@ -107,9 +108,9 @@ class SparkSession private(
    */
   @transient
   private[sql] lazy val sessionState: SessionState = {
-    SparkSession.reflect[SessionState, SparkSession](
+    existingSessionState.getOrElse(SparkSession.reflect[SessionState, SparkSession](
       SparkSession.sessionStateClassName(sparkContext.conf),
-      self)
+      self))
   }
 
   /**
@@ -201,6 +202,8 @@ class SparkSession private(
   /**
    * Start a new session with isolated SQL configurations, temporary tables, registered
    * functions are isolated, but sharing the underlying `SparkContext` and cached data.
+   * If inherit is enabled, then SQL configurations, temporary tables, registered functions
+   * are copied over from parent `SparkSession`.
    *
    * @note Other than the `SparkContext`, all shared state is initialized lazily.
    * This method will force the initialization of the shared state to ensure that parent
@@ -209,8 +212,12 @@ class SparkSession private(
    *
    * @since 2.0.0
    */
-  def newSession(): SparkSession = {
-    new SparkSession(sparkContext, Some(sharedState))
+  def newSession(inheritSessionState: Boolean = false): SparkSession = {
+    if (inheritSessionState) {
+      new SparkSession(sparkContext, Some(sharedState), Some(sessionState.clone))
+    } else {
+      new SparkSession(sparkContext, Some(sharedState))
+    }
   }
 
 

sql/core/src/main/scala/org/apache/spark/sql/internal/SessionState.scala

@@ -38,16 +38,27 @@ import org.apache.spark.sql.util.ExecutionListenerManager
 
 /**
  * A class that holds all session-specific state in a given [[SparkSession]].
+ * If an `existingSessionState` is supplied, then its members will be copied over.
  */
-private[sql] class SessionState(sparkSession: SparkSession) {
+private[sql] class SessionState(
+    sparkSession: SparkSession,
+    existingSessionState: Option[SessionState] = None) {
 
   // Note: These are all lazy vals because they depend on each other (e.g. conf) and we
   // want subclasses to override some of the fields. Otherwise, we would get a lot of NPEs.
 
   /**
    * SQL-specific key-value configurations.
    */
-  lazy val conf: SQLConf = new SQLConf
+  lazy val conf: SQLConf = {
+    val result = new SQLConf
+    if (existingSessionState.nonEmpty) {
+      existingSessionState.get.conf.getAllConfs.foreach {
+        case (k, v) => if (v ne null) result.setConfString(k, v)
+      }
+    }
+    result
+  }
 
   def newHadoopConf(): Configuration = {
     val hadoopConf = new Configuration(sparkSession.sparkContext.hadoopConfiguration)
@@ -65,12 +76,29 @@ private[sql] class SessionState(sparkSession: SparkSession) {
     hadoopConf
   }
 
-  lazy val experimentalMethods = new ExperimentalMethods
+  lazy val experimentalMethods: ExperimentalMethods = {
+    existingSessionState
+      .map(_.experimentalMethods.clone)
+      .getOrElse(new ExperimentalMethods)
+  }
 
   /**
    * Internal catalog for managing functions registered by the user.
    */
-  lazy val functionRegistry: FunctionRegistry = FunctionRegistry.builtin.copy()
+  lazy val functionRegistry: FunctionRegistry = {
+    val registry = FunctionRegistry.builtin.copy()
+
+    if (existingSessionState.nonEmpty) {
+      val sourceRegistry = existingSessionState.get.functionRegistry
+      sourceRegistry
+        .listFunction()
+        .foreach(name => registry.registerFunction(
+          name,
+          sourceRegistry.lookupFunction(name).get,
+          sourceRegistry.lookupFunctionBuilder(name).get))
+    }
+    registry
+  }
 
   /**
    * A class for loading resources specified by a function.
@@ -93,14 +121,18 @@ private[sql] class SessionState(sparkSession: SparkSession) {
   /**
    * Internal catalog for managing table and database states.
    */
-  lazy val catalog = new SessionCatalog(
-    sparkSession.sharedState.externalCatalog,
-    sparkSession.sharedState.globalTempViewManager,
-    functionResourceLoader,
-    functionRegistry,
-    conf,
-    newHadoopConf(),
-    sqlParser)
+  lazy val catalog: SessionCatalog = {
+    existingSessionState
+      .map(_.catalog.clone)
+      .getOrElse(new SessionCatalog(
+        sparkSession.sharedState.externalCatalog,
+        sparkSession.sharedState.globalTempViewManager,
+        functionResourceLoader,
+        functionRegistry,
+        conf,
+        newHadoopConf(),
+        sqlParser))
+  }
 
   /**
    * Interface exposed to the user for registering user-defined functions.
@@ -165,6 +197,13 @@ private[sql] class SessionState(sparkSession: SparkSession) {
     conf.setConfString(k, v)
   }
 
+  /**
+    * Get an identical copy of the `SessionState`.
+    */
+  override def clone: SessionState = {
+    new SessionState(sparkSession, Some(this))
+  }
+
   // ------------------------------------------------------
   //  Helper methods, partially leftover from pre-2.0 days
   // ------------------------------------------------------