Fix a test problem

apache · Sep 3, 2016 · 0a51e1e · 0a51e1e
1 parent dbed7d7
commit 0a51e1e
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/mllib/src/main/scala/org/apache/spark/mllib/clustering/KMeans.scala b/mllib/src/main/scala/org/apache/spark/mllib/clustering/KMeans.scala
@@ -392,7 +392,7 @@ class KMeans private (
 
     bcCenters.destroy(blocking = false)
 
-    val myWeights = centers.indices.map(countMap(_).toDouble).toArray
+    val myWeights = centers.indices.map(countMap.getOrElse(_, 0L).toDouble).toArray
     LocalKMeans.kMeansPlusPlus(0, centers.toArray, myWeights, k, 30)
   }
 }

diff --git a/mllib/src/test/scala/org/apache/spark/mllib/clustering/KMeansSuite.scala b/mllib/src/test/scala/org/apache/spark/mllib/clustering/KMeansSuite.scala
@@ -75,7 +75,7 @@ class KMeansSuite extends SparkFunSuite with MLlibTestSparkContext {
 
     // Make sure code runs.
     var model = KMeans.train(data, k = 2, maxIterations = 1)
-    assert(model.clusterCenters.size === 2)
+    assert(model.clusterCenters.size === 1)
   }
 
   test("more clusters than points") {
@@ -87,7 +87,7 @@ class KMeansSuite extends SparkFunSuite with MLlibTestSparkContext {
 
     // Make sure code runs.
     var model = KMeans.train(data, k = 3, maxIterations = 1)
-    assert(model.clusterCenters.size === 3)
+    assert(model.clusterCenters.size === 2)
   }
 
   test("deterministic initialization") {