Multiple counts in the same sql request in spark scala - scala

Question

I am a new developer in Spark Scala and I have a simple table like this :
City Dept employee_id employee_salary
NY FI 10 10000
NY FI null 20000
WDC IT 30 100000
LA IT 40 500
What I want to do is to calculate :
the number of employees by city & department (a simple count(*))
the number of employees with a non null id
the number of employees with a small (< 500) or medium (< 1000) or high salary (> 1000) => 3 additional counts !
So, as an output, I should have something like this :
City Dept total_emp total_emp total_emp_small total_emp_medium total_emp_high
NY FI 100 90 10 70 10
WDC IT 200 100
LA IT 10 10
The challenge here is that I want to optimize those counts. Because as you can see, we have here many counts to do. Without an optimisation the brute force for me is to do one count by request and generate a new DF result after each count and at the end I do a left joint based on the fixed column (city & dept) to add those new columns. But it will be too heavy since my table contains a lot of data.
I think that the method "window" can simplify this but I am not sure.
Can you help me this at least with just 2 cases (id != null and salary < 500).
Thank you in advance

Answer 1

scala> df.show
+----+----+-----------+---------------+
|City|Dept|employee_id|employee_salary|
+----+----+-----------+---------------+
| NY| FI| 10| 10000|
| NY| FI| null| 20000|
| WDC| IT| 30| 100000|
| LA| IT| 40| 500|
| LA| IT| 40| 600|
| LA| IT| null| 200|
+----+----+-----------+---------------+
scala> val df1 = df.withColumn("NonNullEmpID", when(col("employee_id").isNotNull, lit(1)).otherwise(lit(0)))
.withColumn("Salary_Cat",
when(col("employee_salary") < 500, lit("SmallSalary")).
when(col("employee_salary") < 1000, lit("MediumSalary")).
when(col("employee_salary") >= 1000, lit("HighSalary")))
scala> val SalaryMap = df1.groupBy("Salary_Cat").agg(count(lit(1)).alias("count")).collect.flatMap(x => Map(x(0).toString -> x.toSeq.slice(1, x.length).mkString)).toMap
//the number of employees by city & department (a simple count(*))
//the number of employees with a non null id
//the number of employees with a small (< 500) or medium (< 1000) or high salary (> 1000) => 3 additional counts !
scala> df1.groupBy("City", "Dept")
.agg(count(lit(1)).alias("#Emp_per_City_Dept"), sum(col("NonNullEmpID")).alias("NonNullEmpID"))
.withColumn("total_emp_small", lit(SalaryMap("SmallSalary")))
.withColumn("total_emp_medium", lit(SalaryMap("MediumSalary")))
.withColumn("total_emp_high", lit(SalaryMap("HighSalary")))
.show()
+----+----+------------------+------------+---------------+----------------+--------------+
|City|Dept|#Emp_per_City_Dept|NonNullEmpID|total_emp_small|total_emp_medium|total_emp_high|
+----+----+------------------+------------+---------------+----------------+--------------+
| NY| FI| 2| 1| 1| 2| 3|
| LA| IT| 3| 2| 1| 2| 3|
| WDC| IT| 1| 1| 1| 2| 3|
+----+----+------------------+------------+---------------+----------------+--------------+