从一段代码浅谈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的优化属于坑非常深的领域,并且会经常踩到,欢迎大家有什么好的方法来找我一起讨论~~

展开阅读全文

没有更多推荐了,返回首页