如何按行索引均等拆分Spark DataFrame？规避randomSplit随机问题

嘿，这个场景我太熟悉了！要按原顺序均分Spark DataFrame，又不想用会打乱数据的randomSplit，也不想用take这种把全量数据拉到Driver的高消耗操作，其实有几个高效的分布式方案可以解决，我给你细细道来：

方案一：窗口函数+分桶（推荐，分布式高效） #

这个方法是最稳妥的分布式方案，不会把数据集中到Driver，性能拉满。核心思路是给每行数据分配一个行号，再根据行号计算所属的分组，最后按分组过滤拆分。

步骤说明： #

1. 计算总条数与分组大小：先通过count()获取DataFrame的总行数，再计算每个子DataFrame的目标行数。
2. 添加行号：用row_number()窗口函数生成连续行号（如果需要严格保留原顺序，一定要指定排序键；如果允许大致顺序，用monotonically_increasing_id()更高效）。
3. 分配分组ID：根据行号和分组大小，计算每行所属的分组。
4. 拆分DataFrame：遍历分组ID，过滤出每个分组的数据，最后删除辅助列即可。

代码示例（Scala）： #

import org.apache.spark.sql.functions._
import org.apache.spark.sql.expressions.Window

// 假设你的源DataFrame名为df
val totalRows = df.count()
val numSplits = 4 // 要拆分的份数
val splitSize = totalRows / numSplits

// 生成行号：这里用monotonically_increasing_id保证大致顺序，若需严格顺序，替换为你的业务有序列
val windowSpec = Window.orderBy(monotonically_increasing_id())
val dfWithRowNum = df.withColumn("row_num", row_number().over(windowSpec))

// 计算每行所属的分组ID
val dfWithGroup = dfWithRowNum.withColumn(
  "group_id",
  floor(($"row_num" - 1) / splitSize).cast("int")
)

// 拆分得到子DataFrames
val splitDfs = (0 until numSplits).map { groupId =>
  dfWithGroup.filter($"group_id" === groupId)
             .drop("row_num", "group_id") // 删除辅助列
}

注意事项： #

• 如果总行数无法被拆分份数整除，最后一个分组会包含剩余的所有行，你可以根据需求调整逻辑（比如把多余的行均匀分配到前面的分组）。
• 若必须严格保证原数据的顺序，一定要用业务中本身存在的有序列（比如时间戳、自增ID）来替代monotonically_increasing_id()，避免Spark随机排序。
• 绝对不要为了生成连续行号而用repartition(1)，这会把所有数据挤到一个分区，数据量大时直接OOM。

方案二：limit + exceptAll（适合小数据量或简单场景） #

如果你的数据量不大，或者追求代码简洁，可以用limit和exceptAll的组合来逐步拆分。这个方法不需要添加辅助列，但性能不如窗口函数，因为每次exceptAll都需要做数据比对。

代码示例（Scala）： #

import scala.collection.mutable.ArrayBuffer

val numSplits = 4
val totalRows = df.count()
val splitSize = totalRows / numSplits

var remainingDf = df
val splitDfs = ArrayBuffer[DataFrame]()

// 拆分前numSplits-1个分组
for (_ <- 0 until numSplits - 1) {
  val splitDf = remainingDf.limit(splitSize)
  splitDfs.append(splitDf)
  remainingDf = remainingDf.exceptAll(splitDf)
}
// 添加最后一个分组（包含剩余的所有行）
splitDfs.append(remainingDf)

注意事项： #

• 这个方法依赖Spark的稳定顺序，如果你的DataFrame没有指定排序，可能会出现顺序错乱的情况，所以建议先对DataFrame做排序操作。
• 数据量大时，exceptAll的性能开销会比较明显，优先推荐方案一。

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