map算子源码分析

map算子源码解析

• map算子可以对rdd中的元素进行一对一的映射。如下表示将rdd每个元素乘以2，并且map之后新的rdd分区数不会改变。

val rdd1 = sc.parallelize(List(1,2,3,4,5))
rdd1.map(a=>a*2).collect().foreach(println(_))

• 现在分析一下map算子的底层实现

def map[U: ClassTag](f: T => U): RDD[U] = withScope {
  val cleanF = sc.clean(f)
  new MapPartitionsRDD[U, T](this, (context, pid, iter) => iter.map(cleanF))
}

• 可以看出，map算子接收一个函数类型的参数f，f类型为T => U，表示将T类型的元素转换为U类型的元素。
• 在map算子内部，首先将函数f传入clean方法，先看看clean方法都干了啥，clean方法除了接收函数f还有一个默认参数表示是否进行序列化检查，这很好理解因为函数需要序列化才能分发给每一个task。

private[spark] def clean[F <: AnyRef](f: F, checkSerializable: Boolean = true): F = {
  ClosureCleaner.clean(f, checkSerializable)
  f
}

• 继续跟踪ClosureCleaner.clean方法

def clean(
    closure: AnyRef,
    checkSerializable: Boolean = true,
    cleanTransitively: Boolean = true): Unit = {
  clean(closure, checkSerializable, cleanTransitively, Map.empty)
}

• 继续跟踪，看看究竟做了什么事情

private def clean(
    func: AnyRef,
    checkSerializable: Boolean,
    cleanTransitively: Boolean,
    accessedFields: Map[Class[_], Set[String]]): Unit = {

  val maybeIndylambdaProxy = IndylambdaScalaClosures.getSerializationProxy(func)
   //如果当前函数不是闭包函数  返回
  if (!isClosure(func.getClass) && maybeIndylambdaProxy.isEmpty) {
    logDebug(s"Expected a closure; got ${func.getClass.getName}")
    return
  }
  //否则进行一些清理工作，具体什么工作暂时不细究了
  // TODO: clean all inner closures first. This requires us to find the inner objects.
  // TODO: cache outerClasses / innerClasses / accessedFields

• 然后回头看map方法，clean函数会返回一个cleanF对象，然后创建一个MapPartitionsRDD返回

def map[U: ClassTag](f: T => U): RDD[U] = withScope {
  val cleanF = sc.clean(f)
  new MapPartitionsRDD[U, T](this, (context, pid, iter) => iter.map(cleanF))
}

• 跟踪一下MapPartitionsRDD这个类，其中的getPartitions方法调用了firstParent[T]的分区
  override def getPartitions: Array[Partition] = firstParent[T].partitions

• 通过查看firstParent源码可以看到，新的rdd分区取决于rdd依赖中上一个rdd的分区信息，这也就验证了map算子不会改变分区数量

protected[spark] def firstParent[U: ClassTag]: RDD[U] = {
  dependencies.head.rdd.asInstanceOf[RDD[U]]
}

• 我们在创建MapPartitionsRDD对象时会传入一个函数和上一个rdd对象this
• 然后跟一下MapPartitionsRDD的构造方法，可以看到prev就是上一个rdd，f就是传入的函数

private[spark] class MapPartitionsRDD[U: ClassTag, T: ClassTag](
    var prev: RDD[T],
    f: (TaskContext, Int, Iterator[T]) => Iterator[U],  // (TaskContext, partition index, iterator)
    preservesPartitioning: Boolean = false,
    isFromBarrier: Boolean = false,
    isOrderSensitive: Boolean = false)
  extends RDD[U](prev)

• 然后仔细看看传入的函数f到底在哪用到了，可以发现compute方法调用了f，也就是说明真正的map算子转换逻辑在compute方法实现的。函数f中封装了用户调用map算子传入的那个转换函数。

override def compute(split: Partition, context: TaskContext): Iterator[U] =
  f(context, split.index, firstParent[T].iterator(split, context))