zhengruifeng commented on a change in pull request #26739: [SPARK-29967][ML][PYTHON] KMeans support instance weighting URL: https://github.com/apache/spark/pull/26739#discussion_r353573872
########## File path: mllib/src/main/scala/org/apache/spark/mllib/clustering/KMeans.scala ########## @@ -278,30 +287,32 @@ class KMeans private ( val bcCenters = sc.broadcast(centers) // Find the new centers - val collected = data.mapPartitions { points => + val collected = data.mapPartitions { pointsAndWeights => val thisCenters = bcCenters.value val dims = thisCenters.head.vector.size val sums = Array.fill(thisCenters.length)(Vectors.zeros(dims)) - val counts = Array.fill(thisCenters.length)(0L) - points.foreach { point => - val (bestCenter, cost) = distanceMeasureInstance.findClosest(thisCenters, point) + // clusterWeightSum is needed to calculate cluster center + // cluster center = + // sample1 * weight1/clusterWeightSum + sample2 * weight2/clusterWeightSum + ... + val clusterWeightSum = Array.fill(thisCenters.length)(0.0) + + pointsAndWeights.foreach { case (point, weight) => + var (bestCenter, cost) = distanceMeasureInstance.findClosest(thisCenters, point) + cost *= weight costAccum.add(cost) - distanceMeasureInstance.updateClusterSum(point, sums(bestCenter)) - counts(bestCenter) += 1 + distanceMeasureInstance.updateClusterSum(point, sums(bestCenter), weight) + clusterWeightSum(bestCenter) += weight } - counts.indices.filter(counts(_) > 0).map(j => (j, (sums(j), counts(j)))).iterator - }.reduceByKey { case ((sum1, count1), (sum2, count2)) => + clusterWeightSum.indices.filter(clusterWeightSum(_) > 0) + .map(j => (j, (sums(j), clusterWeightSum(j)))).iterator + }.reduceByKey { case ((sum1, clusterWeightSum1), (sum2, clusterWeightSum2)) => axpy(1.0, sum2, sum1) - (sum1, count1 + count2) + (sum1, clusterWeightSum1 + clusterWeightSum2) }.collectAsMap() - if (iteration == 0) { - instr.foreach(_.logNumExamples(collected.values.map(_._2).sum)) - } - Review comment: I am OK to add new `instr.log` in other PR. Here I prefer to keep `instr.logNumExamples` log the unweighted count, in order to keep it in sync with other algs. ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: us...@infra.apache.org With regards, Apache Git Services --------------------------------------------------------------------- To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org For additional commands, e-mail: reviews-h...@spark.apache.org