如何在Apache Spark生成row_number列时避免OutOfMemory错误

你的思路正确性判断 #

这个思路完全可行，核心逻辑是将大数据集拆分为多个有序小分片，规避单分区计算引发的OOM问题，同时通过分片偏移量累加保证全局行号连续。但关键前提是：分片必须按排序键column_x有序拆分，否则无法保证最终行号的连续性。

具体操作步骤 #

1. 按排序键有序分片并落地（可选） #

如果数据量极大，内存无法承载全量排序后的数据集，建议先按column_x范围分片并写入磁盘，确保每个分片内部数据有序，分片之间column_x范围不重叠：

// Scala示例，Python逻辑一致
val splitNum = 16 // 根据集群资源、数据量调整，比如设为集群核心数的2倍
val sortedDf = originalDf.orderBy("column_x")
// 按column_x范围分片，保证每个分片内数据有序
sortedDf.repartitionByRange(splitNum, "column_x")
  .write
  .mode("overwrite")
  .parquet("/tmp/sharded_ordered_data")

必须用repartitionByRange而非普通repartition，前者能保证分片的有序性和范围不重叠，是后续计算的基础。

2. 计算各分片的行号偏移量 #

读取分片数据后，先统计每个分片的记录数，再计算每个分片的起始行号偏移量（即前面所有分片的总记录数之和）：

// 读取分片数据
val shardedDf = spark.read.parquet("/tmp/sharded_ordered_data")
// 统计每个分区的记录数并按分区索引排序
val partitionCounts = shardedDf.rdd.mapPartitionsWithIndex { (idx, iter) =>
  Iterator((idx, iter.size.toLong))
}.collect().sortBy(_._1)
// 计算每个分区的起始偏移量
val offsetMap = partitionCounts.scanLeft((-1, 0L)) { case ((_, prevOffset), (idx, count)) =>
  (idx, prevOffset + count)
}.tail.toMap

3. 生成全局连续行号 #

通过mapPartitionsWithIndex为每个分片内的记录生成局部行号，再加上对应分片的起始偏移量，得到全局行号：

import org.apache.spark.sql.Row
import org.apache.spark.sql.types.LongType

val finalDf = shardedDf.rdd.mapPartitionsWithIndex { (idx, iter) =>
  val baseOffset = offsetMap(idx)
  iter.zipWithIndex.map { case (row, localIdx) =>
    // 将全局行号追加到原数据末尾，可根据需求调整字段位置
    Row.fromSeq(row.toSeq :+ (baseOffset + localIdx + 1)) // +1使行号从1开始，从0开始则去掉
  }
}.toDF(shardedDf.columns :+ "global_row_number")

4. 结果验证 #

执行查询验证行号连续性：

SELECT column_x, global_row_number 
FROM finalDf 
ORDER BY column_x 
LIMIT 100;

检查global_row_number是否按column_x的顺序连续递增，无断层或重复。

无磁盘落地的优化方案 #

若集群内存足够承载全量排序后的数据集，可跳过磁盘落地，直接在内存中完成分片和行号计算，减少IO开销：

val sortedRdd = originalDf.orderBy("column_x").rdd
// 统计各分区记录数
val partitionCounts = sortedRdd.mapPartitionsWithIndex((idx, iter) => Iterator((idx, iter.size.toLong))).collect().sortBy(_._1)
// 计算偏移量
val offsetMap = partitionCounts.scanLeft((-1, 0L)) { case ((_, prev), (idx, cnt)) => (idx, prev + cnt) }.tail.toMap
// 生成全局行号
val finalRdd = sortedRdd.mapPartitionsWithIndex((idx, iter) => {
  val baseOffset = offsetMap(idx)
  iter.zipWithIndex.map { case (row, localIdx) => Row.fromSeq(row.toSeq :+ (baseOffset + localIdx + 1)) }
})
// 转换为DataFrame
val finalDf = spark.createDataFrame(finalRdd, originalDf.schema.add("global_row_number", LongType))

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