如何在PySpark中将SparseVector特征列转换为CSC格式的SparseMatrix

核心思路 #

采用PySpark原生分布式坐标矩阵CoordinateMatrix作为中间载体，全流程基于RDD分布式计算，无全量数据拉取到Driver端的操作，适配超大规模矩阵场景。

实现代码 #

from pyspark.mllib.linalg.distributed import CoordinateMatrix, MatrixEntry

# 1. 为每行数据分配唯一行索引，对应矩阵的行号
df_with_rowid = df.selectExpr("features", "monotonically_increasing_id() as row_id")

# 2. 分布式拆解每个SparseVector为(行号,列号,值)格式的矩阵条目
entries_rdd = df_with_rowid.rdd.flatMap(
    lambda row: [
        MatrixEntry(row.row_id, col_idx, val) 
        for col_idx, val in zip(row.features.indices, row.features.values)
    ]
)

# 3. 构建分布式稀疏坐标矩阵，可直接用于后续分布式矩阵运算
coord_matrix = CoordinateMatrix(entries_rdd)

可选转换：转CSC格式SparseMatrix #

注意：PySpark原生SparseMatrix为单节点内存对象，仅适用于矩阵可完整放入Driver端内存的小体量场景。超大规模矩阵直接使用上述CoordinateMatrix进行分布式运算即可，支持转置、矩阵乘法、行列统计等全量操作。

# 仅小矩阵场景执行，输出为官方定义的CSC格式SparseMatrix
sparse_csc_matrix = coord_matrix.toSparseMatrix()

方案优势 #

• 无全量数据collect到Driver的操作，完全分布式计算，适配超大规模矩阵
• 基于PySpark原生API实现，无需引入第三方依赖
• 转换过程为窄依赖操作，无额外shuffle开销，执行效率高

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