如何根据spark-scala中的唯一列将数据帧保存到多个文件中

Question

I've the inputDf that I need to divide based on the columns origin and destination and save each unique combination into a different csv file.

(Using Spark 2.4.4)

val spark: SparkSession = SparkSession.builder().appName("Test").getOrCreate()

val inputRdd: RDD[(String, String, String, String, String, String)] = spark.sparkContext.parallelize(Seq(
  ("City1", "City2", "Sedan", "AE1235", "80", "2020-02-01"),
  ("City2", "City3", "Hatchback", "XY5434", "100", "2020-02-01"),
  ("City3", "City1", "Sedan", "YU3456", "120", "2020-02-01"),
  ("City3", "City2", "Sedan", "BV3555", "105", "2020-02-01"),
  ("City2", "City1", "SUV", "PO1234", "75", "2020-02-01"),
  ("City1", "City3", "SUV", "TY4123", "125", "2020-02-01"),
  ("City1", "City2", "Hatchback", "VI3415", "85", "2020-02-01"),
  ("City1", "City2", "SUV", "VF1244", "84", "2020-02-01"),
  ("City3", "City1", "Sedan", "EW1248", "124", "2020-02-01"),
  ("City2", "City1", "Hatchback", "GE576", "82", "2020-02-01"),
  ("City3", "City2", "Sedan", "PK2144", "104", "2020-02-01"),
  ("City3", "City1", "Hatchback", "PJ1244", "118", "2020-02-01"),
  ("City3", "City2", "SUV", "WF0976", "98", "2020-02-01"),
  ("City1", "City2", "Sedan", "WE876", "78", "2020-02-01"),
  ("City2", "City1", "Hatchback", "AB5467", "80", "2020-02-01")
))
val inputDf = spark.createDataFrame(inputRdd).toDF("origin", "destination", "vehicleType", "uniqueId", "distanceTravelled", "date")

Sample Output:

csv file 1:

origin,destination,vehicleType,uniqueId,distanceTravelled,date
City1,City2,Sedan,AE1235,80,2020-02-01
City1,City2,Hatchback,VI3415,85,2020-02-01
City1,City2,SUV,VF1244,84,2020-02-01
City1,City2,Sedan,WE876,78,2020-02-01

csv file 2:

origin,destination,vehicleType,uniqueId,distanceTravelled,date
City3,City1,Sedan,YU3456,120,2020-02-01
City3,City1,Sedan,EW1248,124,2020-02-01
City3,City1,Hatchback,PJ1244,118,2020-02-01

csv file 3:

origin,destination,vehicleType,uniqueId,distanceTravelled,date
City2,City1,SUV,PO1234,75,2020-02-01
City2,City1,Hatchback,GE576,82,2020-02-01
City2,City1,Hatchback,AB5467,80,2020-02-01

So far I've tried getting the unique combinations into a tuple and then using a foreach on it, filtering the inputDf each time saving the filtered dataframe to csv

val tuple = inputDf.groupBy("origin","destination").count()
  .select("origin","destination").rdd.map(r => (r(0),r(1))).collect

tuple.foreach(row => {
  val origin = row._1
  val destination = row._2
  val dataToWrite = inputDf.filter(inputDf.col("origin").equalTo(origin) && inputDf.col("destination").equalTo(destination))
  dataToWrite.repartition(1).write.mode("overwrite").format("csv").option("header", "true").save("/path/to/output/folder/" + origin + "-" + destination + ".csv")
})

This approach takes a lot of time as it involves filtering the inputDf every single time as the number of unique combinations are pretty huge. What would be an optimal way to do it?

EDIT: Each inputDf will have data only for one date.

The output should contain files at date level.

Like:

/output/City1-City2/2020-02-01.csv

/output/City1-City2/2020-02-02.csv

/output/City1-City2/2020-02-03.csv

/output/City3-City1/2020-02-01.csv

/output/City3-City1/2020-02-02.csv

... and so on

Answer 1

You can use partitionBy and divide data in separate csv file as per your combination. I have used coalesce to keep all data into one csv file which is not recommended if you have large data. GO through below code which will write all possible combination into separate csv files.

    scala> df.show()
+------+-----------+-----------+--------+-----------------+----------+
|origin|destination|vehicleType|uniqueId|distanceTravelled|      date|
+------+-----------+-----------+--------+-----------------+----------+
| City1|      City2|      Sedan|  AE1235|               80|2020-02-01|
| City2|      City3|  Hatchback|  XY5434|              100|2020-02-01|
| City3|      City1|      Sedan|  YU3456|              120|2020-02-01|
| City3|      City2|      Sedan|  BV3555|              105|2020-02-01|
| City2|      City1|        SUV|  PO1234|               75|2020-02-01|
| City1|      City3|        SUV|  TY4123|              125|2020-02-01|
| City1|      City2|  Hatchback|  VI3415|               85|2020-02-02|
| City1|      City2|        SUV|  VF1244|               84|2020-02-02|
| City3|      City1|      Sedan|  EW1248|              124|2020-02-02|
| City2|      City1|  Hatchback|   GE576|               82|2020-02-02|
| City3|      City2|      Sedan|  PK2144|              104|2020-02-02|
| City3|      City1|  Hatchback|  PJ1244|              118|2020-02-02|
| City3|      City2|        SUV|  WF0976|               98|2020-02-02|
| City1|      City2|      Sedan|   WE876|               78|2020-02-02|
| City2|      City1|  Hatchback|  AB5467|               80|2020-02-02|
+------+-----------+-----------+--------+-----------------+----------+


scala> val df1 = df.withColumn("combination", concat(col("origin") ,lit("-"), col("destination")))

scala> df1.coalesce(1).write.partitionBy("combination","date").format("csv").option("header", "true").mode("overwrite").save("/stackOut/")

Output will be like: