Spark withColumn 陷阱

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

总结

1. 多次执行 withColumn / withColumnRenamed 时，大部分时间都花费在 catalyse analyse 的反复调用上，且随着迭代次数的增加，逻辑计划的 Project 会增加，耗时会呈指数上升。
2. 完全可以使用 select 取代多次调用 withColumn / withColumnRenamed 的方式。