Spark Scala(2.3.0)高效添加连续RowNumber的最优方案求助

我完全懂你的困扰——用monotonically_increasing_id()会莫名其妙出现超大跳号，而ROW_NUMBER()或者zipWithIndex又因为全局排序/洗牌的开销，处理243MB的数据集都慢到没法用。针对Spark 2.3.0，给你一个既高效又能生成连续无跳变行号的方案：

为什么你之前的方法不行？ #

先快速复盘下问题根源：

• monotonically_increasing_id()：这个函数是基于分区ID和分区内偏移生成的64位整数，并不是严格连续的。当数据分布到不同分区，或者后续分区数变化时，很容易出现大跨度的跳值，你在24万行后遇到的8589934592就是典型的分区ID导致的跳变。
• ROW_NUMBER() OVER (ORDER BY Year)：这个操作需要全局排序，Spark会触发大规模shuffle，把所有数据拉到一起排序，对于中等规模的数据集来说，shuffle的IO和计算开销极大，自然耗时过长。zipWithIndex本质上也需要全局对齐数据，同样存在这个问题。

最优解决方案：分区内编号+预计算分区偏移 #

核心思路是避免全局shuffle，通过预计算每个分区的起始行号偏移，再给分区内的每行分配局部编号，最终相加得到全局连续行号。具体步骤和代码如下：

import org.apache.spark.sql.functions._
import org.apache.spark.sql.types.LongType

// 1. 计算每个分区的行数，然后累加得到每个分区的起始行号偏移（比如第一个分区从0开始，第二个从第一个分区的行数开始）
val partitionOffsets = df.rdd.mapPartitionsWithIndex { case (partitionId, rows) =>
  Iterator((partitionId, rows.size))
}.collect()
// 按分区ID排序，然后用scanLeft计算累加偏移量
.sortBy(_._1)
.scanLeft((-1L, 0L)) { case ((_, prevOffset), (currentId, count)) =>
  (currentId, prevOffset + count)
}.tail // 去掉初始的(-1, 0)
.toMap // 转成分区ID到起始偏移的映射

// 2. 广播分区偏移量，让每个任务都能快速获取，避免重复计算
val broadcastedOffsets = spark.sparkContext.broadcast(partitionOffsets)

// 3. 给每个分区内的行添加局部编号，加上分区起始偏移得到全局行号
val dfWithRowNumber = df.rdd.mapPartitionsWithIndex { case (partitionId, rows) =>
  val startOffset = broadcastedOffsets.value(partitionId)
  // zipWithIndex给分区内每行分配0开始的局部编号，相加后得到全局行号（+1让行号从1开始，不需要可以去掉）
  rows.zipWithIndex.map { case (row, localIdx) =>
    row.copy(row.toSeq :+ (startOffset + localIdx + 1): _*)
  }
}.toDF(df.columns :+ "RowNumber": _*)

这个方案的优势 #

• 性能拉满：全程没有全局shuffle，只有计算分区偏移时的一次小数据量collect（分区数通常远小于数据行数），处理速度比ROW_NUMBER()快几个量级。
• 行号连续无跳变：完全按照数据的原始分区顺序生成连续递增的行号，不会出现异常大值。
• 适配Spark 2.3.0：所有API都是2.3.0支持的，不需要升级版本。

额外注意点 #

• 如果你的行号需要从0开始，去掉代码里的+1即可。
• 如果DataFrame已经按某个字段分区，也可以基于分区字段分组后再用类似逻辑生成组内行号，但如果是全局行号，上面的方案最通用。

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