Spark ML中KMeans模型如何预测新数据点所属集群？

嘿，我明白你的困惑！Spark ML（DataFrame API）和旧的MLlib（RDD API）在使用逻辑上确实有差异，predict(Vector)是MLlib专属的方法，而ML API里我们得用transform，但要注意几个关键细节才能成功跑通：

核心前提：特征列完全匹配 #

Spark ML的KMeans模型是和特征列名称绑定的——训练时你用的特征列（默认是"features"），新数据必须有完全同名的列，而且特征向量的维度要和训练数据保持一致，否则transform要么报错，要么输出错误结果。

正确的预测步骤 #

1. 把新数据转换成符合要求的DataFrame
   不管是单个还是多个数据点，都要包装成包含特征列的DataFrame，比如：

   import org.apache.spark.ml.linalg.Vectors
   // 假设你的训练数据用的是二维特征向量，创建新数据
   val newData = spark.createDataFrame(Seq(
     (Vectors.dense(1.2, 3.4)),
     (Vectors.dense(5.6, 7.8))
   )).toDF("features")
   

   如果你训练时用setFeaturesCol指定了自定义列名（比如"user_features"），那新数据的特征列必须也用这个名字。

2. 调用transform获取集群预测结果
   直接用训练好的模型调用transform，结果会自动新增一个"prediction"列，这个列的值就是数据点所属的集群编号：

   val predictionResults = model.transform(newData)
   // 查看预测结果
   predictionResults.select("features", "prediction").show()
   

想单独预测单个Vector怎么办？ #

Spark ML没有提供直接的model.predict(Vector)方法，但你可以把单个Vector快速转换成DataFrame来实现：

val singlePoint = Vectors.dense(0.5, 0.6)
val singlePointDf = spark.createDataFrame(Seq((singlePoint))).toDF("features")
// 提取单个预测结果
val clusterId = model.transform(singlePointDf).select("prediction").head().getInt(0)

为什么之前transform失败？ #

大概率是以下原因之一：

• 新数据的特征列名称和训练时不一致
• 新数据的特征向量维度和训练数据不匹配（比如训练用3维向量，新数据给了2维）

检查并修正这两点后，应该就能正常完成预测了！

内容的提问来源于stack exchange，提问作者Ishan Kumar