(spark源码)union算子

摘要

1. 问题: spark中, union方法是否重新分区, 是否会触发shuffle
2. 结论: 不会shuffle, 不会划分stage, 但是可能重新分区(窄依赖)
3. 解释:
   (1) 宽窄依赖对应的原称为ShuffleDependency和NarrowDependency, 字面上可以看出来, 只有宽依赖才会发生shuffle. 但是两种依赖都会重新分区, 因此重分区和是否shuffle没有关系
   (2) 如果被union的多个rdd, 分区规则相同, 那么index相同的分区, 会被整合到多数分区所在的节点. 比如a节点有2个index为0的分区, b节点有1个, union之后的0分区会全部转移到a节点
   (3) 如果分区规则不同, union后生成的UnionRDD, 不会进行重新分区, 而是把每个分区合并记录到分区数组中

union方法

// 调用sparkContent的方法, 把this作为参数传进去. 基本操作
def union(other: RDD[T]): RDD[T] = withScope {
  sc.union(this, other)
}
// 实际调用了重载方法
def union[T: ClassTag](first: RDD[T], rest: RDD[T]*): RDD[T] = withScope {
  union(Seq(first) ++ rest)
}
// 解析分区器, 执行不同逻辑
def union[T: ClassTag](rdds: Seq[RDD[T]]): RDD[T] = withScope {
  // 把每个rdd的分区器取出来, 放到Set中去重
  val partitioners = rdds.flatMap(_.partitioner).toSet
  // 如果每个rdd都定义了分区器, 而且全部相同, 就使用此分区器整合所有rdd
  if (rdds.forall(_.partitioner.isDefined) && partitioners.size == 1) {
    new PartitionerAwareUnionRDD(this, rdds)
  } else {
    // 否则就保留各自的分区
    new UnionRDD(this, rdds)
  }
}

PartitionerAwareUnionRDD类

该类中, 数据重新分区, 但是是窄依赖

class PartitionerAwareUnionRDD[T: ClassTag](
    sc: SparkContext,
    var rdds: Seq[RDD[T]]
    // 注意这个继承, OneToOneDependency是窄依赖
  ) extends RDD[T](sc, rdds.map(x => new OneToOneDependency(x))) {
  // 因为保证了全部rdd的分区器一样, 所以head取第一个即可
  override val partitioner = rdds.head.partitioner
  // 整合rdds的分区
  override def getPartitions: Array[Partition] = {
    // 获取分区数
    val numPartitions = partitioner.get.numPartitions
    (0 until numPartitions).map { index =>
      // 将分区抽象为对象, 该类表示分区
      new PartitionerAwareUnionRDDPartition(rdds, index)
    }.toArray
  }
}

UnionRDD类

class UnionRDD[T: ClassTag](
    sc: SparkContext,
    var rdds: Seq[RDD[T]])
  // 传入的参数是空依赖, 但是自定义了依赖方法
  extends RDD[T](sc, Nil) {  // Nil since we implement getDependencies

  // 查看是否超出了并行列表阈值
  private[spark] val isPartitionListingParallel: Boolean =
    rdds.length > conf.getInt("spark.rdd.parallelListingThreshold", 10)

  override def getPartitions: Array[Partition] = {
    // 如果超出了并行列表阈值
    val parRDDs = if (isPartitionListingParallel) {
      // par方法, 用于获取和原集合并行的副本
      val parArray = rdds.par
      // 如果并行度太高, 就是靠线程池(fork join pool)拆分任务
      parArray.tasksupport = UnionRDD.partitionEvalTaskSupport
      parArray
    } else {
      rdds
    }
    
    // array的长度是rdds的分区数之和
    val array = new Array[Partition](parRDDs.map(_.partitions.length).seq.sum)
    var pos = 0
    // 双循环
    for ((rdd, rddIndex) <- rdds.zipWithIndex; split <- rdd.partitions) {
      // 参数是: (自增的角标, rdd, rdd在rdds中的角标, 原始分区的rdd的角标)
      // 根据这几个参数, 重新记录各分区. 实际上就是做了个和
      array(pos) = new UnionPartition(pos, rdd, rddIndex, split.index)
      pos += 1
    }
    array
  }
  
  // 自定义的获取依赖的方法
  override def getDependencies: Seq[Dependency[_]] = {
    val deps = new ArrayBuffer[Dependency[_]]
    var pos = 0
    for (rdd <- rdds) {
      // RangeDependency是窄依赖
      deps += new RangeDependency(rdd, 0, pos, rdd.partitions.length)
      pos += rdd.partitions.length
    }
    deps
  }
}