spark源码（七）RDD cache、persist、checkpoint功能区别和源码解析

1、理论上功能介绍

cache：将数据缓存到本地内存中。cache底层调用的也是persist方法，只不过其设置的存储级别为MEMEORY_ONLY

persist：将数据存储到本地内存或者本地磁盘中。目前常用的持久化级别有：
        DISK_ONLY：持久化到磁盘
        MEMORY_ONLY：持久化到本地内存
        MEMORY_ONLY_SER：序列化后持久到本地内存
        MEMORY_AND_DISK：内存放不下，则将数据溢写到磁盘上
        MEMORY_AND_DISK_SER：内存放序列化后数据，如果放不下，则将数据溢写到磁盘上
        DISK_ONLY：数据持久化到本地磁盘
如果要缓存的数据太多，内存中放不下， Spark 会自动利用最近最少使用（LRU）的缓存策略把最老的分区从内存中移除。但是对于使用 MEMORY_AND_DISK 缓存级别的分区来说，被移除的分区都会写入磁盘。

checkpoint：将数据持久化到分布式存储系统中，常用的是HDFS。它可以避免节点故障导致的持久化数据丢失（如persist将持久化数据存储在本地内存和磁盘中，一旦节点故障，数据将丢失），可用性和容错性更高。

下面从源码角度看一下三者的功能实现。

2、源码展示

2.1 cache

cache就是将数据持久化到本地内存中，它底层调用的还是persist，只不过其设置的存储级别为MEMORY_ONLY级别，所以这里不过多介绍，统一放到persist中讲：

2.2 persist

2.2.1 设置存储级别

persist是将数据持久化到本地，具体是持久化到本地内存还是本地磁盘则根据设置的级别来定，下面我们看下persist方法的执行逻辑：

//重点一：在RDD被首次计算出来后，根据设置的存储级别进行持久化  
def persist(newLevel: StorageLevel): this.type = {
    //重点二：判断该RDD是否需要checkpoint
    if (isLocallyCheckpointed) {
      //重点三：如果需要checkpoint，则转换存储级别为disk，不然一些存储在内存的数据如果在job执行中被清理掉，那么checkpoint还需要再重新计算。
      persist(LocalRDDCheckpointData.transformStorageLevel(newLevel), allowOverride = true)
    } else {
      persist(newLevel, allowOverride = false)
    }
  }

可以看到persist方法中虽然代码行数不多，但是要注意的细节点很多：

重点一：首先是方法本身的注释，该注释提醒我们真正的持久化要在RDD计算出来的时候执行，所以这块的方法只是设置一个存储级别，真正的持久化逻辑要到iterator中去看。

重点二：isLocallyCheckpointed方法实现很简单，如下：

 这里就是判断当前RDD是否需要checkpoint，这里有个潜在的使用逻辑，那就是如果同一个RDD继续用持久化到本地内存或磁盘，又需要checkpoint到hdfs一类的分布式文件存储系统中，那么最好是先checkpoint，再persist，不然persist时很大可能仍会重新计算当前RDD。

重点三：如果需要checkpoint，则转换存储级别为disk，不然一些存储在内存的数据如果在job执行中被清理掉，那么checkpoint还需要再重新计算。

 回过头我们继续深入看下persist方法：

  private def persist(newLevel: StorageLevel, allowOverride: Boolean): this.type = {
    if (storageLevel != StorageLevel.NONE && newLevel != storageLevel && !allowOverride) {
      throw new UnsupportedOperationException(
        "Cannot change storage level of an RDD after it was already assigned a level")
    }
    // RDD默认情况下storageLevel 为NONE，这表明是第一次对该RDD设置持久化，此时将该RDD注册到sparkContext上下文中
    if (storageLevel == StorageLevel.NONE) {
      sc.cleaner.foreach(_.registerRDDForCleanup(this))
      sc.persistRDD(this)
    }
    //设置当前该RDD的存储级别为指定的存储级别
    storageLevel = newLevel
    this
  }

