How to add total count of DataFrame to an already grouped DataFrame?

Here is our data. We have an employee DataFrame with 3 columns, name, project and cost_to_project. An employee can belong to multiple projects and for each project a cost_to_project is assigned.

val data = Seq(
      ("Ingestion", "Jerry", 1000), ("Ingestion", "Arya", 2000), ("Ingestion", "Emily", 3000),
      ("ML", "Riley", 9000), ("ML", "Patrick", 1000), ("ML", "Mickey", 8000),
      ("Analytics", "Donald", 1000), ("Ingestion", "John", 1000), ("Analytics", "Emily", 8000),
      ("Analytics", "Arya", 10000), ("BI", "Mickey", 12000), ("BI", "Martin", 5000))

import spark.sqlContext.implicits._

val df = data.toDF("Project", "Name", "Cost_To_Project")


scala> df.show()
+---------+-------+---------------+
|  Project|   Name|Cost_To_Project|
+---------+-------+---------------+
|Ingestion|  Jerry|           1000|
|Ingestion|   Arya|           2000|
|Ingestion|  Emily|           3000|
|       ML|  Riley|           9000|
|       ML|Patrick|           1000|
|       ML| Mickey|           8000|
|Analytics| Donald|           1000|
|Ingestion|   John|           1000|
|Analytics|  Emily|           8000|
|Analytics|   Arya|          10000|
|       BI| Mickey|          12000|
|       BI| Martin|           5000|
+---------+-------+---------------+

We want to group the dataset by Name and get a count to see the employee and the number of projects they are assigned to. In addition to that sub count, we also want to add a column with a total count like below.

+-------+------------------+-----------+
|   Name|number_of_projects|Total Count|
+-------+------------------+-----------+
| Mickey|                 2|         12|
| Martin|                 1|         12|
|  Jerry|                 1|         12|
|  Riley|                 1|         12|
| Donald|                 1|         12|
|   John|                 1|         12|
|Patrick|                 1|         12|
|  Emily|                 2|         12|
|   Arya|                 2|         12|
+-------+------------------+-----------+

Solution

It is pretty simple to achieve this. Simply add a column with counting the dataFrame and convert the value to a literal.

val groupBy = df.groupBy("Name").agg(count("*").alias("number_of_projects")).withColumn("Total Count", lit(df.count))
groupBy.show()

+-------+------------------+-----------+
|   Name|number_of_projects|Total Count|
+-------+------------------+-----------+
| Mickey|                 2|         12|
| Martin|                 1|         12|
|  Jerry|                 1|         12|
|  Riley|                 1|         12|
| Donald|                 1|         12|
|   John|                 1|         12|
|Patrick|                 1|         12|
|  Emily|                 2|         12|
|   Arya|                 2|         12|
+-------+------------------+-----------+