记一次spark数据倾斜问题解决过程

1. 背景

最近运行spark任务时，经常出现任务失败，查看原因都是shuffle过程中某些文件不存在，无法读取。但是这些任务长期运行，会产生通常那种疑问：“以前没问题，怎么最近就有问题了，难道不是任务的问题，是集群又有什么问题了”。
由于没有开启history server，所以重新运行了一次查看原因，发现以下现象：

• 某个Stage的Task大部分很快完成，只有一个Task一直不能完成，而且shuffle过程中数据量过大（该Task数据量超过千万行，大小在10+GB）。
• Task对应的Executor日志报出内存溢出。

从现象很容易看出是发生数据倾斜了，以前看过一篇关于数据倾斜优化的文章（Spark性能优化指南——高级篇），正好实践一下。

2. 问题描述

该问题发生在Spark SQL两表join的过程中，以“大表” join “小表”表示，其中连表条件以“join_key”表示。我们知道Spark SQL在join时会发生shuffle，如果某个key对应的数据量过大，就会发生数据倾斜，即上述问题。现象即整个Spark任务运行超过2小时，最终失败。
从具体SQL来讲，可简化成：

select 大表.join_key, xxx
from 大表 
join 小表
on 大表.join_key = 小表.new_join_keyoin_key;

2.1 问题定位

定位上述问题，可以对数据进行采样，查看是否如我们猜测，真的是大表的某个key对应的数据量过大。可通过以下代码进行定位。

from pyspark.sql.functions import desc

df = spark.sql('select join_key from 大表 where 筛选条件')
df.sample(False, 0.2).groupby('join key').count().sort(desc("count")).show()

这里使用sample对数据进行抽样（首个参数为False，意思是不放回抽样；第二个参数0.2代表取出近似20%的数据）；再使用count计数；按照count倒排查看排在前几名的key对应的是数据量是否远超出其他key。
就本文提到的问题来讲，按照此定位方法，发生排在首位的key数量级远超其他key，所以基本定位是这个原因导致的。

3. 问题解决

3.1节仅是尝试，最终选择3.2节的方案。

3.1 提高shuffle操作的并行度

最简单的尝试无非是提高shuffle操作的并行度，即提高spark.sql.shuffle.partitions（默认200）。
尝试下来发现以下现象：

• 因为增加了spark.sql.shuffle.partitions，所以整体Task的数量增加。
• 对问题没有帮助，问题依然存在。

分析下来就是，对于上述key仍然被分在同一个Task中。

3.2 join时打散key，select时恢复key

这个方法也是参考了（Spark性能优化指南——高级篇中的方案六），有兴趣可以参考原文，本节给出我的解法。

# 步骤一：产生新大表，对大表的join_key做随机后缀，此处添加了后缀>>>10 ~ 20随机数
select concat(join_key, '>>>', floor(rand() * 10 + 10)) as new_join_key from 大表

# 步骤二：对小表的join_key扩展n倍，产生新小表
SELECT
    concat(join_key, '>>>', num) as new_join_key
FROM
(
    SELECT
        join_key, '1' AS inner_join_key
    FROM
        小表
) AS tmp_base
JOIN
(
    SELECT explode(array(10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21)) AS num, '1' AS join_key
) AS tmp_explode
ON
    tmp_base.join_key = tmp_explode.join_key
    
# 步骤三：新大表和新小表join完，恢复join_key
select regexp_replace(新大表.new_join_key, ">>>[0-9]+", "") AS join_key
from 新大表 
join 新小表
on 新大表.new_join_key = 新小表.new_join_key

• 步骤一：对原join_key添加一个随机后缀，这样原来数量级较大的key会被打散到多个key，举例来说，比如原来值为“软件一班”的join_key数据量为10000000（一千万），使用上述方式（10 ~ 20随机数）打散后，变成了10个new_join_key，每个new_join_key被分到大概1000000（一百万）。
• 步骤二：扩展小表，使用explode + array，生成固定的10 ~ 20后缀的表，将其和小表join，结果即是小表每个join_key都被扩大了10倍，扩大的数据用于和新的大表join。
  • 为什么要这样：大表每个join_key被附上随机后缀，为了和这些随机的new_join_key去join上，小表只能遍历生成所有可能的new_join_key，否则就会存在和新大表join不上的情况（但是原来可以join的上），
• 步骤三：新大表 join 新小表，恢复join_key。

4. 效果

通过3.2节的方式，new_join_key将原有的1个join_key打散到了10个，在join时原来被分到1个Task的数据量也会被分到10个。
实际效果也是如此，原来问题中超过千万行，大小在10+GB的数据，被分到了10个Task中，每个Task分到了百万行，1GB左右的数据量；时间也从原来的2小时多，优化到了10分钟，不会再失败。

5. 后续优化

3.2节的方法是对所有join_key进行打散，进一步优化可以指针对发生数据倾斜的key做该优化，其他key不需要做。考虑到这样的做法会对代码有一定的定制化，同时当前的做法已经优化到了10分钟以内，所以没有做进一步的优化。有兴趣的读者可以在需要时尝试。