address comments

markhamstra · Apr 20, 2015 · f80ec1b · f80ec1b
1 parent 49600cc
commit f80ec1b
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 deletions.
diff --git a/mllib/src/main/scala/org/apache/spark/mllib/regression/IsotonicRegression.scala b/mllib/src/main/scala/org/apache/spark/mllib/regression/IsotonicRegression.scala
@@ -143,7 +143,7 @@ object IsotonicRegressionModel extends Loader[IsotonicRegressionModel] {
 
   import org.apache.spark.mllib.util.Loader._
 
-  private  object SaveLoadV1_0 {
+  private object SaveLoadV1_0 {
 
     def thisFormatVersion: String = "1.0"
 
@@ -167,7 +167,7 @@ object IsotonicRegressionModel extends Loader[IsotonicRegressionModel] {
           ("isotonic" -> isotonic)))
       sc.parallelize(Seq(metadata), 1).saveAsTextFile(metadataPath(path))
 
-      sqlContext.createDataFrame(boundaries.toList.zip(predictions.toList)
+      sqlContext.createDataFrame(boundaries.toSeq.zip(predictions)
         .map { case (b, p) => Data(b, p) }).saveAsParquetFile(dataPath(path))
     }
 
@@ -177,8 +177,8 @@ object IsotonicRegressionModel extends Loader[IsotonicRegressionModel] {
 
       checkSchema[Data](dataRDD.schema)
       val dataArray = dataRDD.select("boundary", "prediction").collect()
-      val (boundaries, predictions) = dataArray.map { 
-        x => (x.getAs[Double](0), x.getAs[Double](1)) }.toList.unzip
+      val (boundaries, predictions) = dataArray.map {
+        x => (x.getDouble(0), x.getDouble(1)) }.toList.sortBy(_._1).unzip
       (boundaries.toArray, predictions.toArray)
     }
   }

diff --git a/mllib/src/test/scala/org/apache/spark/mllib/regression/IsotonicRegressionSuite.scala b/mllib/src/test/scala/org/apache/spark/mllib/regression/IsotonicRegressionSuite.scala
@@ -75,7 +75,9 @@ class IsotonicRegressionSuite extends FunSuite with MLlibTestSparkContext with M
   }
 
   test("model save/load") {
-    val model = runIsotonicRegression(Seq(1, 2, 3, 1, 6, 17, 16, 17, 18), true)
+    val boundaries = Array(0.0, 1.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0)
+    val predictions = Array(1, 2, 2, 6, 16.5, 16.5, 17.0, 18.0)
+    val model = new IsotonicRegressionModel(boundaries, predictions, true)
 
     val tempDir = Utils.createTempDir()
     val path = tempDir.toURI.toString