可以看到这块的的代码也很简单，就是根据是否初次设置持久化，将该RDD记录到sparkContext中，然后再更新下存储级别变量即可。接下来我们到iterator中看下具体的持久化逻辑。

2.2.2 具体的计算处理

首先我们到分区计算的入口iterator方法中看下：

  final def iterator(split: Partition, context: TaskContext): Iterator[T] = {
    // 如果存储级别不为NONE，则进入getOrCompute，否则进入computeOrReadCheckpoint
    if (storageLevel != StorageLevel.NONE) {
      getOrCompute(split, context)
    } else {
      computeOrReadCheckpoint(split, context)
    }
  }

因为我们追踪的是设置存储级别的处理逻辑，所以这里我们只看getOrCompute方法的处理逻辑，另外一个方法不是我们研究的重点，这里不再细述，感兴趣的可以自己看下，好了进入正文：

  private[spark] def getOrCompute(partition: Partition, context: TaskContext): Iterator[T] = {
    val blockId = RDDBlockId(id, partition.index)
    var readCachedBlock = true
    // 通过blockManager获取或者更新数据
    SparkEnv.get.blockManager.getOrElseUpdate(blockId, storageLevel, elementClassTag, () => {
      readCachedBlock = false
      computeOrReadCheckpoint(partition, context)
    }) match {
      //根据BlockManager返回的结果的不同，统一封装成InterruptibleIterator迭代对象
      case Left(blockResult) =>
        if (readCachedBlock) {
          val existingMetrics = context.taskMetrics().inputMetrics
          existingMetrics.incBytesRead(blockResult.bytes)
          new InterruptibleIterator[T](context, blockResult.data.asInstanceOf[Iterator[T]]) {
            override def next(): T = {
              existingMetrics.incRecordsRead(1)
              delegate.next()
            }
          }
        } else {
          new InterruptibleIterator(context, blockResult.data.asInstanceOf[Iterator[T]])
        }
      case Right(iter) =>
        new InterruptibleIterator(context, iter.asInstanceOf[Iterator[T]])
    }
  }

这块的逻辑有两块，一个是通过BlockManager获取或者更新数据，一个是对获取的结果封装返回，结果的封装返回跟我们的源码阅读目的关联不大，这里不再细述，我们深入看下BlockManager的处理逻辑：

  def getOrElseUpdate[T](
      blockId: BlockId,
      level: StorageLevel,
      classTag: ClassTag[T],
      makeIterator: () => Iterator[T]): Either[BlockResult, Iterator[T]] = {
    // 通过BlockManager从本地或者远程获取数据，如果获取到则直接返回
    get[T](blockId)(classTag) match {
      case Some(block) =>
        return Left(block)
      case _ =>
    }
    // 通过BlockManager没有获取到数据时，通过makeIterator重新计算并持久化数据
    doPutIterator(blockId, makeIterator, level, classTag, keepReadLock = true) match {
      case None =>
        val blockResult = getLocalValues(blockId).getOrElse {
          releaseLock(blockId)
          throw new SparkException(s"get() failed for block $blockId even though we held a lock")
        }
        releaseLock(blockId)
        Left(blockResult)
      case Some(iter) =>
       Right(iter)
    }
  }

这里是看BlockManager能不能从本地或远程获取到数据，如果能获取到则直接返回，否则调用doPutIterator进一步处理：

