Spark Dataframe : Pivot with sorting - scala

Question

I am reading the following json file into Dataframe in spark:
{"id" : "a", "country" : "uk", "date" : "2016-01-01"}
{"id" : "b", "country" : "uk", "date" : "2016-01-02"}
{"id" : "c", "country" : "fr", "date" : "2016-02-01"}
{"id" : "d", "country" : "de", "date" : "2016-03-01"}
{"id" : "e", "country" : "tk", "date" : "2016-04-01"}
{"id" : "f", "country" : "be", "date" : "2016-05-01"}
{"id" : "g", "country" : "nl", "date" : "2016-06-01"}
{"id" : "h", "country" : "uk", "date" : "2016-06-01"}
I then apply groupBy on it and pivot it on date, here's the (pseudo) code:
val df = spark.read.json("file.json")
val dfWithFormattedDate = df.withColumn("date", date_format(col("date"), "yyyy-MM"))
dfWithFormattedDate.groupBy("country").pivot("date").agg(countDistinct("id").alias("count")).orderBy("country")
This gives me the Dataframe with country and pivoted dates (months) as columns. I would then like to order the results in descending order of total count. However, I don't have count as one of the columns and I can't apply pivot after applying count() on groupBy as it returns Dataset and not RelationalGroupedDataset. I have tried the following as well:
dfWithFormattedDate.groupBy("country").pivot("date").count()
This does not give me count column either. Is there any way I can gave both count and pivoted dates in resultant Dataset so that I can order by count desc?
Update
Here's the current output:
country|2016-01|2016-02|2016-03| ....
fr | null | 1 | null |
be | null | null | null |
uk | 2 | null | null |
Here's the expected output:
country|count|2016-01|2016-02|2016-03| ....
uk | 3 | 2 | null | null |
fr | 1 | null | 1 | null |
be | 1 | null | null | null |
As you can see, I need the count column in the result and order the rows in descending order of count. Ordering without explicitly having count column is fine as well.

Answer 1

If our starting point is this DataFrame :
import org.apache.spark.sql.functions.{date_format ,col, countDistinct}
val result = df.withColumn("date", date_format(col("date"), "yyyy-MM"))
.groupBy("country").pivot("date").agg(countDistinct("id").alias("count"))
.na.fill(0)
We then can simply calculate the rowsum for all the columns excluding the country column:
import org.apache.spark.sql.functions.desc
val test = result.withColumn("count",
result.columns.filter(_!="country")
.map(c => col(c))
.reduce((x, y) => x + y))
.orderBy(desc("count"))
test.show()
+-------+-------+-------+-------+-------+-------+-------+-----+
|country|2016-01|2016-02|2016-03|2016-04|2016-05|2016-06|count|
+-------+-------+-------+-------+-------+-------+-------+-----+
| uk| 2| 0| 0| 0| 0| 1| 3|
| be| 0| 0| 0| 0| 1| 0| 1|
| de| 0| 0| 1| 0| 0| 0| 1|
| tk| 0| 0| 0| 1| 0| 0| 1|
| nl| 0| 0| 0| 0| 0| 1| 1|
| fr| 0| 1| 0| 0| 0| 0| 1|
+-------+-------+-------+-------+-------+-------+-------+-----+