PySpark Groupby DataFrame

Similar to SQL “GROUP BY” clause, PySpark groupBy() function is used to collect identical data into groups on DataFrame/Dataset and perform aggregate functions on the grouped data.

Syntax:

groupBy(*cols): pyspark.sql.group.GroupedData

When we perform groupBy() on a Spark DataFrame, it returns a GroupedData object that contains the following aggregate functions:

• count(): Returns the count of rows for each group.
• mean(): Returns the mean of values for each group.
• max(): Returns the maximum of values for each group.
• min(): Returns the minimum of values for each group.
• sum(): Returns the total for values for each group.
• avg(): Returns the average for values for each group.
• agg(): Allows calculation of multiple aggregates simultaneously.
• pivot(): This function is used to pivot the DataFrame, which is not covered in this lesson but in a separate article on Pivot & Unpivot DataFrame.

Preparing Data & DataFrame

Before we start, let's create the DataFrame from a sequence of data to work with. This DataFrame contains columns "employee_name", "department", "state", "salary", "age", and "bonus".

We will use this PySpark DataFrame to run groupBy() on the "department" column and calculate aggregates like minimum, maximum, average, and total salary for each group using min(), max(), and sum() aggregate functions respectively. Finally, we will also see how to group and aggregate on multiple columns.

from pyspark.sql import SparkSession

spark = SparkSession.builder \
    .appName('AIStudyLabs.com') \
    .getOrCreate()

simpleData = [
    ("James", "Sales", "NY", 90000, 34, 10000),
    ("Michael", "Sales", "NY", 86000, 56, 20000),
    ("Robert", "Sales", "CA", 81000, 30, 23000),
    ("Maria", "Finance", "CA", 90000, 24, 23000),
    ("Raman", "Finance", "CA", 99000, 40, 24000),
    ("Scott", "Finance", "NY", 83000, 36, 19000),
    ("Jen", "Finance", "NY", 79000, 53, 15000),
    ("Jeff", "Marketing", "CA", 80000, 25, 18000),
    ("Kumar", "Marketing", "NY", 91000, 50, 21000)
]

columns = ["employee_name", "department", "state", "salary", "age", "bonus"]
df = spark.createDataFrame(data=simpleData, schema=columns)
df.show(truncate=False)

+-------------+----------+-----+------+---+-----+
|employee_name|department|state|salary|age|bonus|
+-------------+----------+-----+------+---+-----+
|        James|     Sales|   NY| 90000| 34|10000|
|      Michael|     Sales|   NY| 86000| 56|20000|
|       Robert|     Sales|   CA| 81000| 30|23000|
|        Maria|   Finance|   CA| 90000| 24|23000|
|        Raman|   Finance|   CA| 99000| 40|24000|
|        Scott|   Finance|   NY| 83000| 36|19000|
|          Jen|   Finance|   NY| 79000| 53|15000|
|         Jeff| Marketing|   CA| 80000| 25|18000|
|        Kumar| Marketing|   NY| 91000| 50|21000|
+-------------+----------+-----+------+---+-----+

groupBy and aggregate on DataFrame columns

Let's do the groupBy() on the department column of the DataFrame and find the sum of salary for each department using the sum() aggregate function.

df.groupBy("department").sum("salary").show(truncate=False)

+----------+-----------+
|department|sum(salary)|
+----------+-----------+
|Sales     |257000     |
|Finance   |351000     |
|Marketing |171000     |
+----------+-----------+

Similarly, we can calculate the number of employees in each department using count()

df.groupBy("department").count().show()

Calculate the minimum salary of each department using min()

df.groupBy("department").min("salary").show()

Calculate the maximum salary of each department using max()

df.groupBy("department").max("salary").show()

Calculate the average salary of each department using avg()

df.groupBy("department").avg("salary").show()

Calculate the mean salary of each department using mean()

df.groupBy("department").mean("salary").show()

groupBy and aggregate on DataFrame multiple columns

Similarly, we can also run groupBy and aggregate on two or more DataFrame columns. The example below does a group by on department, state and applies sum() on salary and bonus columns.

df.groupBy("department", "state") \
  .sum("salary", "bonus") \
  .show(truncate=False)

+-----------+-----+------------+----------+
|department |state|sum(salary) |sum(bonus)|
+-----------+-----+------------+----------+
|Finance    |NY   |162000      |34000     |
|Marketing  |NY   |91000       |21000     |
|Sales      |CA   |81000       |23000     |
|Marketing  |CA   |80000       |18000     |
|Finance    |CA   |189000      |47000     |
|Sales      |NY   |176000      |30000     |
+-----------+-----+------------+----------+

Similarly, we can run group by and aggregate on two or more columns for other aggregate functions. Please refer to the source code for an example.

Running more aggregates at a time

Using the agg() aggregate function, we can calculate many aggregations at once using Spark SQL aggregate functions such as sum(), avg(), min(), max(), mean(). To use these, we should import:

from pyspark.sql.functions import sum, avg, min, max, mean

df.groupBy("department") \
  .agg(
    sum("salary").alias("sum_salary"),
    avg("salary").alias("avg_salary"),
    sum("bonus").alias("sum_bonus"),
    max("bonus").alias("max_bonus")
  ) \
  .show(truncate=False)

This example groups by the department column and calculates sum() and avg() of salary for each department and calculates sum() and max() of bonus for each department.

+-----------+----------+-----------------+---------+----------+
|department |sum_salary|avg_salary       |sum_bonus|max_bonus |
+-----------+----------+-----------------+---------+----------+
|Sales      |257000    |85666.66666666667|53000    |23000     |
|Finance    |351000    |87750.0          |81000    |24000     |
|Marketing  |171000    |85500.0          |39000    |21000     |
+-----------+----------+-----------------+---------+----------+

Using a filter on aggregate data

Similar to the SQL “HAVING” clause, in a Spark DataFrame, we can use either the where() or filter() functions to filter the rows of aggregated data.

from pyspark.sql.functions import col

df.groupBy("department") \
  .agg(
    sum("salary").alias("sum_salary"),
    avg("salary").alias("avg_salary"),
    sum("bonus").alias("sum_bonus"),
    max("bonus").alias("max_bonus")
  ) \
  .filter(col("sum_bonus") >= 50000) \
  .show(truncate=False)

This removes the sum of a bonus that is less than 50000 and yields the output below.

+-----------+----------+-----------------+---------+----------+
|department |sum_salary|avg_salary       |sum_bonus|max_bonus |
+-----------+----------+-----------------+---------+----------+
|Sales      |257000    |85666.66666666667|53000    |23000     |
|Finance    |351000    |87750.0          |81000    |24000     |
+-----------+----------+-----------------+---------+----------+