Spark学习记录|RDD分区的那些事

以前在工作中主要写Spark SQL相关的代码，对于RDD的学习有些疏漏。本周工作中学习了一些简单的RDD的知识，主要是关于RDD分区相关的内容。下面的内容都是自己亲身实践所得，如果有错误的地方，还希望大家批评指正。

本文的目录如下：

1、RDD特性-分区列表

2、宽／窄依赖

3、RDD的创建

4、查看分区数据：mapPartitionsWithIndex

5、笛卡尔积：cartesian

6、数据分组：groupByKey

7、重新分区：repartition VS coalesce

8、map Vs mapPartitions

在正式开始文章内容之前，先介绍一下咱们的背景。考虑一下机器学习中网格搜索策略，比如随机森林中，我们想得到n_estimators和max_depth两个参数的最优组合，我们会对给出的参数取值范围计算笛卡尔积，然后对每一种组合训练得到一个效果，并选取效果最好的一组参数。假设我们想使用spark把这个过程并行化，但是参数组合数量太多，没有足够的计算资源，只能一个task上运行几组参数。

举例来说，假设n_estimators有10个取值，max_depth有5个取值，共有5*10=50种组合，最好的方法就是并行50个task。但是由于资源不足，我们只能并行执行10个task，也就是说一个task上执行五组参数组合。

好了，介绍完了背景，是不是很简单？接下来就介绍一下在这一过程中的一些学习收获。

1、RDD特性-分区列表

Spark中的RDD是被分区的，每一个分区都会被一个计算任务（Task处理），分区数决定了并行计算的数量。

2、宽／窄依赖

RDD中的一些算子，会将一个RDD转换为另一个RDD，在这一过程中，由于RDD是分区的，就会产生不同的依赖关系，主要分为宽依赖和窄依赖。

2.1 窄依赖

窄依赖如下图所示：

先定义一下，上图中每一组中左边的称做父RDD、右边的称为子RDD，那么窄依赖就是说：每一个父RDD中的Partition最多被子RDD中的1个Partition所使用。窄依赖最常见的就是map算子。

2.2 宽依赖

宽依赖的示意图如下：

在宽依赖中，一个父RDD的Partition会被多个子RDD所使用。宽依赖也很常见，如我们下文要介绍的groupByKey和repartition。

介绍完一些简单的知识之后，我们开始进入正题了，你应该还没有忘记刚才的背景介绍吧！

3、RDD的创建

首先创建一个sparkSession的对象：

val spark = SparkSession
  .builder()
  .appName("Spark SQL basic example")
  .enableHiveSupport()
  .getOrCreate()

spark.sparkContext.setLogLevel("WARN")

接下来，咱们建立两个RDD，分别是n_estimators和max_depth对应的取值

val n_estimators = Array(10,20,30,40,50,60,70,80,90,100)
val max_depth = Array(3,4,5,6,7)

val n_estimators_rdd = spark.sparkContext.parallelize(n_estimators)
val max_depth_rdd =  spark.sparkContext.parallelize(max_depth)

先来查看一下分区数量：

println(n_estimators_rdd.partitions.length)
println(max_depth_rdd.partitions.length)

输出为：

1
1

可以看到两个RDD的分区数量都是1，那是不是RDD的分区数量默认是1呢？答案当然是否定的，有关于RDD默认的分区数量，可以参考：https://www.jianshu.com/p/fe987f6d2018?utm_source=oschina-app

当然，我们也可以在创建时指定RDD的分区数量：

val n_estimators_rdd = spark.sparkContext.parallelize(n_estimators,10)
val max_depth_rdd =  spark.sparkContext.parallelize(max_depth,5)

此时分区数量就是10和5了。

4、查看分区数据：mapPartitionsWithIndex

接下来你可能会想查看一下每个分区的内容，我们使用mapPartitionsWithIndex函数，先直接给出代码，随后再详细介绍：

n_estimators_rdd.mapPartitionsWithIndex((partid,iter)=>{
      iter.map(n_estimator => {
        println(partid + "," + n_estimator)
        n_estimator
      })
    }).collect()

结果输出如下：

可以看到，每个分区对应了一个n_estimator的数值，关于上述代码，有以下几点需要注意的点：

1）mapPartitionsWithIndex对每一对(分区id,分区内容)进行操作，partid即是分区id，上面的iter即是分区内容。 2）分区内容iter是一个Iterator，在我们这里是Iterator[Int]类型的，对其每个元素进行查看，需要通过map操作。

5、笛卡尔积：cartesian

接下来，我们需要计算两组参数的笛卡尔积，RDD间的笛卡尔积操作示意图如下：

可以看到，经过笛卡尔积后的RDD的Partition数量应该是两个父RDD的分区数量的乘积：

val cartesian_rdd = n_estimators_rdd.cartesian(max_depth_rdd)
println(cartesian_rdd.partitions.length)

由于n_estimators_rdd的分区数量是10，max_depth_rdd的分区数量是5，因此cartesian_rdd的分区数量是50。

6、数据分组：groupByKey

在对n_estimators和max_depth计算笛卡尔积之后，我们已经得到了50组参数组合，并且分布在50个Partition上。无论是通过map还是mapPartitions去并行计算每组参数对应的GBDT拟合效果，都会创建50个task。如果我们只想要10个task，每个task去执行5组参数，那么需要对数据进行分组，使用groupByKey算子。代码如下：