private def doPutIterator[T](
      blockId: BlockId,
      iterator: () => Iterator[T],
      level: StorageLevel,
      classTag: ClassTag[T],
      tellMaster: Boolean = true,
      keepReadLock: Boolean = false): Option[PartiallyUnrolledIterator[T]] = {
    doPut(blockId, level, classTag, tellMaster = tellMaster, keepReadLock = keepReadLock) { info =>
      val startTimeNs = System.nanoTime()
      var iteratorFromFailedMemoryStorePut: Option[PartiallyUnrolledIterator[T]] = None
      var size = 0L
      // 如果用到内存，则首先将数据存储到内存，当内存存不下时，再写到磁盘上
      if (level.useMemory) {
        // 序列化和非序列化分开处理
        if (level.deserialized) {
          // 通过memoryStore将数据存储到内存中
          memoryStore.putIteratorAsValues(blockId, iterator(), classTag) match {
            case Right(s) =>
              size = s
            case Left(iter) =>
              // 如果内存不够，则溢写到磁盘上
              if (level.useDisk) {
                logWarning(s"Persisting block $blockId to disk instead.")
                // 通过 diskStore 将数据存储到磁盘上
                diskStore.put(blockId) { channel =>
                  val out = Channels.newOutputStream(channel)
                  serializerManager.dataSerializeStream(blockId, out, iter)(classTag)
                }
                size = diskStore.getSize(blockId)
              } else {
                iteratorFromFailedMemoryStorePut = Some(iter)
              }
          }
        } else { 
          memoryStore.putIteratorAsBytes(blockId, iterator(), classTag, level.memoryMode) match {
            case Right(s) =>
              size = s
            case Left(partiallySerializedValues) =>
              if (level.useDisk) {
                logWarning(s"Persisting block $blockId to disk instead.")
                diskStore.put(blockId) { channel =>
                  val out = Channels.newOutputStream(channel)
                  partiallySerializedValues.finishWritingToStream(out)
                }
                size = diskStore.getSize(blockId)
              } else {
                iteratorFromFailedMemoryStorePut = Some(partiallySerializedValues.valuesIterator)
              }
          }
        }

      } else if (level.useDisk) {
        diskStore.put(blockId) { channel =>
          val out = Channels.newOutputStream(channel)
          serializerManager.dataSerializeStream(blockId, out, iterator())(classTag)
        }
        size = diskStore.getSize(blockId)
      }

      val putBlockStatus = getCurrentBlockStatus(blockId, info)
      val blockWasSuccessfullyStored = putBlockStatus.storageLevel.isValid
      if (blockWasSuccessfullyStored) {
        info.size = size
        if (tellMaster && info.tellMaster) {
          reportBlockStatus(blockId, putBlockStatus)
        }
        addUpdatedBlockStatusToTaskMetrics(blockId, putBlockStatus)
        logDebug(s"Put block $blockId locally took ${Utils.getUsedTimeNs(startTimeNs)}")
        if (level.replication > 1) {
          val remoteStartTimeNs = System.nanoTime()
          val bytesToReplicate = doGetLocalBytes(blockId, info)
          val remoteClassTag = if (!serializerManager.canUseKryo(classTag)) {
            scala.reflect.classTag[Any]
          } else {
            classTag
          }
          try {
            replicate(blockId, bytesToReplicate, level, remoteClassTag)
          } finally {
            bytesToReplicate.dispose()
          }
          logDebug(s"Put block $blockId remotely took ${Utils.getUsedTimeNs(remoteStartTimeNs)}")
        }
      }
      assert(blockWasSuccessfullyStored == iteratorFromFailedMemoryStorePut.isEmpty)
      iteratorFromFailedMemoryStorePut
    }
  }

该方法的处理逻辑很多，我们并没有全部注释，但是关键点也就是我们注释的那几行。大体流程是查看是否有内存存储，如果有则优先存储到内存中，当内存不够时再将数据写到磁盘中。而这些写操作功能的实现全依赖于BlockManager的memoryStore和diskStore。

至此，我们persist方法的介绍基本结束，我们总结下：persist持久化的大体逻辑是先设置存储级别，然后RDD计算的时候会通过BlockManager检查本地或者远程是否已有数据直接使用，如果有则直接返回，如果没有则检查是否有checkpoint数据使用，如果有则使用，如果没有则进行计算，并根据存储级别，通过BlockManager的memoryStore和diskStore将数据持久化到本地。

2.3 checkpoint

