Spark上的稀疏矩阵乘法

时间:2019-04-07 20:11:55

标签: apache-spark pyspark distributed-computing

问题

我想在Spark基础架构下以可扩展方式有效地乘以2个稀疏矩阵,并假设两个矩阵都可以放入内存中。

可能接近

开始时,为了获得可比较的基线,我在一个带有约100,000个稀疏向量的数据帧中,在一台机器上进行了微不足道的内部乘法运算(通过转换为scipy csr matrix,分成多个批次并依次多个批次)。 我得到的总运行时间为5分钟。

接下来,我使用here所述的CoordinateMatrix在spark上实现了一个完全基于映射约简的算法,并获得了非常糟糕的性能-整个乘法超过50分钟(!):

def coordinate_matrix_mul(mat_left, mat_right):
    mat_left_cols = mat_left.entries.map(lambda entry: (entry.j, (entry.i, entry.value)))
    mat_right_rows = mat_right.entries.map(lambda entry: (entry.i, (entry.j, entry.value)))

    product_entries = mat_left_cols.join(mat_right_rows)\
                .map(lambda pair: ((pair[1][0][0], pair[1][1][0]), pair[1][0][1]*pair[1][1][1]))\
                .reduceByKey(lambda x,y: x+y)\
                .map(lambda cell: MatrixEntry(cell[0][0], cell[0][1], cell[1])) 

    return CoordinateMatrix(product_entries)

我的最终方法是进行第一次天真尝试(5分钟的运行时间未使用火花),并通过以下步骤并行化批处理乘法:

  • 向工人广播左右矩阵。
  • 以左矩阵[[from_row,to_row),...]的形式将批次指标与rdd平行。
  • 映射批次索引并同时执行乘法。

但是,在上述所有技术之间,我的运行时间最差-超过60分钟。

我尝试更改批处理大小,并检查是否存在内存不足导致交换和脱粒的情况,但这不是事实。

有人知道我想念什么吗?

谢谢。

1 个答案:

答案 0 :(得分:0)

from pyspark.sql import SparkSession, types, Window, functions as F

a = spark.createDataFrame([
    [0, 0, 1],
    [0, 1, 2],
    [1, 1, 4],
    [2, 0, 5],
    [2, 1, 6]], ['a_row', 'a_column', 'a_value']).cache()

b = spark.createDataFrame([
    [0, 0, 1],
    [0, 1, 2],
    [1, 0, 3],
    [1, 1, 4]], ['b_row', 'b_column', 'b_value']).cache()

c = a.join(b, a.a_column == b.b_row) \
    .withColumn('product', F.udf(lambda x, y: x * y, types.IntegerType())('a_value', 'b_value')) \
    .groupBy(['a_row', 'b_column']).agg(F.sum('product').alias('c_value')) \
    .select(F.column('a_row').alias('c_row'), F.column('b_column').alias('c_column'), 'c_value').cache()

c.show()

+-----+--------+-------+
|c_row|c_column|c_value|
+-----+--------+-------+
|    1|       0|     15|
|    1|       1|     22|
|    0|       1|     10|
|    2|       0|     23|
|    0|       0|      7|
|    2|       1|     34|
+-----+--------+-------+