address feedback

Author: DB Tsai

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

@@ -23,6 +23,7 @@ import org.apache.spark.annotation.Experimental
 import org.apache.spark.Logging
 import org.apache.spark.SparkContext._
 import org.apache.spark.mllib.linalg.{Vector, Vectors}
+import org.apache.spark.mllib.linalg.BLAS.{axpy, scal}
 import org.apache.spark.mllib.util.MLUtils
 import org.apache.spark.rdd.RDD
 import org.apache.spark.storage.StorageLevel
@@ -168,15 +169,10 @@ class KMeans private (
 
     // Execute iterations of Lloyd's algorithm until all runs have converged
     while (iteration < maxIterations && !activeRuns.isEmpty) {
-      type WeightedPoint = (Array[Double], Long)
-      def mergeContribs(p1: WeightedPoint, p2: WeightedPoint): WeightedPoint = {
-        require(p1._1.size == p2._1.size)
-        var i = 0
-        while(i < p1._1.size) {
-          p1._1(i) += p2._1(i)
-          i += 1
-        }
-        (p1._1, p1._2 + p2._2)
+      type WeightedPoint = (Vector, Long)
+      def mergeContribs(x: WeightedPoint, y: WeightedPoint): WeightedPoint = {
+        axpy(1.0, x._1, y._1)
+        (y._1, x._2 + y._2)
       }
 
       val activeCenters = activeRuns.map(r => centers(r)).toArray
@@ -191,15 +187,15 @@ class KMeans private (
         val k = thisActiveCenters(0).length
         val dims = thisActiveCenters(0)(0).vector.size
 
-        val sums = Array.fill(runs, k)(Array.ofDim[Double](dims))
+        val sums = Array.fill(runs, k)(Vectors.zeros(dims))
         val counts = Array.fill(runs, k)(0L)
 
         points.foreach { point =>
           (0 until runs).foreach { i =>
             val (bestCenter, cost) = KMeans.findClosest(thisActiveCenters(i), point)
             costAccums(i) += cost
             val sum = sums(i)(bestCenter)
-            point.vector.foreachActive((index, value) => sum(index) += value)
+            axpy(1.0, point.vector, sum)
             counts(i)(bestCenter) += 1
           }
         }
@@ -217,12 +213,7 @@ class KMeans private (
         while (j < k) {
           val (sum, count) = totalContribs((i, j))
           if (count != 0) {
-            val size = sum.size
-            var i = 0
-            while(i < sum.size) {
-              sum(i) /= count
-              i += 1
-            }
+            scal(1.0 / count, sum)
             val newCenter = new VectorWithNorm(sum)
             if (KMeans.fastSquaredDistance(newCenter, centers(run)(j)) > epsilon * epsilon) {
               changed = true

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

@@ -21,7 +21,7 @@ import scala.util.Random
 
 import org.apache.spark.Logging
 import org.apache.spark.mllib.linalg.Vectors
-import org.apache.spark.mllib.linalg.BLAS.axpy
+import org.apache.spark.mllib.linalg.BLAS.{axpy, scal}
 
 /**
  * An utility object to run K-means locally. This is private to the ML package because it's used
@@ -75,12 +75,12 @@ private[mllib] object LocalKMeans extends Logging {
     while (moved && iteration < maxIterations) {
       moved = false
       val counts = Array.fill(k)(0.0)
-      val sums = Array.fill(k)(Array.ofDim[Double](dimensions))
+      val sums = Array.fill(k)(Vectors.zeros(dimensions))
       var i = 0
       while (i < points.length) {
         val p = points(i)
         val index = KMeans.findClosest(centers, p)._1
-        axpy(weights(i), p.vector, Vectors.dense(sums(index)))
+        axpy(weights(i), p.vector, sums(index))
         counts(index) += weights(i)
         if (index != oldClosest(i)) {
           moved = true
@@ -95,14 +95,7 @@ private[mllib] object LocalKMeans extends Logging {
           // Assign center to a random point
           centers(j) = points(rand.nextInt(points.length)).toDense
         } else {
-          val sum = sums(j)
-          val count = counts(j)
-          val size = sum.size
-          var i = 0
-          while(i < size) {
-            sum(i) /= count
-            i += 1
-          }
+          scal(1.0 / counts(j), sums(j))
           centers(j) = new VectorWithNorm(sums(j))
         }
         j += 1

mllib/src/main/scala/org/apache/spark/mllib/util/MLUtils.scala

@@ -311,9 +311,11 @@ object MLUtils {
     } else if (v1.isInstanceOf[SparseVector] || v2.isInstanceOf[SparseVector]) {
       val dotValue = dot(v1, v2)
       sqDist = math.max(sumSquaredNorm - 2.0 * dotValue, 0.0)
-      val precisionBound2 = EPSILON * (sumSquaredNorm + 2.0 * math.abs(dotValue)) / (sqDist + EPSILON)
+      val precisionBound2 = EPSILON * (sumSquaredNorm + 2.0 * math.abs(dotValue)) /
+        (sqDist + EPSILON)
       if (precisionBound2 > precision) {
-        // TODO: breezeSquaredDistance is slow, so we should replace it with our own implementation.
+        // TODO: breezeSquaredDistance is slow,
+        // so we should replace it with our own implementation.
         sqDist = breezeSquaredDistance(v1.toBreeze, v2.toBreeze)
       }
     } else {