apache-spark - 如何在pyspark中计算两个矩阵的外积

Question

假设我有一个如下的火花数据框。我有两个 20M 行的向量。一个列大小为 15000，另一个列大小为 200。

>>> df.printSchema()
root
 |-- id: string (nullable = true)
 |-- vec1: vector (nullable = true)
 |-- vec2: vector (nullable = true)

>>> df.count()
20000000
>>> df.rdd.first()[1].size
15000
>>> df.rdd.first()[2].size
200

将两组向量表示为矩阵 A 和 B，我想计算 A'B，这将是一个 15000 X 200 的小矩阵。

我试图通过蛮力

df.rdd \
.map(lambda row: np.outer(row["vec1"].toArray(), row["vec2"].toArray())) \
.reduce(lambda a,b: a+b)

这适用于较小的列大小向量，但不适用于我正在处理的维度。

我还尝试使用 mapPartitions 方法，首先计算每个分区内的乘积，这似乎效果不佳。

在这种情况下，有没有更有效的方法来计算 A'B？

谢谢！

标签： apache-sparkpyspark