Spark MLLib Kmeans 来自数据帧，然后再次返回

我的目标是使用 Spark (1.3.1) MLLib 将 kmeans 聚类算法应用于非常大的数据集。我已经使用 Spark 中的 hiveContext 调用了 HDFS 中的数据，并且最终希望以这种方式将其放回原处 - 以这种格式

    |I.D     |cluster |
    ===================
    |546     |2       |
    |6534    |4       |
    |236     |5       |
    |875     |2       |

我运行了以下代码，其中“data”是双精度数的数据框，以及第一列的 ID。

    val parsedData = data.rdd.map(s => Vectors.dense(s.getDouble(1),s.getDouble(2))).cache()
    val clusters = KMeans.train(parsedData, 3, 20)

它运行成功，我现在被困在如上所述的数据帧中将集群映射回它们各自的 ID。我可以使用以下方法将其转换为数据框：

    sc.makeRDD(clusters.predict(parsedData).toArray()).toDF()

但据我所知。这个帖子走在正确的轨道上，并且这个帖子我认为正在问与我类似的问题。

我怀疑需要labeledPoint 库。任何评论，答案将不胜感激，干杯。

编辑：刚刚发现this在 Spark 用户列表中，看起来很有前途

我知道你想最后得到 DataFrame 。我看到两种可能的解决方案。我想说，在它们之间进行选择是品味问题。

从 RDD 创建列

以 RDD 的形式获取 id 和 cluster 对非常容易：

val idPointRDD = data.rdd.map(s => (s.getInt(0), Vectors.dense(s.getDouble(1),s.getDouble(2)))).cache()
val clusters = KMeans.train(idPointRDD.map(_._2), 3, 20)
val clustersRDD = clusters.predict(idPointRDD.map(_._2))
val idClusterRDD = idPointRDD.map(_._1).zip(clustersRDD)

然后你从中创建 DataFrame

val idCluster = idClusterRDD.toDF("id", "cluster")

它之所以有效，是因为地图不会改变 RDD 中数据的顺序，这就是为什么你可以只用预测结果压缩 id。

使用UDF（用户定义函数）

第二种方法涉及使用clusters.predict方法作为 UDF：

val bcClusters = sc.broadcast(clusters)
def predict(x: Double, y: Double): Int = {
    bcClusters.value.predict(Vectors.dense(x, y))
}
sqlContext.udf.register("predict", predict _)

现在我们可以使用它来向数据添加预测：

val idCluster = data.selectExpr("id", "predict(x, y) as cluster")

请记住，Spark API 不允许取消 UDF 注册。这意味着闭包数据将保存在内存中。

错误/非最佳解决方案

• Using clusters.predict without broadcasting

It won't work in the distributed setup. Edit: actually it will work, I was confused by implementation of predict for RDD, which uses broadcast.

• sc.makeRDD(clusters.predict(parsedData).toArray()).toDF()

toArray收集驱动程序中的所有数据。这意味着在分布式模式下，您将把集群 ID 复制到一个节点中。