pyspark 解决数据倾斜demo（两端聚合)（全网之最）

什么是数据倾斜

数据倾斜是一种很常见的问题（依据二八定律），简单来说，比方WordCount中某个Key对应的数据量非常大的话，就会产生数据倾斜，导致两个后果：

• OOM（单或少数的节点）；
• 拖慢整个Job执行时间（其他已经完成的节点都在等这个还在做的节点）。

Shuffle时，需将各节点的相同key的数据拉取到某节点上的一个task来处理，若某个key对应的数据量很大就会发生数据倾斜。比方说大部分key对应10条数据，某key对应10万条，大部分task只会被分配10条数据，很快做完，个别task分配10万条数据，不仅运行时间长，且整个stage的作业时间由最慢的task决定。

什么是两端聚合

• 两阶段聚合（局部聚合+全局聚合）。
• 场景：对RDD进行reduceByKey等聚合类shuffle算子，SparkSQL的groupBy做分组聚合这两种情况。
• 思路：首先通过map给每个key打上n以内的随机数的前缀并进行局部聚合，即(hello, 1) (hello, 1) (hello, 1) (hello, 1)变为(1_hello, 1) (1_hello, 1) (2_hello, 1)，并进行reduceByKey的局部聚合，然后再次map将key的前缀随机数去掉再次进行全局聚合。本文demo案例给出df下的倾斜处理方式。
• 原理：对原本相同的key进行随机数附加，变成不同key，让原本一个task处理的数据分摊到多个task做局部聚合，规避单task数据过量。之后再去随机前缀进行全局聚合。
• 优点：效果非常好（对聚合类Shuffle操作的倾斜问题）。
• 缺点：范围窄（仅适用于聚合类的Shuffle操作，join类的Shuffle还需其它方案）。

案例图说明数据倾斜

本文demo案例是根据df的，不是rdd的，特此说明

如图所示，大量相同的key（上图的A）被分到同一分区，导致数据倾斜。

两端聚合解决数据倾斜案例图

将key加盐，实际上就是对key加随机整数。避免大量相同的key出现在同一分区。

案例code

假数据生成

from pyspark import SparkContext, SQLContext, SparkConf
from pyspark.sql import SparkSession
from pyspark.sql import functions as fn
from pyspark.sql.functions import udf

tmpdict = [
    {'Col1': 'A', 'Col2': 1},
    {'Col1': 'A', 'Col2': 1},
    {'Col1': 'A', 'Col2': 1},
    {'Col1': 'A', 'Col2': 1},
    {'Col1': 'A', 'Col2': 1},
    {'Col1': 'A', 'Col2': 1},
    {'Col1': 'A', 'Col2': 1},
    {'Col1': 'A', 'Col2': 1},
    {'Col1': 'B', 'Col2': 1},
    {'Col1': 'B', 'Col2': 1},
    {'Col1': 'B', 'Col2': 1},
    {'Col1': 'A', 'Col2': 1},
    {'Col1': 'A', 'Col2': 1},
    {'Col1': 'A', 'Col2': 1}
]
df = ss.createDataFrame(tmpdict)
df.show()

+----+----+
|Col1|Col2|
+----+----+
|   A|   1|
|   A|   1|
|   A|   1|
|   A|   1|
|   A|   1|
|   A|   1|
|   A|   1|
|   A|   1|
|   B|   1|
|   B|   1|
|   B|   1|
|   A|   1|
|   A|   1|
|   A|   1|
+----+----+

加盐

@udf
def adSalt(adid):
    salt = random.sample(range(0, 4), 1)[0]
    salt = str(salt) + '_'
    return salt + adid


df = df.withColumn('Col1', adSalt(fn.col('Col1')))
df.show()

+----+----+
|Col1|Col2|
+----+----+
| 1_A|   1|
| 0_A|   1|
| 1_A|   1|
| 1_A|   1|
| 1_A|   1|
| 1_A|   1|
| 2_A|   1|
| 1_A|   1|
| 3_B|   1|
| 0_B|   1|
| 0_B|   1|
| 3_A|   1|
| 1_A|   1|
| 2_A|   1|
+----+----+

局部聚合

def row_dealWith(data):
    Col1, Col2 = data[0], data[1]
    
    tups = (
        str(Col1),
        int(sum(Col2))
    )
    return tups
    
    
df = df.groupBy('Col1').agg(fn.collect_list('Col2').alias('Col2')).rdd.map(row_dealWith).toDF(schema=['Col1', 'Col2'])
df.show()

+----+----+
|Col1|Col2|
+----+----+
| 3_A|   4|
| 1_A|   1|
| 0_B|   3|
| 0_A|   3|
| 2_A|   3|
+----+----+

去掉随机串

getidUDF = fn.udf(lambda x: x.split('_')[1])
df = df.withColumn('Col1', getidUDF(fn.col('Col1')))
df.show()

+----+----+
|Col1|Col2|
+----+----+
|   A|   4|
|   A|   1|
|   B|   3|
|   A|   3|
|   A|   3|
+----+----+

全局聚合

df = df.groupBy('Col1').agg(fn.collect_list('Col2').alias('Col2')).rdd.map(row_dealWith).toDF(schema=['Col1', 'Col2'])
df.show()

+----+----+
|Col1|Col2|
+----+----+
|   B|   3|
|   A|  11|
+----+----+

汇总案例code

from pyspark import SparkContext, SQLContext, SparkConf
from pyspark.sql import SparkSession
from pyspark.sql import functions as fn
from pyspark.sql.functions import udf


tmpdict = [
    {'Col1': 'A', 'Col2': 1},
    {'Col1': 'A', 'Col2': 1},
    {'Col1': 'A', 'Col2': 1},
    {'Col1': 'A', 'Col2': 1},
    {'Col1': 'A', 'Col2': 1},
    {'Col1': 'A', 'Col2': 1},
    {'Col1': 'A', 'Col2': 1},
    {'Col1': 'A', 'Col2': 1},
    {'Col1': 'B', 'Col2': 1},
    {'Col1': 'B', 'Col2': 1},
    {'Col1': 'B', 'Col2': 1},
    {'Col1': 'A', 'Col2': 1},
    {'Col1': 'A', 'Col2': 1},
    {'Col1': 'A', 'Col2': 1}
]
df = ss.createDataFrame(tmpdict)

@udf
def adSalt(adid):
    salt = random.sample(range(0, 4), 1)[0]
    salt = str(salt) + '_'
    return salt + adid

def row_dealWith(data):
    Col1, Col2 = data[0], data[1]
    
    tups = (
        str(Col1),
        int(sum(Col2))
    )
    return tups    

df = df.withColumn('Col1', adSalt(fn.col('Col1')))

df = df.groupBy('Col1').agg(fn.collect_list('Col2').alias('Col2')).rdd.map(row_dealWith).toDF(schema=['Col1', 'Col2'])

getidUDF = fn.udf(lambda x: x.split('_')[1])
df = df.withColumn('Col1', getidUDF(fn.col('Col1')))

df = df.groupBy('Col1').agg(fn.collect_list('Col2').alias('Col2')).rdd.map(row_dealWith).toDF(schema=['Col1', 'Col2'])

两端聚合解决数据倾斜模板思路已给出，具体结合自己需求修改即可

原创不易，转载请注明出处