PySpark含数百独热编码向量列的DataFrame如何保存？

嘿，这个问题我之前处理过！完全不用手动遍历所有列去转字符串，有几个更高效的方案，看你的实际需求来选：

这是最省心的方案，因为Parquet和ORC是Spark原生支持的列式存储格式，完美兼容向量类型，而且有以下优势：

• 不需要做任何数据转换，直接保存就能完整保留所有列的类型信息
• 读写速度比CSV快很多，还支持压缩，能大幅节省存储空间
• 后续用Spark加载时，直接就能用向量列做后续计算，不需要反向转换

代码示例超简单：

# 保存为Parquet
df.write.mode("overwrite").parquet("/path/to/your_data.parquet")

# 或者保存为ORC
df.write.mode("overwrite").orc("/path/to/your_data.orc")

如果你的场景必须输出CSV（比如要给不支持列式格式的工具使用），可以用Spark的内置函数批量处理向量列，不用手动循环转字符串：

方案1：将向量转为数组后保存 #

Spark 3.0及以上提供了vector_to_array函数，可以直接把向量列转成数组类型，保存CSV时数组会自动序列化为逗号分隔的字符串（比如[1,0,1]会变成"1,0,1"）。

代码示例：

from pyspark.sql.functions import vector_to_array, col

# 先找出所有向量类型的列
vector_columns = [col_name for col_name, dtype in df.dtypes if dtype == "vector"]

# 批量将向量列转为数组
for col_name in vector_columns:
    df = df.withColumn(col_name, vector_to_array(col(col_name)))

# 保存为CSV
df.write.mode("overwrite").csv("/path/to/your_data.csv", header=True)

方案2：将向量展开为单独的列 #

如果需要把向量的每个元素拆成单独的列（比如方便其他工具直接读取单个特征），可以先转数组再展开：

from pyspark.sql.functions import vector_to_array, col

vector_columns = [col_name for col_name, dtype in df.dtypes if dtype == "vector"]

# 先把向量转成数组
df = df.selectExpr(
    "*",
    *[f"vector_to_array({col}) as {col}_array" for col in vector_columns]
).drop(*vector_columns)

# 遍历每个数组列，展开为单独的特征列
for col_name in vector_columns:
    array_col = f"{col_name}_array"
    # 获取向量的长度（假设所有同类型向量长度一致）
    vec_length = df.select(f"size({array_col})").first()[0]
    # 逐个提取数组元素作为新列
    for idx in range(vec_length):
        df = df.withColumn(f"{col_name}_{idx}", col(array_col)[idx])
    # 删除临时的数组列
    df = df.drop(array_col)

# 现在所有列都是基本类型，直接保存CSV
df.write.mode("overwrite").csv("/path/to/your_data_expanded.csv", header=True)

如果你用的是Spark 3.0以下的版本，没有vector_to_array函数，可以用自定义UDF来实现向量转数组：

from pyspark.sql.functions import udf
from pyspark.sql.types import ArrayType, DoubleType

# 定义UDF：把向量转为Python列表
vec_to_array_udf = udf(lambda vec: vec.toArray().tolist(), ArrayType(DoubleType()))

vector_columns = [col_name for col_name, dtype in df.dtypes if dtype == "vector"]
for col_name in vector_columns:
    df = df.withColumn(col_name, vec_to_array_udf(col(col_name)))

# 之后就可以按上面的方法保存CSV了

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