[SPARK-15644] [MLlib] [SQL] Replace SQLContext with SparkSession in MLlib

mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala

@@ -33,7 +33,7 @@ import org.apache.spark.ml.classification.{OneVsRest, OneVsRestModel}
 import org.apache.spark.ml.feature.RFormulaModel
 import org.apache.spark.ml.param.{ParamPair, Params}
 import org.apache.spark.ml.tuning.ValidatorParams
-import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.{SparkSession, SQLContext}
 import org.apache.spark.util.Utils
 
 /**
@@ -43,9 +43,18 @@ private[util] sealed trait BaseReadWrite {
   private var optionSparkSession: Option[SparkSession] = None
 
   /**
-   * Sets the Spark Session to use for saving/loading.
+   * Sets the Spark SQLContext to use for saving/loading.
    */
   @Since("1.6.0")
+  def context(sqlContext: SQLContext): this.type = {
+    optionSparkSession = Option(sqlContext.sparkSession)
+    this
+  }
+
+  /**
+   * Sets the Spark Session to use for saving/loading.
+   */
+  @Since("2.0.0")
   def context(sparkSession: SparkSession): this.type = {
     optionSparkSession = Option(sparkSession)
     this
@@ -61,6 +70,11 @@ private[util] sealed trait BaseReadWrite {
     optionSparkSession.get
   }
 
+  /**
+   * Returns the user-specified SQL context or the default.
+   */
+  protected final def sqlContext: SQLContext = sparkSession.sqlContext
+
   /** Returns the underlying [[SparkContext]]. */
   protected final def sc: SparkContext = sparkSession.sparkContext
 }
@@ -115,6 +129,9 @@ abstract class MLWriter extends BaseReadWrite with Logging {
 
   // override for Java compatibility
   override def context(sparkSession: SparkSession): this.type = super.context(sparkSession)
+
+  // override for Java compatibility
+  override def context(sqlContext: SQLContext): this.type = super.context(sqlContext)
 }
 
 /**
@@ -159,6 +176,9 @@ abstract class MLReader[T] extends BaseReadWrite {
 
   // override for Java compatibility
   override def context(sparkSession: SparkSession): this.type = super.context(sparkSession)
+
+  // override for Java compatibility
+  override def context(sqlContext: SQLContext): this.type = super.context(sqlContext)
 }
 
 /**

project/MimaExcludes.scala

@@ -778,12 +778,6 @@ object MimaExcludes {
       ) ++ Seq(
         ProblemFilters.exclude[ReversedMissingMethodProblem]("org.apache.spark.mllib.linalg.Vector.asBreeze"),
         ProblemFilters.exclude[ReversedMissingMethodProblem]("org.apache.spark.mllib.linalg.Matrix.asBreeze")
-      ) ++ Seq(
-        // [SPARK-15644] Replace SQLContext by SparkSession in MLLIB
-        ProblemFilters.exclude[DirectMissingMethodProblem]("org.apache.spark.ml.util.MLWriter.sqlContext"),
-        ProblemFilters.exclude[IncompatibleMethTypeProblem]("org.apache.spark.ml.util.MLWriter.context"),
-        ProblemFilters.exclude[DirectMissingMethodProblem]("org.apache.spark.ml.util.MLReader.sqlContext"),
-        ProblemFilters.exclude[IncompatibleMethTypeProblem]("org.apache.spark.ml.util.MLReader.context")
       )
     case v if v.startsWith("1.6") =>
       Seq(

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

@@ -109,7 +109,7 @@ class SparkSession private(
    * A wrapped version of this session in the form of a [[SQLContext]], for backward compatibility.
    */
   @transient
-  private[sql] val sqlContext: SQLContext = new SQLContext(this)
+  private[spark] val sqlContext: SQLContext = new SQLContext(this)
 
   /**
    * Runtime configuration interface for Spark.