这块其实是我写这篇文章的初始目的，当时想探究checkpoint的两个技术点，一是使用checkpoint的时，BlockManager如何将数据存储到hdfs中；二是如果RDD没有持久化，则checkpoint需要重新计算RDD数据并存储到HDFS中。

在具体介绍前，我们先大致讲解下checkpoint的使用流程，不然可能追踪源码都不知道去哪追踪。如下：

即使用checkpoint，首先要想使用checkpoint需要在HDFS上设置一个checkpoint目录，该目录用于存储持久化的数据，然后调用RDD的checkpoint方法，将rdd的数据进行持久化。我们分别深入看一下这两步代码：

 第一步设置目录很简单，就是通过FileSystem在hdfs上创建设定的目录。

 第二步也很简单，就是返回一个封装的RDD数据。

到这你可能会很懵，checkpoint的存储逻辑呢，其实checkpoint会在RDD所处的job运行结束之后，启动一个单独的job来将checkpoint过的RDD的数据写入之前设置的文件系统。我们到runJob（RDD的行为操作会导致sparkContext调用runJob方法触发job的运行，前面的系列文章有介绍，不清楚的可以看看我前面的文章）的方法中看一下：

 可以看到在DAGScheduler下发执行完所有的计算流程后，会调用rdd的doCheckpoint方法，这个方法才是核心存储方法的入口，下面我们看一下：

//它将存储RDD数据成一个文件到checkpoint目录中，并且该RDD的所有父依赖关系都将被移除。这个函数会再job执行调用完该RDD后执行  
private[spark] def doCheckpoint(): Unit = {
    RDDOperationScope.withScope(sc, "checkpoint", allowNesting = false, ignoreParent = true) {
      //checkpoint过则不再进行处理，防止被多次处理
      if (!doCheckpointCalled) {
        //设置处理标识位true
        doCheckpointCalled = true
        //如果没有封装后的checkpointData数据，则当前rdd不再处理，直接遍历其父rdd
        if (checkpointData.isDefined) {
          if (checkpointAllMarkedAncestors) {
            //遍历rdd的父类，调用它们的doCheckpoint方法
            dependencies.foreach(_.rdd.doCheckpoint())
          }
          //获取封装后的rdd数据，然后checkpoint存储
          checkpointData.get.checkpoint()
        } else {
          dependencies.foreach(_.rdd.doCheckpoint())
        }
      }
    }
  }

 可以看到，checkpoint会倒着从最后一个rdd开始，一一遍历其所有的父rdd，当父rdd处理结束，且当前节点有封装后的checkpointData数据时，再进行checkpoint处理。这里有两个点关注下，一是之所以要先处理父RDD，是因为checkpoint的计算会斩断rdd的血缘关系，故要先处理父RDD。二是checkpointData是我们程序中调用checkpoint方法时创建ReliableRDDCheckpointData（checkpoint使用流程第二步有介绍），所以checkpointData.get.checkpoint()最终调用的是ReliableRDDCheckpointData中的方法，下面我们到源码中看一看：

  final def checkpoint(): Unit = {
    // 通过全局锁，安全的修改checkpoint的状态，如果是初始状态则修改为处理态，其它状态就直接返回
    RDDCheckpointData.synchronized {
      if (cpState == Initialized) {
        cpState = CheckpointingInProgress
      } else {
        return
      }
    }
    // 核心方法，将RDD写到分布式文件系统中，并返回一个新的RDD，该RDD没有父依赖信息
    val newRDD = doCheckpoint()

    // 更新checkpoint的状态，并移除原有rdd的分区和依赖信息
    RDDCheckpointData.synchronized {
      cpRDD = Some(newRDD)
      cpState = Checkpointed
      //移除原有rdd的分区和依赖信息
      rdd.markCheckpointed()
    }
  }

