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())

而使用 select 的方式只會調用一次

進一步做了一個迭代次數和時間的關係測試，發現耗時並不是隨着次數線性增長，這是因爲每次迭代生成的邏輯計劃中會多增加一個 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 的方式。