address review comments

apache · Oct 25, 2017 · 752b685 · 752b685
1 parent 3b0f71c
commit 752b685
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/mllib/src/main/scala/org/apache/spark/ml/clustering/PowerIterationClustering.scala b/mllib/src/main/scala/org/apache/spark/ml/clustering/PowerIterationClustering.scala
@@ -180,8 +180,8 @@ class PowerIterationClustering private[clustering] (
         case Row(id: Long, nbr: Vector, weight: Vector) =>
         require(nbr.size == weight.size,
           "The length of neighbor list must be equal to the the length of the weight list.")
-        val ids = Array.fill(nbr.size)(id)
-        for (i <- 0 until ids.size) yield (ids(i), nbr(i).toLong, weight(i))}
+        nbr.toArray.toIterator.zip(weight.toArray.toIterator)
+          .map(x => (id, x._1.toLong, x._2.toLong))}
     val algorithm = new MLlibPowerIterationClustering()
       .setK($(k))
       .setInitializationMode($(initMode))