如何用PySpark提取并绘制向量格式的PCA特征？

报错原因说明 #

你遇到的AttributeError: 'DataFrame' object has no attribute 'map'是因为map是PySpark RDD的方法，DataFrame对象并不支持该操作，需要用DataFrame专属的API来处理。

方法一：使用PySpark内置函数提取向量元素 #

假设你的PCA特征列名为pca_features，可以通过getItem方法直接提取向量中指定索引的元素（索引从0开始）：

from pyspark.sql.functions import col

# 依次提取三个特征列
df = df.withColumn("feature1", col("pca_features").getItem(0))
df = df.withColumn("feature2", col("pca_features").getItem(1))
df = df.withColumn("feature3", col("pca_features").getItem(2))

# 验证结果
df.select("Article", "feature1", "feature2", "feature3").show(5)

如果使用Spark 3.0及以上版本，也可以用element_at函数（注意此方法索引从1开始）：

from pyspark.sql.functions import element_at

df = df.withColumn("feature1", element_at(col("pca_features"), 1))
df = df.withColumn("feature2", element_at(col("pca_features"), 2))
df = df.withColumn("feature3", element_at(col("pca_features"), 3))

方法二：自定义UDF提取向量元素 #

如果需要更灵活的向量处理逻辑，可以定义用户自定义函数（UDF）：

from pyspark.sql.functions import udf, lit
from pyspark.sql.types import FloatType

# 定义提取向量指定位置元素的函数
def extract_vector_element(vec, idx):
    return float(vec[idx])

# 注册UDF
extract_udf = udf(extract_vector_element, FloatType())

# 添加特征列
df = df.withColumn("feature1", extract_udf(col("pca_features"), lit(0)))
df = df.withColumn("feature2", extract_udf(col("pca_features"), lit(1)))
df = df.withColumn("feature3", extract_udf(col("pca_features"), lit(2)))

方法三：转Pandas DataFrame后处理（适合小数据量场景） #

如果数据量不大，直接转成Pandas DataFrame后处理会更便捷，还能直接对接可视化库：

import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt

# 将PySpark DataFrame转换为Pandas DataFrame
pd_df = df.select("Article", "pca_features").toPandas()

# 拆分向量列为独立特征列
pd_df[["feature1", "feature2", "feature3"]] = pd_df["pca_features"].apply(
    lambda vec: pd.Series([float(vec[i]) for i in range(3)])
)

# 以Article为色调绘制散点图示例
sns.scatterplot(data=pd_df, x="feature1", y="feature2", hue="Article")
plt.title("PCA Features by Article")
plt.show()

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