Add tests for forking new session with inherit config enabled. Update overloaded functions for Java bytecode compatibility.

Author: kunalkhamar

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

@@ -1181,10 +1181,10 @@ 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.
-    */
+   * 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,

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

@@ -49,12 +49,12 @@ 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.
-    */
+   * 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]

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

@@ -44,7 +44,7 @@ import org.apache.spark.sql.internal.{CatalogImpl, SessionState, SharedState}
 import org.apache.spark.sql.internal.StaticSQLConf.CATALOG_IMPLEMENTATION
 import org.apache.spark.sql.sources.BaseRelation
 import org.apache.spark.sql.streaming._
-import org.apache.spark.sql.types.{DataType, LongType, StructType}
+import org.apache.spark.sql.types.{DataType, StructType}
 import org.apache.spark.sql.util.ExecutionListenerManager
 import org.apache.spark.util.Utils
 
@@ -202,8 +202,6 @@ 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
@@ -212,11 +210,25 @@ class SparkSession private(
    *
    * @since 2.0.0
    */
-  def newSession(inheritSessionState: Boolean = false): SparkSession = {
+  def newSession(): SparkSession = {
+    new SparkSession(sparkContext, Some(sharedState))
+  }
+
+  /**
+   * Start a new session, sharing the underlying `SparkContext` and cached data.
+   * If inheritSessionState 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
+   * and child sessions are set up with the same shared state. If the underlying catalog
+   * implementation is Hive, this will initialize the metastore, which may take some time.
+   */
+  def newSession(inheritSessionState: Boolean): SparkSession = {
     if (inheritSessionState) {
       new SparkSession(sparkContext, Some(sharedState), Some(sessionState.clone))
     } else {
-      new SparkSession(sparkContext, Some(sharedState))
+      newSession()
     }
   }
 

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

@@ -42,7 +42,11 @@ import org.apache.spark.sql.util.ExecutionListenerManager
  */
 private[sql] class SessionState(
     sparkSession: SparkSession,
-    existingSessionState: Option[SessionState] = None) {
+    existingSessionState: Option[SessionState]) {
+
+  private[sql] def this(sparkSession: SparkSession) = {
+    this(sparkSession, 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.
@@ -198,8 +202,8 @@ private[sql] class SessionState(
   }
 
   /**
-    * Get an identical copy of the `SessionState`.
-    */
+   * Get an identical copy of the `SessionState`.
+   */
   override def clone: SessionState = {
     new SessionState(sparkSession, Some(this))
   }

sql/core/src/test/scala/org/apache/spark/sql/SparkSessionBuilderSuite.scala

@@ -18,6 +18,8 @@
 package org.apache.spark.sql
 
 import org.apache.spark.{SparkConf, SparkContext, SparkFunSuite}
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.catalyst.rules.Rule
 
 /**
  * Test cases for the builder pattern of [[SparkSession]].
@@ -123,4 +125,70 @@ class SparkSessionBuilderSuite extends SparkFunSuite {
       session.stop()
     }
   }
+
+  test("fork new session and inherit a copy of the session state") {
+    val activeSession = SparkSession.builder().master("local").getOrCreate()
+    val forkedSession = activeSession.newSession(inheritSessionState = true)
+
+    assert(forkedSession ne activeSession)
+    assert(forkedSession.sessionState ne activeSession.sessionState)
+
+    forkedSession.stop()
+    activeSession.stop()
+  }
+
+  test("fork new session and inherit sql config options") {
+    val activeSession = SparkSession
+      .builder()
+      .master("local")
+      .config("spark-configb", "b")
+      .getOrCreate()
+    val forkedSession = activeSession.newSession(inheritSessionState = true)
+
+    assert(forkedSession ne activeSession)
+    assert(forkedSession.conf ne activeSession.conf)
+    assert(forkedSession.conf.get("spark-configb") == "b")
+
+    forkedSession.stop()
+    activeSession.stop()
+  }
+
+  test("fork new session and inherit function registry and udf") {
+    val activeSession = SparkSession.builder().master("local").getOrCreate()
+    activeSession.udf.register("strlenScala", (_: String).length + (_: Int))
+    val forkedSession = activeSession.newSession(inheritSessionState = true)
+
+    assert(forkedSession ne activeSession)
+    assert(forkedSession.sessionState.functionRegistry ne
+      activeSession.sessionState.functionRegistry)
+    assert(forkedSession.sessionState.functionRegistry.lookupFunction("strlenScala").nonEmpty)
+
+    forkedSession.stop()
+    activeSession.stop()
+  }
+
+  test("fork new session and inherit experimental methods") {
+    object DummyRule1 extends Rule[LogicalPlan] {
+      def apply(p: LogicalPlan): LogicalPlan = p
+    }
+    object DummyRule2 extends Rule[LogicalPlan] {
+      def apply(p: LogicalPlan): LogicalPlan = p
+    }
+    val optimizations = List(DummyRule1, DummyRule2)
+
+    val activeSession = SparkSession.builder().master("local").getOrCreate()
+    activeSession.experimental.extraOptimizations = optimizations
+
+    val forkedSession = activeSession.newSession(inheritSessionState = true)
+
+    assert(forkedSession ne activeSession)
+    assert(forkedSession.experimental ne activeSession.experimental)
+    assert(forkedSession.experimental.extraOptimizations ne
+      activeSession.experimental.extraOptimizations)
+    assert(forkedSession.experimental.extraOptimizations.toSet ==
+      activeSession.experimental.extraOptimizations.toSet)
+
+    forkedSession.stop()
+    activeSession.stop()
+  }
 }

sql/core/src/test/scala/org/apache/spark/sql/internal/CatalogSuite.scala

@@ -493,6 +493,21 @@ class CatalogSuite
     }
   }
 
+  test("clone SessionCatalog") {
+    // need to test tempTables are cloned
+    assert(spark.catalog.listTables().collect().isEmpty)
+
+    createTempTable("my_temp_table")
+    assert(spark.catalog.listTables().collect().map(_.name).toSet == Set("my_temp_table"))
+
+    val forkedSession = spark.newSession(inheritSessionState = true)
+    assert(forkedSession.catalog.listTables().collect().map(_.name).toSet == Set("my_temp_table"))
+
+    dropTable("my_temp_table")
+    assert(spark.catalog.listTables().collect().map(_.name).toSet == Set())
+    assert(forkedSession.catalog.listTables().collect().map(_.name).toSet == Set("my_temp_table"))
+  }
+
   // TODO: add tests for the rest of them
 
 }

sql/hive/src/main/scala/org/apache/spark/sql/hive/test/TestHive.scala

@@ -32,13 +32,11 @@ import org.apache.spark.{SparkConf, SparkContext}
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.{SparkSession, SQLContext}
 import org.apache.spark.sql.catalyst.analysis._
-import org.apache.spark.sql.catalyst.analysis.FunctionRegistry.FunctionBuilder
-import org.apache.spark.sql.catalyst.expressions.ExpressionInfo
 import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
 import org.apache.spark.sql.execution.QueryExecution
 import org.apache.spark.sql.execution.command.CacheTableCommand
 import org.apache.spark.sql.hive._
-import org.apache.spark.sql.internal.{SharedState, SQLConf}
+import org.apache.spark.sql.internal.{SessionState, SharedState, SQLConf}
 import org.apache.spark.sql.internal.StaticSQLConf.CATALOG_IMPLEMENTATION
 import org.apache.spark.util.{ShutdownHookManager, Utils}
 
@@ -115,16 +113,22 @@ class TestHiveContext(
 private[hive] class TestHiveSparkSession(
     @transient private val sc: SparkContext,
     @transient private val existingSharedState: Option[SharedState],
+    existingSessionState: Option[SessionState],
     private val loadTestTables: Boolean)
   extends SparkSession(sc) with Logging { self =>
 
   def this(sc: SparkContext, loadTestTables: Boolean) {
     this(
       sc,
       existingSharedState = None,
+      existingSessionState = None,
       loadTestTables)
   }
 
+  def this(sc: SparkContext, existingSharedState: Option[SharedState], loadTestTables: Boolean) {
+    this(sc, existingSharedState, existingSessionState = None, loadTestTables)
+  }
+
   { // set the metastore temporary configuration
     val metastoreTempConf = HiveUtils.newTemporaryConfiguration(useInMemoryDerby = false) ++ Map(
       ConfVars.METASTORE_INTEGER_JDO_PUSHDOWN.varname -> "true",
@@ -151,7 +155,7 @@ private[hive] class TestHiveSparkSession(
     new TestHiveSessionState(self)
 
   override def newSession(): TestHiveSparkSession = {
-    new TestHiveSparkSession(sc, Some(sharedState), loadTestTables)
+    new TestHiveSparkSession(sc, Some(sharedState), None, loadTestTables)
   }
 
   private var cacheTables: Boolean = false