improve unit test

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

@@ -23,21 +23,22 @@ import org.apache.spark.sql.types.{StringType, StructField, StructType}
 
 /**
  * :: Experimental ::
- * Summary of grid search tuning.
+ * Summary for the grid search tuning.
  *
- * @param params  estimator param maps
- * @param metrics  Corresponding evaluation metrics for the param maps
+ * @param params  ParamMaps for the Estimator
+ * @param metrics  corresponding evaluation metrics for the params
+ * @param bestIndex  index in params for the ParamMap of the best model.
  */
 @Since("2.3.0")
 @Experimental
 private[tuning] class TuningSummary private[tuning](
-    val params: Array[ParamMap],
-    val metrics: Array[Double],
-    val bestIndex: Int) {
+    private[tuning] val params: Array[ParamMap],
+    private[tuning] val metrics: Array[Double],
+    private[tuning] val bestIndex: Int) {
 
   /**
    * Summary of grid search tuning in the format of DataFrame. Each row contains one candidate
-   * paramMap and its corresponding metrics.
+   * paramMap and its corresponding metric.
    */
   def trainingMetrics: DataFrame = {
     require(params.nonEmpty, "estimator param maps should not be empty")

mllib/src/test/scala/org/apache/spark/ml/tuning/CrossValidatorSuite.scala

@@ -29,7 +29,7 @@ import org.apache.spark.ml.param.shared.HasInputCol
 import org.apache.spark.ml.regression.LinearRegression
 import org.apache.spark.ml.util.{DefaultReadWriteTest, MLTestingUtils}
 import org.apache.spark.mllib.util.{LinearDataGenerator, MLlibTestSparkContext}
-import org.apache.spark.sql.Dataset
+import org.apache.spark.sql.{Dataset, Row}
 import org.apache.spark.sql.types.StructType
 
 class CrossValidatorSuite
@@ -82,16 +82,11 @@ class CrossValidatorSuite
     assert(cvModel.hasSummary)
     assert(cvModel.summary.params === lrParamMaps)
     assert(cvModel.summary.trainingMetrics.count() === lrParamMaps.length)
-    val expectedSummary = spark.createDataFrame(Seq(
-      (0, 0.001),
-      (2, 0.001),
-      (0, 1.0),
-      (2, 1.0),
-      (0, 1000.0),
-      (2, 1000.0)
-    ).map(t => (t._1.toString, t._2.toString))).toDF("maxIter", "regParam")
-    assert(cvModel.summary.trainingMetrics.select("maxIter", "regParam").collect().toSet
-      .equals(expectedSummary.collect().toSet))
+
+    val expected = lrParamMaps.zip(cvModel.avgMetrics).map { case (map, metric) =>
+      Row.fromSeq(map.toSeq.sortBy(_.param.name).map(_.value.toString) ++ Seq(metric.toString))
+    }
+    assert(cvModel.summary.trainingMetrics.collect().toSet === expected.toSet)
   }
 
   test("cross validation with linear regression") {

mllib/src/test/scala/org/apache/spark/ml/tuning/TrainValidationSplitSuite.scala

@@ -23,12 +23,12 @@ import org.apache.spark.ml.classification.{LogisticRegression, LogisticRegressio
 import org.apache.spark.ml.classification.LogisticRegressionSuite.generateLogisticInput
 import org.apache.spark.ml.evaluation.{BinaryClassificationEvaluator, Evaluator, RegressionEvaluator}
 import org.apache.spark.ml.linalg.Vectors
-import org.apache.spark.ml.param.{ParamMap}
+import org.apache.spark.ml.param.ParamMap
 import org.apache.spark.ml.param.shared.HasInputCol
 import org.apache.spark.ml.regression.LinearRegression
 import org.apache.spark.ml.util.{DefaultReadWriteTest, MLTestingUtils}
 import org.apache.spark.mllib.util.{LinearDataGenerator, MLlibTestSparkContext}
-import org.apache.spark.sql.Dataset
+import org.apache.spark.sql.{Dataset, Row}
 import org.apache.spark.sql.types.StructType
 
 class TrainValidationSplitSuite
@@ -75,16 +75,11 @@ class TrainValidationSplitSuite
     assert(tvsModel.hasSummary)
     assert(tvsModel.summary.params === lrParamMaps)
     assert(tvsModel.summary.trainingMetrics.count() === lrParamMaps.length)
-    val expectedSummary = spark.createDataFrame(Seq(
-      (0, 0.001),
-      (2, 0.001),
-      (0, 1.0),
-      (2, 1.0),
-      (0, 1000.0),
-      (2, 1000.0)
-    ).map(t => (t._1.toString, t._2.toString))).toDF("maxIter", "regParam")
-    assert(tvsModel.summary.trainingMetrics.select("maxIter", "regParam").collect().toSet
-      .equals(expectedSummary.collect().toSet))
+
+    val expected = lrParamMaps.zip(tvsModel.validationMetrics).map { case (map, metric) =>
+      Row.fromSeq(map.toSeq.sortBy(_.param.name).map(_.value.toString) ++ Seq(metric.toString))
+    }
+    assert(tvsModel.summary.trainingMetrics.collect().toSet === expected.toSet)
   }
 
   test("train validation with linear regression") {