From 49717a72dd55eec57e393cb025c115ee4c4c2162 Mon Sep 17 00:00:00 2001
From: ludatabricks <38018689+ludatabricks@users.noreply.github.com>
Date: Wed, 13 Jun 2018 15:21:50 -0700
Subject: [PATCH] put additionalTests to mlmetrics (#153)

The time for additionalTests is missing in MLMetrics. Now add it back to MLMetrics so that we can test the time for other methods.
---
 .../sql/perf/mllib/MLPipelineStageBenchmarkable.scala      | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/src/main/scala/com/databricks/spark/sql/perf/mllib/MLPipelineStageBenchmarkable.scala b/src/main/scala/com/databricks/spark/sql/perf/mllib/MLPipelineStageBenchmarkable.scala
index 32605bd..e9ff623 100644
--- a/src/main/scala/com/databricks/spark/sql/perf/mllib/MLPipelineStageBenchmarkable.scala
+++ b/src/main/scala/com/databricks/spark/sql/perf/mllib/MLPipelineStageBenchmarkable.scala
@@ -91,10 +91,11 @@ class MLPipelineStageBenchmarkable(
       val additionalTests = test.testAdditionalMethods(param, model).map {
         tuple =>
           val (additionalMethodTime, _) = measureTime { tuple._2() }
-          tuple._1 -> additionalMethodTime.toMillis.toDouble
-      }
+          MLMetric(tuple._1, additionalMethodTime.toMillis, false)
+      }.toArray
 
-      val mlMetrics = Array(metricTrainingTime, metricTraining, metricTestTime, metricTest)
+      val mlMetrics = Array(metricTrainingTime, metricTraining, metricTestTime, metricTest) ++
+        additionalTests
       val paramsMap = params.toMap
       val benchmarkId = name.split('.').last + "_" + paramsMap.hashCode.abs