withColumn / withColumnRenamed是 spark 中常用的 API抡秆，可以用于添加新字段 / 字段重命名 / 修改字段類型坊夫，但是當列的數(shù)量增加時逮壁，會出現(xiàn)嚴重的性能下降現(xiàn)象村斟，本文將分析該現(xiàn)象出現(xiàn)的原因以及該如何解決它奶陈。
背景
在日常工作中怯屉，有時候會有建奈狄ǎ或分析的同學問我，為什么用 withColumn / withColumnRenamed 會這么慢锨络，明明數(shù)據(jù)量也不大赌躺，應該怎么解決。初步分析會發(fā)現(xiàn)羡儿，出現(xiàn)這種情況的時往往伴隨著大量的列礼患，難道是 spark 處理不了大寬表的場景嗎？
現(xiàn)象及探究
對真實場景做了一個簡化掠归，下面是對一個10行的數(shù)據(jù)增加500列的一個操作缅叠，從代碼上看好像沒有什么問題，執(zhí)行一下拂到，卻發(fā)現(xiàn)耗時14秒痪署。
var df = spark.range(10).toDF()
for (i <- 1 to 500) {
df = df.withColumn("id_" + i, col("id") + i)
}
同樣的邏輯使用 select 來實現(xiàn)，只需要0.1秒兄旬。
var df = spark.range(10).toDF()
df = df.select((1 to 500).map { i =>
(col("id") + i).as("id_" + i)
}: _*)
是什么導致了這么大差距狼犯，withColumn 時間花到哪去了？查看 withColumn 源碼领铐，每次執(zhí)行完返回一個新的 DataFrame悯森，好像也沒有什么問題 。
def withColumn(colName: String, col: Column): DataFrame = withColumns(Seq(colName), Seq(col))
private[spark] def withColumns(colNames: Seq[String], cols: Seq[Column]): DataFrame = {
require(colNames.size == cols.size,
s"The size of column names: ${colNames.size} isn't equal to " +
s"the size of columns: ${cols.size}")
SchemaUtils.checkColumnNameDuplication(
colNames,
"in given column names",
sparkSession.sessionState.conf.caseSensitiveAnalysis)
val resolver = sparkSession.sessionState.analyzer.resolver
val output = queryExecution.analyzed.output
val columnMap = colNames.zip(cols).toMap
val replacedAndExistingColumns = output.map { field =>
columnMap.find { case (colName, _) =>
resolver(field.name, colName)
} match {
case Some((colName: String, col: Column)) => col.as(colName)
case _ => Column(field)
}
}
val newColumns = columnMap.filter { case (colName, col) =>
!output.exists(f => resolver(f.name, colName))
}.map { case (colName, col) => col.as(colName) }
select(replacedAndExistingColumns ++ newColumns : _*)
}
使用 df.explain(true) 就能發(fā)現(xiàn)一些端倪绪撵，雖然他們最終生成的物理計劃是一致的瓢姻，但是邏輯計劃存在著巨大的差異，使用 withColumn 方式的邏輯計劃存在 500個 Project 音诈，而 select 只有1個幻碱。
再用 RuleExecutor 查看 catalyst analysis 的統(tǒng)計信息，會發(fā)現(xiàn) withColumn 中調(diào)用了 500 次 analyse细溅，情況逐漸開始明朗了褥傍。
import org.apache.spark.sql.catalyst.rules.RuleExecutor
var df = spark.range(10).toDF()
RuleExecutor.resetMetrics()
for (i <- 1 to 500) {
df = df.withColumn("id_" + i, col("id") + i)
}
println(RuleExecutor.dumpTimeSpent())
而使用 select 的方式只會調(diào)用一次
進一步做了一個迭代次數(shù)和時間的關系測試，發(fā)現(xiàn)耗時并不是隨著次數(shù)線性增長喇聊，這是因為每次迭代生成的邏輯計劃中會多增加一個 Project 恍风，因此下一次的 analyse 時間會比上一次要長。
| 次數(shù) | analyse 耗時(s) |
|---|---|
| 1 | 0.4 |
| 10 | 0.4 |
| 100 | 0.9 |
| 500 | 14 |
| 1000 | 65 |
總結
- 多次執(zhí)行
withColumn / withColumnRenamed時，大部分時間都花費在 catalyse analyse 的反復調(diào)用上朋贬，且隨著迭代次數(shù)的增加凯楔，邏輯計劃的 Project 會增加，耗時會呈指數(shù)上升锦募。 - 完全可以使用
select取代多次調(diào)用withColumn / withColumnRenamed的方式摆屯。
