python - 使用具有特征的原始 RDD 项将 Pyspark Python k-means 模型预测插入 DF 中

Question

我有一个带有 ID 和功能的 Rdd。简而言之，我试图输出与标签(“id”)与其所属的簇号(0、1、2等)相匹配的内容

rdd 数据集中的三行看起来像这样(虽然它更像是 100 行，第一项是字符串，其余的是 float ):

rdd = ["id1",2,12,3.4,19], ["id2",4,17,3.6,40] ["id3",5,14,2.3,47]...

我通过创建一个仅包含特征的 RDD 来运行该模型的特征(id 破坏了直接在原始 RDD 上运行的模型):

feature_rdd = [2,12,3.4,19], [4,17,3.6,40] [5,14,2.3,47]...

model = KMeans.train(parsedData, num_clusters, maxIterations=max_iterations, initializationMode=initialization_mode, seed=seed)

我预测使用:

predictions = model.predict(feature_rdd)

并得到一个看起来像这样的RDD，对应于该行的预测的簇号:

[0, 0, 1, 2, 0...]

我想以某种方式将 id 与预测结合起来，这样我就可以报告哪些 ID 属于哪个集群。我找不到一个很好的方法来做到这一点。我尝试合并两个 RDD，但随后它只给出新 Rdd 中的另一个项目，而不是将每个预测与每个 ID 配对。我还尝试过转换两个数据框，但在变量的混合转换方面遇到了问题。我想做一些类似于数据框的事情:

*****************
* id  * cluster *
*****************
* "id1" *    0  * 
* "id2" *    0  *
* "id3" *    1  *
*****************

或者只是以某种方式配对在一起并可导出到列表等。

["id1", 0],["id2", 1]...

但是，我们非常感谢任何有关如何解决此问题的帮助。

Answer 1

您可以使用map获取具有特征的rdd的第一个条目，然后使用zip添加预测的集群。您可以使用 createDataFrame 转换生成的 rdd。下面显示了一个示例，希望这对您有所帮助。

from pyspark.sql.types import StructType, StructField, StringType, IntegerType

data = sc.parallelize(
[
    ('id1',1,2,3),
    ('id2',2,3,4),
    ('id3',3,4,5)
])
predictions = sc.parallelize(
[
    (1),
    (0),
    (1)
])

# zip the id's (first element of each entry in the rdd) and the predictions into one rdd.
id_and_predictions = data.map(lambda x: x[0]).zip(predictions)

# Convert to DataFrame
schema = StructType([
    StructField('id',StringType()), StructField('cluster',IntegerType())
])
df = sqlContext.createDataFrame(id_and_predictions,schema)
df.show()

输出:

+---+-------+
| id|cluster|
+---+-------+
|id1|      1|
|id2|      0|
|id3|      1|
+---+-------+