From 0a51e1eae6baefcc7e9360ebe1cd82ac24604507 Mon Sep 17 00:00:00 2001 From: Sean Owen Date: Sat, 3 Sep 2016 11:55:57 +0100 Subject: [PATCH] Fix a test problem --- .../main/scala/org/apache/spark/mllib/clustering/KMeans.scala | 2 +- .../scala/org/apache/spark/mllib/clustering/KMeansSuite.scala | 4 ++-- 2 files changed, 3 insertions(+), 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 index e00494c322319..711d904ca2881 100644 --- 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 index 3003c62d9876c..909c5b45b9fa4 100644 --- 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") {