如何重命名Spark DataFrame中已存在的重复列？

当拿到的DataFrame本身存在重复列、且无法控制上游生成逻辑时，无法直接通过withColumnRenamed按名称修改列名（重名时列名匹配存在歧义），可通过按列索引批量重命名的方式解决。

通用实现代码 #

import org.apache.spark.sql.DataFrame

def renameDuplicateColumns(df: DataFrame): DataFrame = {
  // 遍历列名生成带序号后缀的新列名列表
  val newColumnNames = df.columns.foldLeft((List.empty[String], Map.empty[String, Int])) {
    case ((processedNames, countRecorder), currentCol) =>
      val appearCount = countRecorder.getOrElse(currentCol, 0)
      val newName = if (appearCount == 0) currentCol else s"${currentCol}_${appearCount + 1}"
      (processedNames :+ newName, countRecorder + (currentCol -> (appearCount + 1)))
  }._1
  // 按位置一一对应重命名列
  df.toDF(newColumnNames: _*)
}

测试效果 #

对应用户给出的示例数据，调用方法如下：

val data = Seq((1,2),(3,4)).toDF("a","a")
val processedDf = renameDuplicateColumns(data)
processedDf.show()

输出结果符合预期：

+---+---+
|  a|a_2|
+---+---+
|  1|  2|
|  3|  4|
+---+---+

方案说明 #

• toDF()方法传入列名数组时，会按照顺序和DataFrame的列位置一一对应重命名，不会受重复列名干扰，执行逻辑稳定。
• 上述实现自动识别重复列并追加序号后缀，无需提前感知重复列的数量和名称，可直接复用在所有存在重复列的场景。

内容的提问来源于stack exchange，提问作者Jeremy