Spark Sql 之 cacheTable

1. 前言

spark sql中使用DataFrame/DataSet来抽象表示结构化数据(关系数据库中的table),DataSet上支持和RDD类似的操作,和RDD上的操作生成新的RDD一样,DataSet上的操作生成新的DataSet来表示新的数据抽象。最终DataSet上的这些操作经过:
logical plan -> analyzed logical plan -> optimized logical pan -> physical plan -> rdd dag的转化提交rdd 运行。这里plan(执行计划)就是DataSet上的转换操作,一个DataSet也就是对应一个logical plan生成的数据。

cacheTable也就是缓存DataSet抽象表示的数据,也就是DataSet的plan生成的数据。

2. cacheTable

从上面的介绍可以看出DataSet只是数据的抽象,它描述了从数据源头开始经过怎样的执行计划(plan)才能得到当前的DataSet表示的真实数据,也就是必须等到执行计划提交spark job运行结束后才能得到数据。spark实现cacheTable时,并没有立即提交table(DataSet)对应的plan去运行,然后得到运行结果数据去缓存,而是采用一种lazy模式:最终在DataSet上调用一些触发任务提交的方法时(类似RDD的action操作),发现plan对应的抽象语法树中发现子树是表缓存plan,如果这个时候数据已经缓存了,直接使用缓存的数据,没有则触发缓存表的plan去执行,然后采用按列缓存的方式缓存数据。

看看代码实现:
调用SQLContext # cacheTable(tableName : String)最终会走到下面的调用:

// query 即缓存的Dataset
// storageLevel 可以使用memory和disk缓存
def cacheQuery(
      query: Dataset[_],
      tableName: Option[String] = None,
      storageLevel: StorageLevel = MEMORY_AND_DISK): Unit = writeLock {
    // 拿到dataset的plan
    val planToCache = query.logicalPlan
    if (lookupCachedData(planToCache).nonEmpty) {
      logWarning("Asked to cache already cached data.")
    } else {
      val sparkSession = query.sparkSession
      // 缓存是建立一个plan到InMemoryRelation的映射
      cachedData.add(CachedData(
        planToCache,
        InMemoryRelation(
          sparkSession.sessionState.conf.useCompression,
          sparkSession.sessionState.conf.columnBatchSize,
          storageLevel,
          sparkSession.sessionState.executePlan(planToCache).executedPlan,
          tableName)))
    }
  }

上面InMemoryRelation是执行计划中一个节点,当出现select * from table_a语句(或者任何逻辑执行计划中有table_a出现),假设table_a被缓存了,那么这条语句生成的逻辑执行计划中,table_a对应的Relation节点会被执行计划优化器(optimizer)替换成InMemoryRelation。

InMemoryRelation构造参数:

  • columnBatchSize,后面会提到,table缓存是按列缓存的,然后数据又被按行分为一个个batch,这个参数用来控制一个batch里行数,通过配置项spark.sql.inMemoryColumnarStorage.batchSize设置,默认是10000行。

下面是InMemoryRelation中和缓存相关的代码:


  private def buildBuffers(): Unit = {
    // output输出的是Seq[Attribute],也就是表的schema,包含所有列名,列类型等信息
    // child也就是缓存的dataset对应的plan
    val output = child.output
    // 调用逻辑执行计划的execute返回的是RDD[InternalRow],返回RDD是整个执行计划分析的最后一步了,接下来就rdd的提交运行。那么这个rdd也就是dataset表示的数据的rdd形式的抽象。
   // 这里在rdd上调用mapPartitionsInternal,实现的是将遍历每一行数据,然后按列缓存。
   // 这个地方返回新的RDD的数据类型CachedBatch,CachedBatch是一个batch内若干行上的按列缓存。
    val cached = child.execute().mapPartitionsInternal { rowIterator =>
      new Iterator[CachedBatch] {
        def next(): CachedBatch = {
          // 按照每一列的类型生成ColumnBuilder,内部使用数组来保存列数据
          val columnBuilders = output.map { attribute =>
            ColumnBuilder(attribute.dataType, batchSize, attribute.name, useCompression)
          }.toArray

          var rowCount = 0
          var totalSize = 0L
         // 遍历每一行数据,控制当前batch行数 rowCount不超过batchSize,且同时batch中数据大小不超过MAX_BATCH_SIZE_IN_BYTE(4MB)
          while (rowIterator.hasNext && rowCount < batchSize
            && totalSize < ColumnBuilder.MAX_BATCH_SIZE_IN_BYTE) {
            val row = rowIterator.next()

            assert(
              row.numFields == columnBuilders.length,
              s"Row column number mismatch, expected ${output.size} columns, " +
                s"but got ${row.numFields}." +
                s"\nRow content: $row")

            var i = 0
            totalSize = 0
            while (i < row.numFields) {
              columnBuilders(i).appendFrom(row, i)
              totalSize += columnBuilders(i).columnStats.sizeInBytes
              i += 1
            }
            rowCount += 1
          }

          batchStats.add(totalSize)

          val stats = InternalRow.fromSeq(columnBuilders.map(_.columnStats.collectedStatistics)
            .flatMap(_.values))
          CachedBatch(rowCount, columnBuilders.map { builder =>
            JavaUtils.bufferToArray(builder.build())
          }, stats)
        }

        def hasNext: Boolean = rowIterator.hasNext
      }
      // 调用persist缓存RDD,所以cacheTable最终还是调用rdd的缓存接口完成缓存的
    }.persist(storageLevel)

    cached.setName(
      tableName.map(n => s"In-memory table $n")
        .getOrElse(StringUtils.abbreviate(child.toString, 1024)))
  
    _cachedColumnBuffers = cached
  }

上面代码可以看出cacheTable实际上还是通过cache rdd实现的。上面InMemoryRelation只是逻辑执行计划中一个节点,逻辑执行计划需要转换成物理执行计划,再转换成RDD dag才能执行,加上spark中RDD的计算是lazy模式的,所以上面的缓存rdd并没有提交运行,所以数据还没有缓存下来。

真正缓存还得看InMemoryRelation所在的执行计划真正提交后,这个缓存rdd被计算,数据才会被缓存在内存中。

推荐阅读更多精彩内容