scala - Spark DataFrame/DataSet 分页或一次迭代 N 行的 block

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 到 9)、(10 到 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得到所有的splits，比如

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|
+-----+