可以看到checkpoint方法会继续调用doCheckPoint方法将数据写入分布式文件并返回一个新的RDD，且该新RDD没有父依赖信息。另外还会清除当前RDD的依赖关系，具体的实现方法在markCheckpointed中，因为逻辑比较简单，这里就不展示出来了，感兴趣的可以自己追踪点击看看。我们接着进入到核心方法doCheckPoint看下后续的处理逻辑：

  protected override def doCheckpoint(): CheckpointRDD[T] = {
    //将rdd写入到设置的checkpoint目录中
    val newRDD = ReliableCheckpointRDD.writeRDDToCheckpointDirectory(rdd, cpDir)

    // Optionally clean our checkpoint files if the reference is out of scope
    if (rdd.conf.get(CLEANER_REFERENCE_TRACKING_CLEAN_CHECKPOINTS)) {
      rdd.context.cleaner.foreach { cleaner =>
        cleaner.registerRDDCheckpointDataForCleanup(newRDD, rdd.id)
      }
    }

    logInfo(s"Done checkpointing RDD ${rdd.id} to $cpDir, new parent is RDD ${newRDD.id}")
    newRDD
  }

继续深入：

  def writeRDDToCheckpointDirectory[T: ClassTag](
      originalRDD: RDD[T],
      checkpointDir: String,
      blockSize: Int = -1): ReliableCheckpointRDD[T] = {
    val checkpointStartTimeNs = System.nanoTime()

    val sc = originalRDD.sparkContext

    // 创建要输出的路径对象
    val checkpointDirPath = new Path(checkpointDir)
    val fs = checkpointDirPath.getFileSystem(sc.hadoopConfiguration)
    if (!fs.mkdirs(checkpointDirPath)) {
      throw new SparkException(s"Failed to create checkpoint path $checkpointDirPath")
    }

    val broadcastedConf = sc.broadcast(
      new SerializableConfiguration(sc.hadoopConfiguration))
    
    // 发起一个job任务写RDD数据到checkpoint目录
    sc.runJob(originalRDD,
      writePartitionToCheckpointFile[T](checkpointDirPath.toString, broadcastedConf) _)

    // 根据自定义的分区器将checkpoint目录中的数据进行重写（根据源码推测的含义，如果不对，欢迎指正）
    if (originalRDD.partitioner.nonEmpty) {
      writePartitionerToCheckpointDir(sc, originalRDD.partitioner.get, checkpointDirPath)
    }

    val checkpointDurationMs =
      TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - checkpointStartTimeNs)
    logInfo(s"Checkpointing took $checkpointDurationMs ms.")

    //通过checkpoint目录下的文件信息创建新的RDD
    val newRDD = new ReliableCheckpointRDD[T](
      sc, checkpointDirPath.toString, originalRDD.partitioner)
    if (newRDD.partitions.length != originalRDD.partitions.length) {
      throw new SparkException(
        "Checkpoint RDD has a different number of partitions from original RDD. Original " +
          s"RDD [ID: ${originalRDD.id}, num of partitions: ${originalRDD.partitions.length}]; " +
          s"Checkpoint RDD [ID: ${newRDD.id}, num of partitions: " +
          s"${newRDD.partitions.length}].")
    }
    newRDD
  }

这块的重点是起了一个job任务用于checkpoint数据的写入，还有一个重点则是checkpoint数据写入分布式文件系统后，还可以根据自定义的分区器重写checkpoint的数据（这块是根据源码推测出的大致功能，还不是很确定，如果有知道的大佬，欢迎指正，我后期构思如何验证后，也会给一个定论）。接下来我们看下runJob中核心方法writePartitionToCheckpointFile的处理逻辑：

