如何用PySpark创建n*k零矩阵？解决80K规模矩阵内存报错问题

首先得明确：你之前用numpy创建矩阵再转Spark DataFrame的方式行不通，完全是因为numpy是单机内存计算框架。80K×80K的矩阵光是存储就需要80000×80000×8字节=512GB内存，这远远超出了单机的内存上限，必然会触发OOM错误。Spark的核心优势是分布式计算，所以我们必须用分布式的思路来创建矩阵，而不是依赖单机的numpy。

一、创建n×n的零矩阵（针对你的相似度矩阵场景） #

直接创建宽表形式的n×n矩阵（80K列）非常不推荐——Spark对超宽表的支持很差，会导致性能暴跌甚至报错。更合理的方式是用**长表（三元组）**存储，也就是每行记录(行索引, 列索引, 0)，这种形式天然适合分布式存储。

步骤1：给原DataFrame添加行索引 #

先给你的data1分配唯一的行索引，方便后续生成所有行对：

from pyspark.sql import functions as F
from pyspark.sql.window import Window

# 给data1添加从0开始的行索引
indexed_data = data1.withColumn(
    "row_idx",
    F.row_number().over(Window.orderBy(F.monotonically_increasing_id())) - 1
)
n = indexed_data.count()  # 获取行数n

步骤2：生成所有行对并赋值0 #

通过笛卡尔积生成所有(i,j)组合，再添加值为0的列：

# 提取唯一的行索引，生成单独的DataFrame
row_indices = indexed_data.select("row_idx").distinct()

# 笛卡尔积得到所有行-列索引对
all_pairs = row_indices.crossJoin(row_indices.withColumnRenamed("row_idx", "col_idx"))

# 添加零值列，得到长表形式的零矩阵
zero_matrix_long = all_pairs.withColumn("value", F.lit(0.0))

（可选）转成宽表（不推荐） #

如果你确实需要宽表形式，可通过pivot转换，但要做好性能下降的准备：

# 注意：80K列会极大消耗资源，仅在k很小的时候考虑使用
zero_matrix_wide = zero_matrix_long.groupBy("row_idx").pivot("col_idx").agg(F.first("value"))

二、创建n×k规模的零矩阵 #

分两种场景处理：

场景1：k是指定的固定数值（比如k=100） #

直接生成行索引，再批量添加k个值为0的列：

n = data1.count()
k = 100  # 你需要的列数

# 生成n行的行索引DataFrame
row_df = spark.range(n).withColumnRenamed("id", "row_idx")

# 批量生成k个零列
zero_cols = [F.lit(0.0).alias(f"col_{i}") for i in range(k)]
n_k_matrix = row_df.select("row_idx", *zero_cols)

场景2：k对应另一个DataFrame的行数 #

类似n×n的思路，用两个DataFrame的索引做笛卡尔积：

# 假设data2的行数为k，先给data2添加列索引
indexed_data2 = data2.withColumn(
    "col_idx",
    F.row_number().over(Window.orderBy(F.monotonically_increasing_id())) - 1
)

# 提取两个DataFrame的索引做笛卡尔积，再添加零值
row_indices = indexed_data.select("row_idx").distinct()
col_indices = indexed_data2.select("col_idx").distinct()

n_k_matrix_long = row_indices.crossJoin(col_indices).withColumn("value", F.lit(0.0))

额外建议：相似度计算的更优方案 #

如果你创建零矩阵是为了后续填充相似度值，其实完全不需要提前创建全零矩阵——Spark MLlib提供了分布式的相似度计算工具，直接生成稀疏的相似度矩阵（只存储非零/有效相似度值），效率高得多：

from pyspark.ml.feature import VectorAssembler
from pyspark.mllib.linalg.distributed import RowMatrix

# 1. 把data1的特征列合并成向量列
assembler = VectorAssembler(inputCols=["你的特征列1", "特征列2", ...], outputCol="features")
feature_df = assembler.transform(data1).select("features")

# 2. 转成分布式RowMatrix
row_matrix = RowMatrix(feature_df.rdd.map(lambda x: x.features))

# 3. 计算余弦相似度矩阵（返回稀疏矩阵）
similarity_matrix = row_matrix.columnSimilarities()

# 4. 转成DataFrame方便后续处理
similarity_df = spark.createDataFrame(
    similarity_matrix.entries.map(lambda e: (e.i, e.j, e.value)),
    ["row_idx", "col_idx", "similarity"]
)

这种方式不需要存储全零的无效值，既节省资源又提升计算速度。

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