Spark Scala - DataFrame GroupBy and Aggregate

Key Insights

The agg() function is your workhorse for complex aggregations—learn to combine multiple aggregate functions in a single pass over your data to avoid redundant shuffles.
Conditional aggregations using when() and filter() clauses eliminate the need for pre-filtering DataFrames, making your aggregation logic cleaner and more efficient.
GroupBy operations trigger expensive shuffle operations; strategic partitioning and caching can dramatically improve performance on large datasets.

Introduction to GroupBy Operations in Spark

GroupBy operations form the backbone of data analysis in Spark. When you’re working with distributed datasets spanning gigabytes or terabytes, understanding how to efficiently aggregate data becomes critical. The groupBy() method partitions your DataFrame into groups based on one or more columns, then applies aggregate functions to each group independently.

Unlike single-machine tools like pandas, Spark’s groupBy must coordinate across a cluster. This means data shuffles—moving records between nodes so all rows sharing the same key end up on the same partition. Understanding this distributed nature helps you write performant aggregation code.

Basic GroupBy Syntax and Simple Aggregations

Let’s start with a sales dataset to demonstrate fundamental groupBy operations:

import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.functions._

val spark = SparkSession.builder()
  .appName("GroupByExample")
  .master("local[*]")
  .getOrCreate()

import spark.implicits._

// Sample sales data
val sales = Seq(
  ("North", "Electronics", 1200.0, 5),
  ("North", "Electronics", 800.0, 3),
  ("North", "Clothing", 450.0, 10),
  ("South", "Electronics", 1500.0, 7),
  ("South", "Clothing", 300.0, 15),
  ("South", "Clothing", 600.0, 20)
).toDF("region", "category", "revenue", "quantity")

// Simple count
val regionCounts = sales.groupBy("region").count()
regionCounts.show()
// +------+-----+
// |region|count|
// +------+-----+
// | North|    3|
// | South|    3|
// +------+-----+

// Sum revenue by region
val regionRevenue = sales.groupBy("region").sum("revenue")
regionRevenue.show()

// Multiple columns in groupBy
val categoryByRegion = sales
  .groupBy("region", "category")
  .agg(
    sum("revenue").as("total_revenue"),
    avg("quantity").as("avg_quantity")
  )
categoryByRegion.show()

The basic aggregate methods—count(), sum(), avg(), min(), max()—work directly on grouped data. However, they only allow one aggregation at a time. For real-world analysis, you’ll almost always need agg().

Using the `agg()` Function for Multiple Aggregations

The agg() function accepts multiple aggregate expressions, letting you compute several metrics in a single pass:

val comprehensiveStats = sales
  .groupBy("region", "category")
  .agg(
    sum("revenue").as("total_revenue"),
    avg("revenue").as("avg_revenue"),
    count("*").as("transaction_count"),
    sum("quantity").as("total_units"),
    max("revenue").as("largest_sale")
  )

comprehensiveStats.show()
// +------+-----------+-------------+-----------+-----------------+-----------+------------+
// |region|   category|total_revenue|avg_revenue|transaction_count|total_units|largest_sale|
// +------+-----------+-------------+-----------+-----------------+-----------+------------+
// | North|Electronics|       2000.0|     1000.0|                2|          8|      1200.0|
// | North|   Clothing|        450.0|      450.0|                1|         10|       450.0|
// | South|Electronics|       1500.0|     1500.0|                1|          7|      1500.0|
// | South|   Clothing|        900.0|      450.0|                2|         35|}       600.0|
// +------+-----------+-------------+-----------+-----------------+-----------+------------+

You have two notation styles for referencing columns inside agg():

// Using col() function
sales.groupBy("region").agg(sum(col("revenue")))

// Using $ string interpolation (requires spark.implicits._)
sales.groupBy("region").agg(sum($"revenue"))

// Using column name string directly
sales.groupBy("region").agg(sum("revenue"))

I prefer the string notation for simple cases and col() when building dynamic queries. The $ notation looks clean but can cause issues in certain contexts like UDFs.

Built-in Aggregate Functions Deep Dive

Spark provides aggregate functions beyond basic arithmetic. Here are the ones you’ll use regularly:

val advancedAggregations = sales
  .groupBy("region")
  .agg(
    // Collect all categories into an array (with duplicates)
    collect_list("category").as("all_categories"),
    
    // Collect unique categories only
    collect_set("category").as("unique_categories"),
    
    // First and last values (non-deterministic without ordering!)
    first("category").as("first_category"),
    last("category").as("last_category"),
    
    // Count distinct categories
    countDistinct("category").as("category_count"),
    
    // Approximate distinct count (faster for large datasets)
    approx_count_distinct("category", 0.05).as("approx_category_count")
  )

advancedAggregations.show(false)

For statistical analysis on numerical columns:

