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被计算,数据才会被缓存在内存中。