def writePartitionToCheckpointFile[T: ClassTag](
      path: String,
      broadcastedConf: Broadcast[SerializableConfiguration],
      blockSize: Int = -1)(ctx: TaskContext, iterator: Iterator[T]): Unit = {
    //写文件前的准备操作，如创建输出路径，文件名称，创建输出流等信息
    val env = SparkEnv.get
    val outputDir = new Path(path)
    val fs = outputDir.getFileSystem(broadcastedConf.value.value)

    val finalOutputName = ReliableCheckpointRDD.checkpointFileName(ctx.partitionId())
    val finalOutputPath = new Path(outputDir, finalOutputName)
    val tempOutputPath =
      new Path(outputDir, s".$finalOutputName-attempt-${ctx.attemptNumber()}")

    val bufferSize = env.conf.get(BUFFER_SIZE)

    val fileOutputStream = if (blockSize < 0) {
      val fileStream = fs.create(tempOutputPath, false, bufferSize)
      if (env.conf.get(CHECKPOINT_COMPRESS)) {
        CompressionCodec.createCodec(env.conf).compressedOutputStream(fileStream)
      } else {
        fileStream
      }
    } else {
      fs.create(tempOutputPath, false, bufferSize,
        fs.getDefaultReplication(fs.getWorkingDirectory), blockSize)
    }
    val serializer = env.serializer.newInstance()
    val serializeStream = serializer.serializeStream(fileOutputStream)
    Utils.tryWithSafeFinally {
      // 通过序列化流将RDD数据写入checkpoint目录中
      serializeStream.writeAll(iterator)
    } {
      serializeStream.close()
    }

    if (!fs.rename(tempOutputPath, finalOutputPath)) {
      if (!fs.exists(finalOutputPath)) {
        logInfo(s"Deleting tempOutputPath $tempOutputPath")
        fs.delete(tempOutputPath, false)
        throw new IOException("Checkpoint failed: failed to save output of task: " +
          s"${ctx.attemptNumber()} and final output path does not exist: $finalOutputPath")
      } else {
        logInfo(s"Final output path $finalOutputPath already exists; not overwriting it")
        if (!fs.delete(tempOutputPath, false)) {
          logWarning(s"Error deleting ${tempOutputPath}")
        }
      }
    }
  }

这块有两点需要我们留意，第一点也就是最重要的一点，就是通过流写入的时候传入的是iterator迭代器，通过前面persist的源码我们可以知道，迭代器中会去查找RDD是否已经持久化到本地，如果通过BlockManager获取到了RDD数据，则直接返回，如果获取不到再进行重新计算（如果不理解，可以再读读前面persist的源码）。第二点就是checkpoint数据的写入是通过序列化流实现的，而不是通过BlockManager实现的（闲扯（未进行源码验证，但是理论上spark是这么处理的）：BlockManager的生命周期属于application，所以应用结束的时候persist持久化到本地内存或者磁盘的数据都会删除，而通过流写入checkpoint的数据则没有生命周期的困扰，所以其在application结束后数据仍然存在）。

至此，checkpoint的逻辑也算处理结束了，下面我们做个简单的总结。

3、总结

1、cache底层调用的也是persist方法，只不过其设置的存储级别为MEMEORY_ONLY。

2、Persist 和 Cache，不会丢掉RDD间的依赖链/依赖关系，CheckPoint会斩断依赖链。

3、checkpoint首先要调用SparkContext的setCheckPointDir方法，设置一个容错的文件系统的目录，比如说HDFS；然后对RDD调用checkpoint方法。之后在RDD所处的job运行结束之后，会启动一个单独的job，来将checkpoint过的RDD的数据写入之前设置的文件系统，进行高可用、容错的类持久化操作。

4、如果准备使用checkpoint，推荐将此 RDD 持久化在内存或者磁盘中，以供checkpoint直接使用，否则checkpoint将其保存在HDFS文件系统中将需要重新计算。注意在代码调用顺序上尽量保证checkpoint在persist之前。

5、无论cache、persist、checkpoint，使用他们都会占用资源，所以在持久化时不仅要考虑Lineage是否足够长，也要考虑是否有宽依赖，对宽依赖持久化是最物有所值的。

4、引申：

DataFrame 的cache依然调用的persist，但是persist调用cacheQuery，而cacheQuery的默认存储级别为MEMORY_AND_DISK，这点和rdd是不一样的。（待确认）