如何计算pyspark中两个矩阵的外积

Question

假设我有一个spark数据框，如下所示。我有两个带有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？

谢谢！