Scala - 根据行数拆分数据帧答案

【问题标题】：Scala - splitting dataframe based on number of rowsScala - 根据行数拆分数据帧
【发布时间】：2022-10-02 19:28:24
【问题描述】：

我有一个包含大约一百万条记录的 spark 数据框。我试图将此数据帧拆分为多个小数据帧，其中每个数据帧的最大行数为 20,000（这些数据帧中的每一个的行数都应为 20,000，但最后一个数据帧可能有也可能没有 20,000）。你能帮我解决这个问题吗？谢谢你。

您可以粘贴到目前为止使用过的任何代码吗？数据框是否有一列是每一行的唯一 ID？
自从我对 Spark 做了很多事情以来已经有一段时间了，但这听起来像是您可以通过使用 coalesce 重新分区到“正确”数量的分区并按分区拆分来完成的事情。

标签： scala apache-spark-sql

【解决方案1】：

好的，也许不是最有效的方法，但就是这样。您可以创建一个计算每一行的新列（如果您没有唯一的 Id 列）。在这里，我们基本上是遍历整个数据帧并选择大小为 20k 的批次，将它们添加到数据帧列表中。

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

var index = 0
val subsetSize = 20000
var listOfDF: List[DataFrame] = List()

// withColumn optional if you already have a unique id per row
val df = spark.table("your_table").withColumn("rowNum", monotonically_increasing_id())

def returnSubDF(fromIndex: Int, toIndex: Int) = {
  df.filter($"rowNum" >= fromIndex && $"rowNum" < toIndex)
}

while (index <= 1000000){
  listOfDF = listOfDF :+ returnSubDF(index, index+subsetSize)
  index += subsetSize
}

listOfDF.head.show()

【讨论】：