Distributed Data Scenarios: Handling Skew in Spark’s groupBy Operation

Performance issues in Spark sometimes trace back to the data structure itself, even when the code and the cluster are configured correctly. I discovered this firsthand when a seemingly simple groupBy("city") operation exploded into over 200 shuffle tasks. While most tasks completed in seconds, one dragged on for minutes. My job ran slowly and created a bottleneck for an entire pipeline. Behind an ordinary-looking transformation sat a fundamental system imbalance. In this article, I’ll share a real-world example of shuffle skew in distributed data processing and how I resolved it.

Problem Overview: Skewed Shuffle from a groupBy Operation

flowchart TD
    A["Input DataFrame
(billions of rows)"] --> B["groupBy(city)"]
    B --> C["Shuffle
200+ partitions created"]
    C --> D1["Partition
a few MB
finishes in seconds"]
    C --> D2["Partition
a few MB
finishes in seconds"]
    C --> D3["..."]
    C --> DN["One partition ~1.2 GB
millions of New York rows"]
    DN --> E["Longest task ~12 min
spill to disk, OOM risk
job stalls at ~25 min"]

During a large Spark job, a simple transformation using groupBy("city") led to significant performance issues:

• 200+ shuffle partitions were created by Spark to execute the groupBy operation.
• Due to data skew, one or two cities (e.g., “New York”) contained millions of records, while most others had only a few.
• This resulted in:
  • One task handling a massive partition (1+ GB),
  • Most tasks finishing quickly, and
  • One or two tasks taking much longer, ultimately stalling or failing the entire job.
• The Spark UI showed clear task time imbalance, frequent spill to disk, and occasionally executor memory errors (OOMs).
• The root cause lay in how groupBy and shuffle operations distribute tasks across a distributed system.

Solution: Mitigating Skew with Salting + Repartitioning

The core issue sat in how Spark distributed the data during the shuffle caused by groupBy("city"). High-cardinality cities like “New York” created data skew, overwhelming single tasks while others finished almost instantly.

We solved this using two complementary techniques:

flowchart TD
    A["Skewed DataFrame"] --> S1["Salting
salt = floor(rand() x 10)"]
    S1 --> S2["city_salted =
city plus salt"]
    S2 --> S3["Shuffle spreads the hot key
across 10 sub-partitions"]
    S3 --> S4["Aggregate per salted key"]
    S4 --> S5["Strip salt, re-aggregate
to correct per-city counts"]
    A --> R["Repartition
df.repartition(100, city)"]
    R --> S3

1. Salting the Skewed Key

Salting divides a hot key (like “New York”) into multiple artificial sub-keys, distributing the load more evenly.

from pyspark.sql.functions import concat, col, lit, floor, rand

# Step 1: Add a salt column
salted_df = df.withColumn("salt", floor(rand() * 10))  # 10 buckets

# Step 2: Salt the groupBy key
salted_df = salted_df.withColumn("city_salted", concat(col("city"), lit("_"), col("salt")))

# Step 3: Perform groupBy on salted key
grouped = salted_df.groupBy("city_salted").count()

# Step 4: Unsalt (combine buckets)
from pyspark.sql.functions import split
grouped_final = grouped.withColumn("city", split(col("city_salted"), "_")[0]) \
                       .groupBy("city").sum("count")

2. Repartition Before Shuffle

We also balanced the workload by repartitioning by the skewed key to control task distribution:

df = df.repartition(100, "city")

This ensures Spark pre-distributes the data, preventing hotspot partitions.

Before vs After: Performance Comparison