Spark 的两个常用操作

repartition

repartition 将全局数据进行shuffle为指定分区大小。spark 可以根据用户指定的列或者特定的分区数目以及指定的hash函数进行数据分区。

coalesce

coalesce 通过分区数据移动的方式达到降低数据分区的效果，不会增加数据分区。相比repartition 而言，coalesce发生的数据混洗更小，是一个更高效的操作。

示例

假设原始分区如下

node1 = 1, 2, 3
node2 = 4, 5, 6
node3 = 7, 8, 9
node4 = 10, 11, 12

repartition(2) 之后可能会变成如下

node1 = 4, 1, 3, 2, 5, 6
node2 = 8, 9, 7, 12, 10, 11

coalesce(2) 之后可能会变成如下

node1 = 1, 2, 3 + (10, 11, 12)
node3 = 7, 8, 9 + (4, 5, 6)

基于上面的原理，我们封装一个函数用于将DataFrame 转化为指定分区

def repartition_df(df, num_partitions):
    cur_partitions = df.rdd.getNumPartitions()
    if cur_partitions = num_partitions:
        return df
    if cur_partitions > num_partitions:
        df = df.coalesce(num_partitions)
    else:
        df = df.repartition(num_partitions)
    return df