[SPARK-22538][ML] SQLTransformer should not unpersist possibly cached…

… input dataset

## What changes were proposed in this pull request?

`SQLTransformer.transform` unpersists input dataset when dropping temporary view. We should not change input dataset's cache status.

## How was this patch tested?

Added test.

Author: Liang-Chi Hsieh <viirya@gmail.com>

Closes #19772 from viirya/SPARK-22538.

mllib/src/main/scala/org/apache/spark/ml/feature/SQLTransformer.scala

@@ -70,7 +70,8 @@ class SQLTransformer @Since("1.6.0") (@Since("1.6.0") override val uid: String)
     dataset.createOrReplaceTempView(tableName)
     val realStatement = $(statement).replace(tableIdentifier, tableName)
     val result = dataset.sparkSession.sql(realStatement)
-    dataset.sparkSession.catalog.dropTempView(tableName)
+    // Call SessionCatalog.dropTempView to avoid unpersisting the possibly cached dataset.
+    dataset.sparkSession.sessionState.catalog.dropTempView(tableName)
     result
   }
 

mllib/src/test/scala/org/apache/spark/ml/feature/SQLTransformerSuite.scala

@@ -22,6 +22,7 @@ import org.apache.spark.ml.param.ParamsSuite
 import org.apache.spark.ml.util.DefaultReadWriteTest
 import org.apache.spark.mllib.util.MLlibTestSparkContext
 import org.apache.spark.sql.types.{LongType, StructField, StructType}
+import org.apache.spark.storage.StorageLevel
 
 class SQLTransformerSuite
   extends SparkFunSuite with MLlibTestSparkContext with DefaultReadWriteTest {
@@ -60,4 +61,15 @@ class SQLTransformerSuite
     val expected = StructType(Seq(StructField("id1", LongType, nullable = false)))
     assert(outputSchema === expected)
   }
+
+  test("SPARK-22538: SQLTransformer should not unpersist given dataset") {
+    val df = spark.range(10)
+    df.cache()
+    df.count()
+    assert(df.storageLevel != StorageLevel.NONE)
+    new SQLTransformer()
+      .setStatement("SELECT id + 1 AS id1 FROM __THIS__")
+      .transform(df)
+    assert(df.storageLevel != StorageLevel.NONE)
+  }
 }