Scala:如何按行号拆分数据框?

时间:2019-03-14 10:54:03

标签: scala apache-spark split apache-spark-sql databricks

我想将270万行的数据帧拆分为100000行的小数据帧,所以最终得到27个数据帧,我也希望将它们存储为csv文件。

我已经看过这个partitionBy和groupBy,但是我不必担心任何条件,只需要按日期将它们排序即可。我正在尝试编写自己的代码来完成这项工作,但是如果您知道我可以使用的某些Scala(Spark)函数,那就太好了!

谢谢大家的建议!

1 个答案:

答案 0 :(得分:2)

您可以使用RDD API中的zipWithIndex(不幸的是,SparkSQL中没有等效项)将每行映射到一个索引,索引范围在0rdd.count - 1之间。

因此,如果您假设有一个数据框,我认为该数据帧将进行相应的排序,则需要在两个API之间来回切换,如下所示:

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

// creating mock data
val df = spark.range(100).withColumn("test", 'id % 10)

// zipping the data
val partitionSize = 5 // I use 5 but you can use 100000 in your case
val zipped_rdd = df.rdd
    .zipWithIndex.map{ case (row, id) => 
        Row.fromSeq(row.toSeq :+ id / partitionSize ) 
    }

//back to df
val newField = StructField("partition", LongType, false)
val zipped_df = spark
    .createDataFrame(zipped_rdd, df.schema.add(newField))

让我们看一下数据,我们有一个名为partition的新列,它对应于您想要拆分数据的方式。

zipped_df.show(15) // 5 rows by partition
+---+----+---------+
| id|test|partition|
+---+----+---------+
|  0|   0|        0|
|  1|   1|        0|
|  2|   2|        0|
|  3|   3|        0|
|  4|   4|        0|
|  5|   5|        1|
|  6|   6|        1|
|  7|   7|        1|
|  8|   8|        1|
|  9|   9|        1|
| 10|   0|        2|
| 11|   1|        2|
| 12|   2|        2|
| 13|   3|        2|
| 14|   4|        2|
+---+----+---------+

// using partitionBy to write the data
zipped_df.write
    .partitionBy("partition")
    .csv(".../testPart.csv")