pyspark group by

How to group records in PySpark Azure Databricks?

by

Are you looking to find how to use groupBy() on PySpark Dataframe into Azure Databricks cloud or maybe you are looking for a solution, to do aggregation by grouping identical records of a Dataframe in PySpark Databricks? If you are looking for any of these problem solutions, then you have landed on the correct page. I will also show you what and how to use PySpark to do calculation on records of a Dataframe using groupBy() & agg() function in Azure Databricks. I will explain it by taking a practical example. So don’t waste time let’s start step by step guide to understanding how to group and calculate records of numeric values using PySpark Dataframe.

In this blog, I will teach you the following with practical examples:

  • Single aggregation
  • Multiple aggregations
  • Multiple aggregations using multiple columns
  • Renaming aggregated columns

groupBy() method is used to collect records of Dataframe together based on column identical values specified in PySpark Azure Databricks dataframe.

Syntax: dataframe_name.groupBy(column_name)

What is the syntax of the groupBy() function in PySpark Azure Databricks?

The syntax is as follows:

dataframe_name.groupBy(*cols)
Parameter NameRequiredDescription
*cols (str, list or Column)YesIt represents the columns to be considered for grouping.
Table 1: groupBy() Method in PySpark Databricks Parameter list with Details

Apache Spark Official documentation link: groupBy()

Create a simple DataFrame

Gentle reminder:

In Databricks,

  • sparkSession made available as spark
  • sparkContext made available as sc

In case, you want to create it manually, use the below code.

from pyspark.sql.session import SparkSession

spark = SparkSession.builder 
    .master("local[*]") 
    .appName("azurelib.com") 
    .getOrCreate()

sc = spark.sparkContext

a) Create manual PySpark DataFrame

data = [
    (1, "Anandhi", "IT", 35000, "KL"),
    (2, "Manju", "IT", 50000, "KL"),
    (3, "Kannan", "HR", 34000, "DL"),
    (4, "Bharathi", "HR", 43000, "DL"),
    (5, "Kutty", "Sales", 10000, "KL"),
    (6, "Balaji", "Sales", 15000, "TN"),
    (7, "Elango", "Sales", 25000, "TN"),
]

df = spark.createDataFrame(data, schema=["id","name","department","salary","state"])
df.printSchema()
df.show(truncate=False)

"""
root
 |-- id: long (nullable = true)
 |-- name: string (nullable = true)
 |-- department: string (nullable = true)
 |-- salary: long (nullable = true)
 |-- state: string (nullable = true)

+---+--------+----------+------+-----+
|id |name    |department|salary|state|
+---+--------+----------+------+-----+
|1  |Anandhi |IT        |35000 |KL   |
|2  |Manju   |IT        |50000 |KL   |
|3  |Kannan  |HR        |34000 |DL   |
|4  |Bharathi|HR        |43000 |DL   |
|5  |Kutty   |Sales     |10000 |KL   |
|6  |Balaji  |Sales     |15000 |TN   |
|7  |Elango  |Sales     |25000 |TN   |
+---+--------+----------+------+-----+
"""

b) Creating a DataFrame by reading files

Download and use the below source file.

Download Sample Employees.csv
# replace the file_paths with the source file location which you have downloaded.

df_2 = spark.read.format("csv").option("header", True).load(file_path)
df_2.printSchema()

"""
root
 |-- id: long (nullable = true)
 |-- name: string (nullable = true)
 |-- department: string (nullable = true)
 |-- salary: long (nullable = true)
 |-- state: string (nullable = true)
"""

Note: Here, I will be using the manually created DataFrame.

How to use groupBy() to perform single aggregation in PySpark Azure Databricks?

The PySpark function groupBy() is used for collecting identical data into groups on DataFrame and performing a count on the grouped data.

Group by Examples:

In the below example, we are trying to get a number of employees working in each department.

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

"""
Output:

+----------+-----+
|department|count|
+----------+-----+
|        IT|    2|
|        HR|    2|
|     Sales|    3|
+----------+-----+

"""

How to use groupBy() to perform multiple aggregations in PySpark Azure Databricks?

Example:

In the below example, we are trying to get the maximum and minimum paid employees in each department.

from pyspark.sql.functions import max, min

df.groupBy("department").agg(
    max("salary"),
    min("salary")
).show()

"""
Output:

+----------+-----------+-----------+
|department|max(salary)|min(salary)|
+----------+-----------+-----------+
|        IT|      50000|      35000|
|        HR|      43000|      34000|
|     Sales|      25000|      10000|
+----------+-----------+-----------+

"""

How to use groupBy() to perform multiple aggregations by considering multiple columns in PySpark Azure Databricks?

Example:

In the below example, we are trying to get the maximum and minimum paid employees in each department by state wise.

from pyspark.sql.functions import max, avg

df.groupBy("department", "state").agg(
    avg("salary"),
    max("salary")
).show()

"""
Output:

+----------+-----+-----------+-----------+
|department|state|avg(salary)|max(salary)|
+----------+-----+-----------+-----------+
|        IT|   KL|    42500.0|      50000|
|        HR|   DL|    38500.0|      43000|
|     Sales|   KL|    10000.0|      10000|
|     Sales|   TN|    20000.0|      25000|
+----------+-----+-----------+-----------+

"""

How to change column aggregated names while using groupBy() to perform multiple aggregations in PySpark Azure Databricks?

Example:

In the below example, we are trying to get the maximum and minimum paid employees in each department and changing the column names on the go.

from pyspark.sql.functions import max, avg

df.groupBy("department").agg(
    avg("salary").alias("avg_salary"),
    max("salary").alias("max_salary")
).show()

"""
Output:

+----------+------------------+----------+
|department|        avg_salary|max_salary|
+----------+------------------+----------+
|        IT|           42500.0|     50000|
|        HR|           38500.0|     43000|
|     Sales|16666.666666666668|     25000|
+----------+------------------+----------+

"""

I have attached the complete code used in this blog in notebook format to this GitHub link. You can download and import this notebook in databricks, jupyter notebook, etc.

GitHub Link

When should you use the PySpark group() function in Azure Databricks?

You can use the groupBy() function to collect the identical data into groups on DataFrame and perform count, sum, avg, min, and max functions on the grouped data.

Real World Use Case Scenarios for PySpark DataFrame groupBy() function in Azure Databricks?

  • Assume that in the employee DataFrame, we want to calculate the number of total accounts per department. In this scenario, we need to group based on the DataFrame and then count the number of records. This could be a good example of the group by use case.

What are the alternatives for grouping data in PySpark Azure Databricks?

There are multiple alternatives for creating a DataFrame schema manually, which are as follows:

  • The aggregation function over the Window function helps in calculating column values based on groups.

Final Thoughts

In this article, we have learned about the PySpark groupBy() method to collect identical data of DataFrame based on single or multiple columns in Azure Databricks along with the examples explained clearly. I have also covered different scenarios with practical examples that could be possible. I hope the information that was provided helped in gaining knowledge.

Please share your comments and suggestions in the comment section below and I will try to answer all your queries as time permits.

Join Telegram Group
Download Free Azure Videos
Connect on LinkedIn

Arud Seka Berne S
Arud Seka Berne S

As a big data engineer, I design and build scalable data processing systems and integrate them with various data sources and databases. I have a strong background in Python and am proficient in big data technologies such as Hadoop, Hive, Spark, Databricks, and Azure. My interest lies in working with large datasets and deriving actionable insights to support informed business decisions.