从一段代码浅谈pyspark性能优化

问题引出

我们在日常的特征工程中，常常需要将多张表进行关联操作，也就是所谓的join。现在有三张表A,B,C，其中A表数据总大小约300M, B表总数据大小约15G，C表数据总大小约400G，现在的需求是对这三张表做join，该如何实现？

常规做法

最简单的一种实现，就是先将其中的两张表join，再将剩下的一张表做join，代码如下：

sc = SparkSession\
        .builder\
        .appName("Test")\
        .getOrCreate()

A = sc.sparkContext.textFile("...")
B = sc.sparkContext.textFile("...")
C = sc.sparkContext.textFile("...")
A_rdd = A.map(read_A)
B_rdd = B.map(read_B).reduceByKey(add, 40)
C_rdd = C.map(read_C)
BC_rdd = B_rdd.join(C_rdd).map(merge_BC)
result = A_rdd.join(BC_rdd, 320).mapValues(cal_final_score)
result.saveAsTextFile("...")

上面的代码机器简单，假如数据量小还勉强OK，但在目前的实际场景下是跑不起来的。笔者实测在8台M10(45核，125G内存)的机器所搭建的集群上，运行了3个小时，最终还是挂了。查看日志抛出的异常先是org.apache.spark.shuffle.FetchFailedException和Lost Task等，spark不断的重试，到后面又出现了OOM的问题。该如何处理这块的问题呢？第一反应肯定是加内存加CPU，但是这属于指标不治本的方法，而且不一定能够解决问题。我们首先需要从崩溃的心情中冷静下来，看看代码中是否存在一些可优化的地方，我们可以从这几点考虑：

• 三张表的大小是不一样的，而且差别巨大，该代码在一开始就试图将最大的两张表进行join，这明显是没必要的。
• 是不是三张表的join，一定就需要两次spark的join操作呢？
• 该代码所操作的数据量巨大，而却没有任何类似checkpoint的操作，会导致失败的时候重复计算做无用功。
• 如果是单纯的大表和大表join是否可以有些办法呢？

  优化方法

  常见的优化策略总结如下：
  
  
  所以我们可以采用以下策略，对该代码进行优化：
• 对于大表B和小表A的join。可以将最小的A broadcast到各个节点，之后与B在map端做join，避免shuffle。
• 在两张表join之前，先利用broadcast到各个节点的A，对数据进行filter，减少未来shuffle的规模。
• 在对join操作的两张表采用相同的分区策略，使得同一个key留在相同的节点，减少shuffle和网络IO。
• 加入checkpoint机制，避免重复计算。值得注意的是在checkpoint前，最好对rdd做cache，否则会重复计算两次。
• 优化shuffle相关参数，具体如下：
  • —conf spark.default.parallelism=2000 #提高shuffle阶段的任务并行度，降低单个任务的内存占用
    —conf spark.shuffle.file.buffer=128k #提高shuffle 缓冲区大小
    —conf spark.yarn.executor.memoryOverhead=1g #提高shuffle 缓冲区大小
    —conf spark.shuffle.consolidateFiles=true #合并shuffle write的输出文件，减少磁盘IO

修改后的代码如下：

    sc = SparkSession\
            .builder\
            .appName("...")\
            .getOrCreate()

    sc.sparkContext.setCheckpointDir('...')

    A = sc.sparkContext.textFile(INPUT_PATH, PARTITIONS)
    B = sc.sparkContext.textFile(B_PATH, PARTITIONS)
    C = sc.sparkContext.textFile(C_PATH, PARTITIONS)

    A_rdd = A.map(read_A).filter(lambda x: x != None)
    A_mp = sc.sparkContext.broadcast(A_rdd.collectAsMap())

    B_rdd = B.map(read_B).filter(lambda x: x != None and x[0] in A_mp.value)\
    .reduceByKey(add, PARTITIONS).map(lambda line: join_A(line, A_mp)).partitionBy(PARTITIONS, lambda k: k).cache()
    #printRddInfo(B_rdd)
    B_rdd.checkpoint()

    C_rdd = C.map(read_C).filter(lambda x: x != None and x[0] in A_mp.value).partitionBy(PARTITIONS, lambda k: k)
    # printRddInfo(C_rdd)

    result = B_rdd.join(C_rdd, PARTITIONS).mapValues(cal_final_score).filter(lambda x: x != None)

    result.saveAsTextFile("...")

    B.unpersist()
    C.unpersist()
    B_rdd.unpersist()
    sc.stop()

经过以上修改后，该代码约在7分钟左右便可以运行完成。

大表是否也可以如上做map-join?

以上问题算是解决了，但是心里又有个疑问，既然broadcast后进行map join可以避免shuffle，那是否对于大表（大于2G）也可以采用类似的方法呢？答案是否定的，尤其是对于pyspark来讲更是完全不可取的。除了网络IO之外，由于pyspark是每个partition都会起一个进程去单独对应，并且每个进程会对broadcast的数据保留一份序列化的副本，有n个进程就会保存n份，会导致内存资源的严重不足。
内存占用情况如下：

spark源码中work.py对应代码如下：

结语

spark的优化属于坑非常深的领域，并且会经常踩到，欢迎大家有什么好的方法来找我一起讨论~~