如何用Scala在Spark中实现特征与输出变量的相关性排序

实现代码 #

你可以直接在现有代码的基础上追加以下逻辑，即可得到和Pandas版本一致的排序结果：

// 1. 获取你传入VectorAssembler的所有特征名
val featureNames = assembler.getInputCols
val targetCol = "median_house_value"
// 2. 定位目标列在特征列表中的索引
val targetIndex = featureNames.indexOf(targetCol)
// 3. 从相关系数矩阵中提取目标列和所有特征的相关系数数组
val corrValues = coeff_matrix.col(targetIndex).toArray
// 4. 特征名与相关系数配对，过滤掉目标列自身的相关系数（值为1）
val featureCorrPairs = featureNames.zip(corrValues).filter(_._1 != targetCol)
// 5. 按相关系数降序排序
val sortedCorr = featureCorrPairs.sortBy(-_._2)

// 格式化打印输出结果
sortedCorr.foreach { case (featureName, corrCoefficient) =>
  println(f"$featureName%-20s $corrCoefficient%.4f")
}

逻辑说明 #

• Spark ml.linalg.Matrix 默认是列优先存储，调用col(索引)方法即可提取对应位置特征和所有特征的相关系数序列，顺序和你传入VectorAssembler的特征顺序完全对齐
• 动态获取目标列索引的写法比硬编码索引更稳妥，后续调整特征顺序不需要修改索引数值
• 格式化打印可以对齐列，可读性更好，和Pandas输出效果接近

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