val cartesian_grp_rdd = cartesian_rdd
        .zipWithIndex()
        .map(row=>(row._2 / 5,row._1))
        .groupByKey()

对于上一节中得到的RDD，我们首先使用zipWithIndex()为其添加了编号，此时RDD中每一条数据分为两部分，假设一行用row表示，那么row._2 代表其对应的行号，row._1代表一组实验参数，类型为(Int,Int)。

接下来，使用map把我们的RDD转换成K-V形式的pairRDD，注意这里我们进行row._2 / 5，这样5条数据的Key相同。

最后通过groupByKey()对每5条数据进行了分组。接下来我们来看下经过groupByKey()操作后RDD的分区情况：

cartesian_grp_rdd.mapPartitionsWithIndex((partid,iter)=>{
      iter.map(data => {
        println(partid + "," +data)
        data
      })
    }).collect()

输出如下：

可以看到，分区数量为10，但是，每个分区仅有一条数据，每条数据的Key是我们刚才计算的index，Value是一个包含5组实验参数的CompactBuffer。此时如果我们想对每组参数进行操作的话，还需要将数据转换为List，通过循环进行处理，如下：

cartesian_grp_rdd.mapPartitionsWithIndex((partid,iter)=>{
      iter.map(data => {
        println(partid + "," +data)
        val dataList = data._2.toList
        for(i <- 0 until dataList.length){
          println("do something")
        }
        data
      })
    }).collect()

7、重新分区：repartition VS coalesce

看到这里，你可能会想，能不能不通过groupByKey()来对数据进行重新分区？这里我们介绍两种方法repartition和coalesce。

7.1 coalesce

先上代码：

val cartesian_coalesce_rdd = cartesian_rdd.coalesce(10)

    cartesian_coalesce_rdd.mapPartitionsWithIndex((partid,iter)=>{
      iter.map(data => {
        println(partid + "," +data)
        data
      })
    }).collect()

结果如下：

分区效果还不错，对吧。如果将一个分区较多的RDD重新分区为分区较少的RDD，默认的coalesce是不会进行shuffle过程的（参数中的shuffle默认值为false），其过程类似于如下，是一个分区之间相互组合的过程（窄依赖）：

但是如果想要分区较少的RDD转换为分区较多的RDD，shuffle过程是会有的。

7.2 repartition

而对于repartition来说，它其实就是shuffle=True的coalesce过程，看源码也可以看出来：

def repartition(numPartitions: Int)(implicit ord: Ordering[T] = null): RDD[T] = withScope {
    coalesce(numPartitions, shuffle = true)
  }

先看一下repartition的分区效果：

val repartition_coalesce_rdd = cartesian_rdd.repartition(10)

repartition_coalesce_rdd.mapPartitionsWithIndex((partid,iter)=>{
  iter.map(data => {
    println(partid + "," +data)
    data
  })
}).collect()

感觉分区效果很乱啊，其背后的分区原理是什么呢？还是得看源码：

从源码中可以看出，它是基于HashPartitioner来进行分区的。HashPartitioner分区的原理很简单，对于给定的key，计算其hashCode，并除于分区的个数取余，如果余数小于0，则用余数+分区的个数，最后返回的值就是这个key所属的分区ID。

好了， 既然是基于Key分区的，那我们先给RDD搞一个Key，看看是不是按照我们的想法去进行分区的：

val repartition_coalesce_rdd = cartesian_rdd
  .zipWithIndex()
  .map(row=>(row._2 / 5,row._1))
  .repartition(10)

repartition_coalesce_rdd.mapPartitionsWithIndex((partid,iter)=>{
  iter.map(data => {
    println(partid + "," +data)
    data
  })
}).collect()

输出如下：

结果并不是如我们所想，那么是为什么呢？原因在于repartition所使用的Key，并非是RDD中每一条数据的Key，而是它为每条数据重新生成了一个随机数，作为此条数据的Key：

所以，在将一个分区较多的RDD重新分区为分区较少的RDD时，尽量使用coalesce算子。

8、map Vs mapPartitions

好了，我们已经重新将数据分为10个分区了，接下来就是并行调度了，这里有两个方法，一个是map，一个是mapPartitions。二者有什么区别呢：

map是对rdd中的每一个元素进行操作；mapPartitions则是对rdd中的每个分区的迭代器进行操作。假设我们想把RDD中所有的数变为两倍，如果共有5个Partition，每个Partition有1万条数据。那么map操作会执行5万次function，而mapPartitions操作只会执行5次function。因此mapPartitions的性能较高。

但要注意的一点是，MapPartitions操作，对于大量数据来说，比如一个Partition有100万数据，一次传入一个function以后，那么可能一下子内存不够，但是又没有办法去腾出内存空间来，可能就OOM，内存溢出。

而在我们的场景中，选择mapPartitions即可。二者代码示例如下：

def doubleMap(a:(Int,Int))= { (a._1 * 2,a._2 * 3) }

    cartesian_coalesce_rdd.map(
      data=>{
        doubleMap(data)
      }
    ).foreach(println)

def doubleMapPartition( iter : Iterator[(Int,Int)]) = {
  var res = List[(Int,Int)]()
  while (iter.hasNext){
    val cur = iter.next()
    res .::= (cur._1 * 2, cur._2 * 3)
  }
  res.iterator
}

cartesian_coalesce_rdd.mapPartitions(
  iter=>{
    doubleMapPartition(iter)
  }
).foreach(println)

注意，mapPartitions方法接受不了的是一个Iterator，返回也必须是一个Iterator。