[SPARK-22126][ML][WIP] Fix model-specific optimization support for ML tuning

mllib/src/main/scala/org/apache/spark/ml/Estimator.scala

@@ -18,10 +18,13 @@
 package org.apache.spark.ml
 
 import scala.annotation.varargs
+import scala.concurrent.{ExecutionContext, Future}
+import scala.concurrent.duration.Duration
 
 import org.apache.spark.annotation.{DeveloperApi, Since}
 import org.apache.spark.ml.param.{ParamMap, ParamPair}
 import org.apache.spark.sql.Dataset
+import org.apache.spark.util.ThreadUtils
 
 /**
  * :: DeveloperApi ::
@@ -82,5 +85,32 @@ abstract class Estimator[M <: Model[M]] extends PipelineStage {
     paramMaps.map(fit(dataset, _))
   }
 
+  @Since("2.3.0")
+  def fit(dataset: Dataset[_], paramMaps: Array[ParamMap],
+    unpersistDatasetAfterFitting: Boolean, executionContext: ExecutionContext,
+    modelCallback: (Model[_], ParamMap, Int) => Unit
+    ): Unit = {
+    // Fit models in a Future for training in parallel
+    val modelFutures = paramMaps.map { paramMap =>
+      Future[Model[_]] {
+        fit(dataset, paramMap).asInstanceOf[Model[_]]
+      } (executionContext)
+    }
+
+    if (unpersistDatasetAfterFitting) {
+      // Unpersist training data only when all models have trained
+      Future.sequence[Model[_], Iterable](modelFutures)(implicitly, executionContext)
+        .onComplete { _ => dataset.unpersist() }(executionContext)
+    }
+
+    val modelCallbackFutures = modelFutures.zipWithIndex.map {
+      case (modelFuture, paramMapIndex) =>
+        modelFuture.map { model =>
+          modelCallback(model, paramMaps(paramMapIndex), paramMapIndex)
+        }(executionContext)
+    }
+    modelCallbackFutures.map(ThreadUtils.awaitResult(_, Duration.Inf))
+  }
+
   override def copy(extra: ParamMap): Estimator[M]
 }

mllib/src/main/scala/org/apache/spark/ml/tuning/CrossValidator.scala

@@ -124,30 +124,16 @@ class CrossValidator @Since("1.2.0") (@Since("1.4.0") override val uid: String)
       val validationDataset = sparkSession.createDataFrame(validation, schema).cache()
       logDebug(s"Train split $splitIndex with multiple sets of parameters.")
 
-      // Fit models in a Future for training in parallel
-      val modelFutures = epm.map { paramMap =>
-        Future[Model[_]] {
-          val model = est.fit(trainingDataset, paramMap)
-          model.asInstanceOf[Model[_]]
-        } (executionContext)
-      }
-
-      // Unpersist training data only when all models have trained
-      Future.sequence[Model[_], Iterable](modelFutures)(implicitly, executionContext)
-        .onComplete { _ => trainingDataset.unpersist() } (executionContext)
-
-      // Evaluate models in a Future that will calulate a metric and allow model to be cleaned up
-      val foldMetricFutures = modelFutures.zip(epm).map { case (modelFuture, paramMap) =>
-        modelFuture.map { model =>
+      val foldMetrics = new Array[Double](epm.length)
+      est.fit(trainingDataset, epm, true, executionContext,
+        (model: Model[_], paramMap: ParamMap, paramMapIndex: Int) => {
           // TODO: duplicate evaluator to take extra params from input
           val metric = eval.evaluate(model.transform(validationDataset, paramMap))
           logDebug(s"Got metric $metric for model trained with $paramMap.")
-          metric
-        } (executionContext)
-      }
+          foldMetrics(paramMapIndex) = metric
+        }
+      )
 
-      // Wait for metrics to be calculated before unpersisting validation dataset
-      val foldMetrics = foldMetricFutures.map(ThreadUtils.awaitResult(_, Duration.Inf))
       validationDataset.unpersist()
       foldMetrics
     }.transpose.map(_.sum / $(numFolds)) // Calculate average metric over all splits

mllib/src/main/scala/org/apache/spark/ml/tuning/TrainValidationSplit.scala

@@ -123,29 +123,16 @@ class TrainValidationSplit @Since("1.5.0") (@Since("1.5.0") override val uid: St
 
     // Fit models in a Future for training in parallel
     logDebug(s"Train split with multiple sets of parameters.")
-    val modelFutures = epm.map { paramMap =>
-      Future[Model[_]] {
-        val model = est.fit(trainingDataset, paramMap)
-        model.asInstanceOf[Model[_]]
-      } (executionContext)
-    }
-
-    // Unpersist training data only when all models have trained
-    Future.sequence[Model[_], Iterable](modelFutures)(implicitly, executionContext)
-      .onComplete { _ => trainingDataset.unpersist() } (executionContext)
 
-    // Evaluate models in a Future that will calulate a metric and allow model to be cleaned up
-    val metricFutures = modelFutures.zip(epm).map { case (modelFuture, paramMap) =>
-      modelFuture.map { model =>
+    val metrics = new Array[Double](epm.length)
+    est.fit(trainingDataset, epm, true, executionContext,
+      (model: Model[_], paramMap: ParamMap, paramMapIndex: Int) => {
         // TODO: duplicate evaluator to take extra params from input
         val metric = eval.evaluate(model.transform(validationDataset, paramMap))
         logDebug(s"Got metric $metric for model trained with $paramMap.")
-        metric
-      } (executionContext)
-    }
-
-    // Wait for all metrics to be calculated
-    val metrics = metricFutures.map(ThreadUtils.awaitResult(_, Duration.Inf))
+        metrics(paramMapIndex) = metric
+      }
+    )
 
     // Unpersist validation set once all metrics have been produced
     validationDataset.unpersist()