如何查看PySpark的ChiSqSelector选中特征对应的列名与索引？

ChiSqSelector 的 selectedFeatures 属性返回的是选中特征在VectorAssembler输入特征列表的索引值，索引从0开始计数，和你定义inputCols时的顺序完全对应，按以下步骤即可直接映射得到对应的原始列名：

• 首先获取你传入VectorAssembler的原始特征列列表，调用getInputCols()方法可以直接拿到你定义的输入列，不需要手动复写避免出错
• 用训练得到的选择器模型的selectedFeatures拿到选中的索引列表
• 遍历索引列表从原始特征列列表中取出对应列名即可

适配你现有代码的实现示例 #

from pyspark.ml.feature import VectorAssembler, ChiSqSelector

# 你的原有逻辑
df = df.select("IsBeta","AVProductStatesIdentifier","IsProtected","Firewall","HasDetections")
vec_assembler = VectorAssembler(inputCols = ["IsBeta","AVProductStatesIdentifier","IsProtected","Firewall"], outputCol="features")
vec_df = vec_assembler.transform(df)
selector = ChiSqSelector(featuresCol='features', fpr=0.05, outputCol="selectedFeatures",labelCol= "HasDetections")
model = selector.fit(vec_df)
result = model.transform(vec_df)

# 新增特征映射逻辑
# 获取原始特征列列表
origin_feature_cols = vec_assembler.getInputCols()
# 获取选中的特征索引
selected_idx_list = model.selectedFeatures
# 映射得到选中的原始列名
selected_cols = [origin_feature_cols[idx] for idx in selected_idx_list]

# 输出结果
print("选中的特征列名：", selected_cols)
# 如果需要查看索引和列的对应关系可以加这行
print("索引与列名对应关系：", dict(zip(selected_idx_list, selected_cols)))

这个方法不管你有多少特征列都适用，只要特征列是同一个VectorAssembler生成的，映射关系就完全准确。

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