withColumn / withColumnRenamed是 spark 中常用的 API，可以用于添加新字段 / 字段重命名 / 修改字段类型，但是当列的数量增加时，会出现严重的性能下降现象，本文将分析该现象出现的原因以及该如何解决它。
背景
在日常工作中，有时候会有建模或分析的同学问我，为什么用 withColumn / withColumnRenamed 会这么慢，明明数据量也不大，应该怎么解决。初步分析会发现，出现这种情况的时往往伴随着大量的列，难道是 spark 处理不了大宽表的场景吗？
现象及探究
对真实场景做了一个简化，下面是对一个10行的数据增加500列的一个操作，从代码上看好像没有什么问题，执行一下，却发现耗时14秒。
var df = spark.range(10).toDF()
for (i <- 1 to 500) {
df = df.withColumn("id_" + i, col("id") + i)
}
同样的逻辑使用 select 来实现，只需要0.1秒。
var df = spark.range(10).toDF()
df = df.select((1 to 500).map { i =>
(col("id") + i).as("id_" + i)
}: _*)
是什么导致了这么大差距，withColumn 时间花到哪去了？查看 withColumn 源码，每次执行完返回一个新的 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) 就能发现一些端倪，虽然他们最终生成的物理计划是一致的，但是逻辑计划存在着巨大的差异，使用 withColumn 方式的逻辑计划存在 500个 Project ，而 select 只有1个。
再用 RuleExecutor 查看 catalyst analysis 的统计信息，会发现 withColumn 中调用了 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())
withColumn 调用了500次 analyse
而使用 select 的方式只会调用一次
select 调用了1次 analyse
进一步做了一个迭代次数和时间的关系测试，发现耗时并不是随着次数线性增长，这是因为每次迭代生成的逻辑计划中会多增加一个 Project ，因此下一次的 analyse 时间会比上一次要长。
| 次数 | analyse 耗时(s) |
|---|---|
| 1 | 0.4 |
| 10 | 0.4 |
| 100 | 0.9 |
| 500 | 14 |
| 1000 | 65 |
总结
- 多次执行
withColumn / withColumnRenamed时，大部分时间都花费在 catalyse analyse 的反复调用上，且随着迭代次数的增加，逻辑计划的 Project 会增加，耗时会呈指数上升。 - 完全可以使用
select取代多次调用withColumn / withColumnRenamed的方式。