val statisticalAggregations = sales
  .groupBy("region")
  .agg(
    stddev("revenue").as("revenue_stddev"),
    stddev_pop("revenue").as("revenue_stddev_population"),
    variance("revenue").as("revenue_variance"),
    skewness("revenue").as("revenue_skewness"),
    kurtosis("revenue").as("revenue_kurtosis"),
    percentile_approx($"revenue", lit(0.5)).as("median_revenue")
  )

statisticalAggregations.show()

A word of caution: collect_list() and collect_set() gather all values into a single array on one executor. With high-cardinality groups, this can cause out-of-memory errors. Use them judiciously.

Conditional Aggregations with `when()` and `filter()`

Real analysis often requires conditional logic. Rather than filtering your DataFrame multiple times, embed conditions directly in your aggregations:

val salesWithMargin = sales.withColumn("margin", 
  when($"category" === "Electronics", $"revenue" * 0.15)
    .otherwise($"revenue" * 0.30)
)

// Using when() inside aggregation
val conditionalAgg = salesWithMargin
  .groupBy("region")
  .agg(
    sum(when($"category" === "Electronics", $"revenue").otherwise(0))
      .as("electronics_revenue"),
    sum(when($"category" === "Clothing", $"revenue").otherwise(0))
      .as("clothing_revenue"),
    count(when($"revenue" > 500, true)).as("high_value_transactions")
  )

conditionalAgg.show()

Spark 3.0+ introduced the filter() clause on aggregate functions, which is cleaner:

val filterClauseAgg = sales
  .groupBy("region")
  .agg(
    sum("revenue").as("total_revenue"),
    sum("revenue").filter($"category" === "Electronics").as("electronics_revenue"),
    sum("revenue").filter($"category" === "Clothing").as("clothing_revenue"),
    count("*").filter($"revenue" > 500).as("high_value_count"),
    avg("revenue").filter($"quantity" >= 10).as("avg_bulk_order_revenue")
  )

filterClauseAgg.show()

The filter() approach is more readable and often performs better because Spark can optimize it more effectively.

Window Functions vs GroupBy

GroupBy collapses rows—you get one output row per group. Window functions compute aggregates while preserving every input row. Choose based on your output requirements:

import org.apache.spark.sql.expressions.Window

// GroupBy: One row per region
val groupByResult = sales
  .groupBy("region")
  .agg(sum("revenue").as("region_total"))

// Window: Every row preserved, with region total added
val windowSpec = Window.partitionBy("region")

val windowResult = sales
  .withColumn("region_total", sum("revenue").over(windowSpec))
  .withColumn("pct_of_region", $"revenue" / $"region_total" * 100)

windowResult.show()
// +------+-----------+-------+--------+------------+------------------+
// |region|   category|revenue|quantity|region_total|     pct_of_region|
// +------+-----------+-------+--------+------------+------------------+
// | North|Electronics| 1200.0|       5|      2450.0| 48.97959183673469|
// | North|Electronics|  800.0|       3|      2450.0|32.653061224489796|
// | North|   Clothing|  450.0|      10|      2450.0|18.367346938775512|
// ...

Use groupBy when you need summary statistics. Use window functions when you need row-level calculations that reference group-level aggregates—like percentages of totals, running sums, or rankings.

Performance Considerations and Best Practices

GroupBy operations trigger shuffles, which are expensive. Here’s how to minimize the pain:

Partition strategically before groupBy:

// If you're grouping by region repeatedly, partition by it first
val partitionedSales = sales.repartition($"region")

// Subsequent groupBy operations on region will be faster
val result1 = partitionedSales.groupBy("region").agg(sum("revenue"))
val result2 = partitionedSales.groupBy("region").agg(avg("quantity"))

Cache intermediate results when reusing:

val groupedData = sales
  .groupBy("region", "category")
  .agg(
    sum("revenue").as("total_revenue"),
    sum("quantity").as("total_quantity")
  )
  .cache()

// First action triggers computation and caching
groupedData.count()

// Subsequent operations use cached data
val highRevenue = groupedData.filter($"total_revenue" > 1000)
val lowQuantity = groupedData.filter($"total_quantity" < 10)

Reduce data before grouping:

// Bad: groupBy on full dataset, then filter
val inefficient = sales
  .groupBy("region", "category")
  .agg(sum("revenue").as("total"))
  .filter($"total" > 1000)

// Better: filter early when possible
val efficient = sales
  .filter($"revenue" > 100)  // Reduce rows before shuffle
  .groupBy("region", "category")
  .agg(sum("revenue").as("total"))
  .filter($"total" > 1000)

Handle skewed keys:

When one key has disproportionately more data than others, that partition becomes a bottleneck. Consider salting techniques or using Spark 3.0+’s adaptive query execution with spark.sql.adaptive.skewJoin.enabled.

// Enable AQE for automatic skew handling
spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "true")

GroupBy and aggregation form the foundation of analytical workloads in Spark. Master the agg() function, understand when to use conditional aggregations versus pre-filtering, and always consider the shuffle implications of your grouping strategy. These fundamentals will serve you well whether you’re processing megabytes or petabytes.