在搭建好了Spark开发环境之后，我们就可以编程了。这里介绍Spark自带的MLLib中的KMeans源码，尽管最新的Spark已经逐渐开始使用DataFrame作为数据模型，但是旧的RDD依然适用。下面我们将逐渐讲解KMeans的源码。KMeans主要步骤如下：
1、读取数据 2、对数据进行正规化 3、选择几个初始类中心
首先，我们给出输入的文件kmeans_data.txt，文件包含6行数据，每一行是三个以空格分割的数字：
0.0 0.0 0.0
0.1 0.1 0.1
0.2 0.2 0.2
9.0 9.0 9.0
9.1 9.1 9.1
9.2 9.2 9.2
然后，我们给出Scala版本调用KMeans的方法，这里我们简单改写了一下官网的解析数据的方式，官网提供的是map，我们改成了mapPartitions（原因参考： ）具体如下：
欢迎关注 DataLearner 官方微信，获得最新 AI 技术推送
package rdd.ml.clustering
import org.apache.spark.mllib.clustering.{KMeans, KMeansModel}
import org.apache.spark.mllib.linalg.Vectors
import org.apache.spark.{SparkConf, SparkContext}
/**
* Created by d00454735 on 2018/7/25.
*/
object KMeansTest {
def main(args: Array[String]): Unit = {
// 以本地模式运行，并初始化SparkContext
val masterURL = "local[*]"
val conf = new SparkConf().setAppName("KMeans Test").setMaster(masterURL)
val sc = new SparkContext(conf)
// 载入数据，并将每一行的数据解析，注意，官网给出的解析使用的是map方法，这里改写成了mapPartitions，因为后者是以数据块为单位的方式处理，其效率要远高于map方式（以数据行位单位处理）
val data = sc.textFile("file:/d:/data/kmeans_data.txt")
val parsedData = data.mapPartitions(partition => parseData(partition)).cache()
// 设定KMeans聚类的参数并对模型进行训练，这里定义聚成2类，迭代20次
val numClusters = 2
val numIterations = 20
val clusters = KMeans.train(parsedData, numClusters, numIterations)
// 计算Within Set Sum of Squared Errors以评估聚类效果
val WSSSE = clusters.computeCost(parsedData)
println(s"Within Set Sum of Squared Errors = $WSSSE")
// 保存模型以及后续加载模型的写法
clusters.save(sc, "target/org/apache/spark/KMeansExample/KMeansModel")
val sameModel = KMeansModel.load(sc, "target/org/apache/spark/KMeansExample/KMeansModel")
}
//解析数据，每一行先以空额分割成数组，然后将数组内的元素转化成Double形式，最后变成向量
def parseData(lines : Iterator[String]): Iterator[Vector] = {
lines.map(s => Vectors.dense(s.split(' ').map(_.toDouble)))
}
}
下面我们进入KMeans的源代码继续分析，首先来到运行的方法，这个方法先用private[spark]声明成属于spark下的私有方法，参数有两个，一个是RDD类型的数据，另一个是可选参数Option[Instrumentation[NewKMeans]]，Option在scala里面表示这个参数可能存在可能不存在。Option[A] 是一个类型为 A 的可选值的容器： 如果值存在， Option[A] 就是一个 Some[A] ，如果不存在， Option[A] 就是对象 None （参考：https://windor.gitbooks.io/beginners-guide-to-scala/content/chp5-the-option-type.html ）。Instrumentation是用于监控程序运行的类，本质上是java里面的代码。这里面涉及到一个缓存释放的问题，即upersist()，参考：http://bourneli.github.io/scala/spark/2016/06/17/spark-unpersist-after-action.html
在初始化类中心的时候判断是否给定了类中心（参考：https://www.datalearner.com/blog/1051532740001101 ），如果给定了类中心，就用原来的，否则就要随机初始化k个类中心，这里初始化也给了两个方法，一个是所有的随机初始化initRandom(data: RDD[VectorWithNorm])，另一个是并行初始化 initKMeansParallel(data: RDD[VectorWithNorm])，它们都是返回一个VectorWithNorm的数组，
如下：
private[spark] def run(
data: RDD[Vector],
instr: Option[Instrumentation[NewKMeans]]): KMeansModel = {
// 如果数据没有做缓存的话会影响性能，这里做了一个警告检查
if (data.getStorageLevel == StorageLevel.NONE) {
logWarning("The input data is not directly cached, which may hurt performance if its"
+ " parent RDDs are also uncached.")
}
// 计算平方和并缓存，计算方法就是把每一行的元素平方后加和，然后再开方
val norms = data.map(Vectors.norm(_, 2.0))
norms.persist()
//zip是将元素连接起来，这里的含义就是将数据中的每一行和刚才的norm连接
val zippedData = data.zip(norms).map { case (v, norm) =>
new VectorWithNorm(v, norm)
}
//这里开始真正执行程序了
val model = runAlgorithm(zippedData, instr)
//释放缓存数据
norms.unpersist()
// Warn at the end of the run as well, for increased visibility.
if (data.getStorageLevel == StorageLevel.NONE) {
logWarning("The input data was not directly cached, which may hurt performance if its"
+ " parent RDDs are also uncached.")
}
//返回模型
model
}
上面就是KMeans算法的执行流程，接下来，我们看核心算法的方法，即runAgorithm()，其参数也是两个，第一个是刚才带来norm的数据，第二个还是个监控参数。
KMeans更新的逻辑是首先对于每个数据点更新其所在的类，即对于每个数据点需要计算这个数据点与类中心的距离，然后选择最近的类中心点作为它所属的类，这样更新了所有的数据点所属的类之后，再计算每个类下数据点的中心点，那么该中心点就是新的类中心了。
第$i$个类别下的中心点的计算公式为：
m_{i} = \frac{1}{n_i} \sum_{j=1}^{j\in i} X_i
这里的$n_i$表示第$i$个类下有$n_i$个数据，$X_i$表示所有属于类$i$的数据点对应的向量。因此，KMeans的核心更新代码如下：
//采用mapPartitions的方式循环所有数据点
val newCenters = data.mapPartitions { points =>
//thisCenters表示当前的类中心，dims表示数据的维度
val thisCenters = bcCenters.value
val dims = thisCenters.head.vector.size
//这个变量存储类下所有数据点的向量之和，例如更新数据点所属类别后，类1下有两个数据点(1.0, 2.0, 3.0)和(1.0, 1.0, 1.0)，那么sum(1) = (2.0, 3.0, 4.0)，计算某个类下的中心点用到，所以它的维度是类别数，每个类别下的维度与数据维度一致
val sums = Array.fill(thisCenters.length)(Vectors.zeros(dims))
//这个变量是存放每个类下有多少个数据点的
val counts = Array.fill(thisCenters.length)(0L)
//points是当前分区下所有的数据点，我们要计算这些数据点与哪个中心点最近，然后更新一下它所属的类，使用的是KMeans下的findClosest方法，分析见下面。
points.foreach { point =>
//寻找当前数据点所属的中心点，cost是距离，bestCenter是表示属于第几个中心点的索引。
val (bestCenter, cost) = KMeans.findClosest(thisCenters, point)
costAccum.add(cost) //成本函数累加
val sum = sums(bestCenter) //取出对应中心点的向量和
axpy(1.0, point.vector, sum) //把当前数据加上去，sum += 1 * point.vector
counts(bestCenter) += 1 //该类下的数据点数量加1
}
//由于mapPartitions要求返回的对象是Iterator[]，因此，这里返回的counts最后加了iterator
//这里把couns重新包装了一下，返回的是每个类下的向量总和以及数据点的数量
counts.indices.filter(counts(_) > 0).map(j => (j, (sums(j), counts(j)))).iterator
}.reduceByKey { case ((sum1, count1), (sum2, count2)) => //将相同类下的数据取出来，求和
axpy(1.0, sum2, sum1)
(sum1, count1 + count2)
}.mapValues { case (sum, count) => //对于每个类，这里最终求数据向量总和/数据点总数的结果，也就是新的数据中心点了
scal(1.0 / count, sum)
new VectorWithNorm(sum)
}.collectAsMap()
该方法的全部内容实现如下：
private def runAlgorithm(
data: RDD[VectorWithNorm],
instr: Option[Instrumentation[NewKMeans]]): KMeansModel = {
val sc = data.sparkContext
// 这里获取系统时间来监控程序运行时间，注意不要使用currentTimeMillis ，因为它和系统时间有关，在某些时候容易发生错误，参考 https://github.com/databricks/scala-style-guide/blob/master/README-ZH.md#misc_currentTimeMillis_vs_nanoTime
val initStartTime = System.nanoTime()
//初始化中心点，如果模型是载入之前的模型或者是已经初始化过中心点，那么直接采用，否则的话抽样中心点
//抽样中心点有两种方式，一种是随机初始化，一种是采用分布式K-Means++的方式
//具体可参考论文：https://www.datalearner.com/blog/1051532743231974
val centers = initialModel match {
case Some(kMeansCenters) =>
kMeansCenters.clusterCenters.map(new VectorWithNorm(_))
case None =>
if (initializationMode == KMeans.RANDOM) {
initRandom(data)
} else {
initKMeansParallel(data)
}
}
//统计初始化模型所用的时间
val initTimeInSeconds = (System.nanoTime() - initStartTime) / 1e9
logInfo(f"Initialization with $initializationMode took $initTimeInSeconds%.3f seconds.")
var converged = false
var cost = 0.0
var iteration = 0
val iterationStartTime = System.nanoTime()
instr.foreach(_.logNumFeatures(centers.head.vector.size))
// Execute iterations of Lloyd's algorithm until converged
// 执行迭代
while (iteration < maxIterations && !converged) {
val costAccum = sc.doubleAccumulator
val bcCenters = sc.broadcast(centers) //将中心点广播出去，分到各个计算节点上
// Find the new centers
// 在各自的数据分区下，计算分区内数据点与中心点的距离，并更新数据点所属的中心点，采用mapPartitions循环，速度要快于map
val newCenters = data.mapPartitions { points =>
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) = KMeans.findClosest(thisCenters, point)
costAccum.add(cost) //成本函数累加
val sum = sums(bestCenter) //取出对应中心点的向量和
axpy(1.0, point.vector, sum) //sum += 1 * point.vector
counts(bestCenter) += 1
}
counts.indices.filter(counts(_) > 0).map(j => (j, (sums(j), counts(j)))).iterator
}.reduceByKey { case ((sum1, count1), (sum2, count2)) =>
axpy(1.0, sum2, sum1)
(sum1, count1 + count2)
}.mapValues { case (sum, count) =>
scal(1.0 / count, sum)
new VectorWithNorm(sum)
}.collectAsMap()
//将旧的中心点毁掉
bcCenters.destroy(blocking = false)
// 更新类的中心点和距离值，如果距离小于指定值那么收敛了，停止迭代，否则还要继续迭代
converged = true
newCenters.foreach { case (j, newCenter) =>
if (converged && KMeans.fastSquaredDistance(newCenter, centers(j)) > epsilon * epsilon) {
converged = false
}
centers(j) = newCenter
}
cost = costAccum.value
iteration += 1
}
//统计迭代的时间
val iterationTimeInSeconds = (System.nanoTime() - iterationStartTime) / 1e9
logInfo(f"Iterations took $iterationTimeInSeconds%.3f seconds.")
//统计是否达到最大迭代次数
if (iteration == maxIterations) {
logInfo(s"KMeans reached the max number of iterations: $maxIterations.")
} else {
logInfo(s"KMeans converged in $iteration iterations.")
}
//输出cost值
logInfo(s"The cost is $cost.")
//返回模型
new KMeansModel(centers.map(_.vector))
}
在上述循环中，一个核心的方法是计算当前数据点与各个中心点的距离，并返回所属的距离最小的中心点的索引，这个方法的参数有两个，一个是所有的中心点，一个是当前数据点，前者是VectorWithNorm集合，后者是单独的一个VectorWithNorm，返回两个值一个是中心点索引，一个距离。由于计算两个向量的距离需要循环向量中的元素，因此这里用了一个技巧，就是先用一步计算距离的最小值，然后与最优值比较，如果最小值比最优值小，才要计算距离，否则就放弃该点，直接循环，原理参考：https://www.datalearner.com/blog/1051532749735689 ：
private[mllib] def findClosest(
centers: TraversableOnce[VectorWithNorm],
point: VectorWithNorm): (Int, Double) = {
var bestDistance = Double.PositiveInfinity //最优距离，冒泡排序用，初始化为最大值
var bestIndex = 0 //最优中心点索引
var i = 0 //中心点循环索引
centers.foreach { center =>
// Since `\|a - b\| \geq |\|a\| - \|b\||`, we can use this lower bound to avoid unnecessary
// distance computation.
//这里用了一个技巧，就是用lowerBoundOfSqDist表示数据点与当前中心点的距离的最小值，如果该值小于目前的最优值bestDistance才要计算距离，否则就继续循环，计算距离的成本要更高。
var lowerBoundOfSqDist = center.norm - point.norm
lowerBoundOfSqDist = lowerBoundOfSqDist * lowerBoundOfSqDist
if (lowerBoundOfSqDist < bestDistance) {
val distance: Double = fastSquaredDistance(center, point) //计算欧氏距离
if (distance < bestDistance) { //如果距离小于最优值，那么更新最优值
bestDistance = distance
bestIndex = i
}
}
i += 1
}
(bestIndex, bestDistance)
}