[SPARK-18596][ML] add checking and caching to bisecting kmeans

mllib/src/main/scala/org/apache/spark/ml/clustering/BisectingKMeans.scala

@@ -33,6 +33,7 @@ import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.{DataFrame, Dataset, Row}
 import org.apache.spark.sql.functions.{col, udf}
 import org.apache.spark.sql.types.{IntegerType, StructType}
+import org.apache.spark.storage.StorageLevel
 
 
 /**
@@ -255,10 +256,19 @@ class BisectingKMeans @Since("2.0.0") (
 
   @Since("2.0.0")
   override def fit(dataset: Dataset[_]): BisectingKMeansModel = {
+    val handlePersistence = dataset.storageLevel == StorageLevel.NONE
+    fit(dataset, handlePersistence)
+  }
+
+  @Since("2.2.0")
+  protected def fit(dataset: Dataset[_], handlePersistence: Boolean): BisectingKMeansModel = {
     transformSchema(dataset.schema, logging = true)
     val rdd: RDD[OldVector] = dataset.select(col($(featuresCol))).rdd.map {
       case Row(point: Vector) => OldVectors.fromML(point)
     }
+    if (handlePersistence) {
+      rdd.persist(StorageLevel.MEMORY_AND_DISK)
+    }
 
     val instr = Instrumentation.create(this, rdd)
     instr.logParams(featuresCol, predictionCol, k, maxIter, seed, minDivisibleClusterSize)
@@ -273,6 +283,9 @@ class BisectingKMeans @Since("2.0.0") (
     val summary = new BisectingKMeansSummary(
       model.transform(dataset), $(predictionCol), $(featuresCol), $(k))
     model.setSummary(Some(summary))
+    if (handlePersistence) {
+      rdd.unpersist()
+    }
     instr.logSuccess(model)
     model
   }

mllib/src/main/scala/org/apache/spark/ml/clustering/KMeans.scala

@@ -334,10 +334,10 @@ class KMeans @Since("1.5.0") (
     val summary = new KMeansSummary(
       model.transform(dataset), $(predictionCol), $(featuresCol), $(k))
     model.setSummary(Some(summary))
-    instr.logSuccess(model)
     if (handlePersistence) {
       instances.unpersist()
     }
+    instr.logSuccess(model)
     model
   }