Spark Groupby DataFrame

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

Syntax:

groupBy(String col1, String... cols): RelationalGroupedDataset

When we perform groupBy() on a Spark DataFrame, it returns a RelationalGroupedDataset object that allows us to carry out various 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.

Preparing Data & DataFrame

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

We will use this Spark 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.

import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.types.DataTypes;
import org.apache.spark.sql.RowFactory;
import java.util.Arrays;
import static org.apache.spark.sql.functions.col;

SparkSession spark = SparkSession.builder()
        .appName("SparkGroupByExample")
        .master("local")
        .getOrCreate();

List<Row> data = Arrays.asList(
        RowFactory.create("James", "Sales", "NY", 90000, 34, 10000),
        RowFactory.create("Michael", "Sales", "NY", 86000, 56, 20000),
        RowFactory.create("Robert", "Sales", "CA", 81000, 30, 23000),
        RowFactory.create("Maria", "Finance", "CA", 90000, 24, 23000),
        RowFactory.create("Raman", "Finance", "CA", 99000, 40, 24000),
        RowFactory.create("Scott", "Finance", "NY", 83000, 36, 19000),
        RowFactory.create("Jen", "Finance", "NY", 79000, 53, 15000),
        RowFactory.create("Jeff", "Marketing", "CA", 80000, 25, 18000),
        RowFactory.create("Kumar", "Marketing", "NY", 91000, 50, 21000)
);

Dataset<Row> df = spark.createDataFrame(data, DataTypes.createStructType(Arrays.asList(
        DataTypes.createStructField("employee_name", DataTypes.StringType, false),
        DataTypes.createStructField("department", DataTypes.StringType, false),
        DataTypes.createStructField("state", DataTypes.StringType, false),
        DataTypes.createStructField("salary", DataTypes.IntegerType, false),
        DataTypes.createStructField("age", DataTypes.IntegerType, false),
        DataTypes.createStructField("bonus", DataTypes.IntegerType, false)
)));

df.show();

+-------------+----------+-----+------+---+-----+
|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 sum() aggregate function.

df.groupBy("department").sum("salary").show(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(false);

Calculate the minimum salary of each department using min()

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

Calculate the maximum salary of each department using max()

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

Calculate the average salary of each department using avg()

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

Calculate the mean salary of each department using mean()

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

groupBy and aggregate on DataFrame multiple columns

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.

// GroupBy on multiple columns
df.groupBy("department", "state")
  .sum("salary", "bonus")
  .show(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.

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 "import org.apache.spark.sql.functions.*".

import static org.apache.spark.sql.functions.*;

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(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.

df.groupBy("department")
  .agg(
    sum("salary").alias("sum_salary"),
    avg("salary").alias("avg_salary"),
    sum("bonus").alias("sum_bonus"),
    max("bonus").alias("max_bonus")
  )
  .where(col("sum_bonus").geq(50000))
  .show(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     |
+-----------+----------+-----------------+---------+----------+

This example is also available in the GitHub project for reference.