rewrite some functions and classes

Author: yinxusen

mllib/src/main/scala/org/apache/spark/mllib/clustering/KMeansModel.scala

@@ -22,9 +22,8 @@ import org.json4s.JsonDSL._
 import org.json4s.jackson.JsonMethods._
 
 import org.apache.spark.api.java.JavaRDD
-import org.apache.spark.mllib.linalg._
+import org.apache.spark.mllib.linalg.Vector
 import org.apache.spark.mllib.util.{Loader, Saveable}
-import org.apache.spark.mllib.util.Loader._
 import org.apache.spark.rdd.RDD
 import org.apache.spark.SparkContext
 import org.apache.spark.sql.SQLContext
@@ -79,11 +78,11 @@ object KMeansModel extends Loader[KMeansModel] {
     KMeansModel.SaveLoadV1_0.load(sc, path)
   }
 
-  case class IndexedPoint(id: Int, point: Vector)
+  private case class Cluster(id: Int, point: Vector)
 
-  object IndexedPoint {
-    def apply(r: Row): IndexedPoint = {
-      IndexedPoint(r.getInt(0), r.getAs[Vector](1))
+  private object Cluster {
+    def apply(r: Row): Cluster = {
+      Cluster(r.getInt(0), r.getAs[Vector](1))
     }
   }
 
@@ -102,21 +101,21 @@ object KMeansModel extends Loader[KMeansModel] {
         ("class" -> thisClassName) ~ ("version" -> thisFormatVersion) ~ ("k" -> model.k)))
       sc.parallelize(Seq(metadata), 1).saveAsTextFile(Loader.metadataPath(path))
       val dataRDD = sc.parallelize(model.clusterCenters.zipWithIndex).map { case (point, id) =>
-        IndexedPoint(id, point)
+        Cluster(id, point)
       }.toDF()
       dataRDD.saveAsParquetFile(Loader.dataPath(path))
     }
 
     def load(sc: SparkContext, path: String): KMeansModel = {
       implicit val formats = DefaultFormats
       val sqlContext = new SQLContext(sc)
-      val (className, formatVersion, metadata) = loadMetadata(sc, path)
+      val (className, formatVersion, metadata) = Loader.loadMetadata(sc, path)
       assert(className == thisClassName)
       assert(formatVersion == thisFormatVersion)
       val k = (metadata \ "k").extract[Int]
-      val centriods = sqlContext.parquetFile(dataPath(path))
-      Loader.checkSchema[IndexedPoint](centriods.schema)
-      val localCentriods = centriods.map(IndexedPoint.apply).collect()
+      val centriods = sqlContext.parquetFile(Loader.dataPath(path))
+      Loader.checkSchema[Cluster](centriods.schema)
+      val localCentriods = centriods.map(Cluster.apply).collect()
       assert(k == localCentriods.size)
       new KMeansModel(localCentriods.sortBy(_.id).map(_.point))
     }

mllib/src/test/scala/org/apache/spark/mllib/clustering/KMeansSuite.scala

@@ -269,7 +269,7 @@ class KMeansSuite extends FunSuite with MLlibTestSparkContext {
       try {
         model.save(sc, path)
         val sameModel = KMeansModel.load(sc, path)
-        KMeansSuite.checkEqual(model, sameModel, selector)
+        KMeansSuite.checkEqual(model, sameModel)
       } finally {
         Utils.deleteRecursively(tempDir)
       }
@@ -288,22 +288,15 @@ object KMeansSuite extends FunSuite {
     new KMeansModel(Array.fill[Vector](k)(singlePoint))
   }
 
-  def checkEqual(a: KMeansModel, b: KMeansModel, isSparse: Boolean): Unit = {
+  def checkEqual(a: KMeansModel, b: KMeansModel): Unit = {
     assert(a.k === b.k)
-    isSparse match {
-      case true =>
-        a.clusterCenters.zip(b.clusterCenters).foreach { case (pointA, pointB) =>
-          assert(pointA.asInstanceOf[SparseVector].size === pointB.asInstanceOf[SparseVector].size)
-          assert(
-            pointA.asInstanceOf[SparseVector].indices === pointB.asInstanceOf[SparseVector].indices)
-          assert(
-            pointA.asInstanceOf[SparseVector].values === pointB.asInstanceOf[SparseVector].values)
-        }
+    a.clusterCenters.zip(b.clusterCenters).foreach {
+      case (ca: SparseVector, cb: SparseVector) =>
+        assert(ca === cb)
+      case (ca: DenseVector, cb: DenseVector) =>
+        assert(ca === cb)
       case _ =>
-        a.clusterCenters.zip(b.clusterCenters).foreach { case (pointA, pointB) =>
-          assert(
-            pointA.asInstanceOf[DenseVector].toArray === pointB.asInstanceOf[DenseVector].toArray)
-        }
+        throw new AssertionError("checkEqual failed since the two clusters were not identical.\n")
     }
   }
 }