Spark DataFrame/DataSet 分页或一次迭代 N 行块

Question

我需要为我的数据集（在 Spark scala 中）实现分页。

如果 Spark 数据集中有 100 条记录，那么我需要分成 20 个批次，每个批次 5 个元素。

请问如何将 Spark 数据集/数据框拆分为 N 行？

--NS

Answer 1

不确定是否有更好的方法，但您可以尝试：将数据帧转换为 rdd，使用zipWithIndex、过滤，然后再次转换为数据帧。

例如，假设您的数据框由以下给出

scala> val df=sc.parallelize(1 to 100).toDF("value")
df: org.apache.spark.sql.DataFrame = [value: int]

scala> df.show()
+-----+
|value|
+-----+
|    1|
|    2|
|    3|
|    4|
|    5|
|    6|
|    7|
|    8|
|    9|
|   10|
|   11|
|   12|
|   13|
|   14|
|   15|
|   16|
|   17|
|   18|
|   19|
|   20|
+-----+
only showing top 20 rows

转换为带有索引的rdd和zip如下：val dfRDD=df.rdd.zipWithIndex

scala> val dfRDD=df.rdd.zipWithIndex
dfRDD: org.apache.spark.rdd.RDD[(org.apache.spark.sql.Row, Long)] = ZippedWithIndexRDD[81] at zipWithIndex at <console>:69

对于包含第 1 行到第 5 行的第一个数据框，按如下方式进行过滤：

val firstDF=dfRDD.filter{case(datum,index)=>(0 to 4).contains(index)}.map(_._1)
scala> val firstDF=dfRDD.filter{case(datum,index)=>(0 to 4).contains(index)}.map(_._1)
firstDF: org.apache.spark.rdd.RDD[org.apache.spark.sql.Row] = MapPartitionsRDD[85] at map at <console>:71

最后，转换为dataframe如下：sqlContext.createDataFrame(firstDF,df.schema)

scala> sqlContext.createDataFrame(firstDF,df.schema).show()
+-----+
|value|
+-----+
|    1|
|    2|
|    3|
|    4|
|    5|
+-----+

您必须对其余行重复这些步骤(5 to 9)，(10 to 14)等等。

编辑：为了让事情更快一点，我定义了一个方法

def splitDF(range:scala.collection.immutable.Range.Inclusive):org.apache.spark.sql.DataFrame={
    val mySplitRDD=dfRDD.filter{case(datum,index)=>range.contains(index)}.map(_._1)
    val mySplitDF=sqlContext.createDataFrame(mySplitRDD,df.schema)
    mySplitDF
}

然后使用map来获取所有的split，例如

val dataframes=List((0 to 4), (5 to 9),(10 to 14)).map(i=>splitDF(i))

scala> val dataframes=List((0 to 4), (5 to 9),(10 to 14)).map(i=>splitDF(i))
dataframes: List[org.apache.spark.sql.DataFrame] = List([value: int], [value: int], [value: int])

数据框2：

scala> dataframes(1).show()
+-----+
|value|
+-----+
|    6|
|    7|
|    8|
|    9|
|   10|
+-----+

数据框1：

scala> dataframes(0).show()
+-----+
|value|
+-----+
|    1|
|    2|
|    3|
|    4|
|    5|
+-----+

数据框3：

scala> dataframes(2).show()
+-----+
|value|
+-----+
|   11|
|   12|
|   13|
|   14|
|   15|
+-----+