pysparkadvanced50 minPySpark

Data Skew Mitigation with Key Salting

Eliminate distributed join stragglers and executor OOMs by splitting heavily skewed keys across parallel partitions using random salting.

Start buildingOpen related lesson

The incident

Company
StreamMedia Prime
Engineer
Karan Mehta, Principal Distributed Systems Engineer
Trigger
Video playback event aggregation jobs stalled at 99% progress bar for 45 minutes because a celebrity premiere generated 15 million views on a single video ID.
Stakes
A single overloaded executor repeatedly ran out of memory, wasting 120 EC2 worker hours and stalling streaming dashboards.

On a Friday evening, Karan Mehta responded to a critical performance alert at StreamMedia Prime, an over-the-top video streaming platform serving millions of concurrent viewers. The hourly video playback metrics aggregation job was completely stalled. The Spark application progress bar reached 99 percent within 5 minutes, but remained stuck on the final task for over 45 minutes before the executor terminated with a java.lang.OutOfMemoryError heap dump. Platform engineers were unable to refresh executive dashboards or calculate real-time content licensing royalties.

Karan inspected the Spark Web UI Stage detail view. The cluster had allocated 100 executor tasks. Ninety-nine tasks finished processing in under 15 seconds, each handling approximately 10,000 playback records. However, task 47 was assigned a single partition containing over 15 million playback events. A viral celebrity series premiere had generated massive viewing spikes on a single content identifier: 'video_null' and celebrity title 'video_hot'. The extreme data concentration created an acute bottleneck that paralyzed the entire cluster.

Because standard Spark hash partitioning assigns all records with identical join or grouping keys to the exact same partition (hash(key) % num_partitions), all 15 million events for 'video_hot' were routed to a single executor node. The individual executor was overwhelmed by memory pressure, thrashing garbage collection cycles for 45 minutes while the remaining 99 cluster nodes sat idle, wasting 120 EC2 worker hours. Simply throwing larger cloud hardware at the problem failed because vertical scaling could not resolve the mathematical skew.

Karan resolved to eliminate the partition skew using the distributed key salting technique. Key salting redistributes heavily skewed keys across multiple parallel tasks by appending a random integer suffix (the salt) between 0 and N minus 1 to the skewed key. In a two-stage aggregation architecture, Spark first aggregates records by the salted key in parallel across executors, and then performs a lightweight final aggregation on the unsalted key, balancing cluster workloads evenly.

Karan implemented the salting pipeline in PySpark. For each record, the script computed a salt integer between 0 and 3 using F.floor(F.rand() * 4) and concatenated it to the video_id. The 15 million records were evenly distributed across 4 parallel partitions. Stage 1 computed partial counts per salted key, and Stage 2 aggregated the partial totals back into the final true video count of 8 events and 1,475 total watch seconds, keeping the aggregation on the executors.

Karan deployed the salting logic to production. The hourly aggregation job completed in 3 minutes instead of 45 minutes, with task duration variance dropping from 45 minutes to under 8 seconds across all executors. Karan documented the key salting pattern as the enterprise standard for high-cardinality skew mitigation across StreamMedia Prime distributed streaming pipelines, establishing an architectural template for handling viral viewership surges.

What you will learn

  • Identify data skew bottlenecks in Spark UI task metrics and execution distributions
  • Generate random salt suffixes using F.floor(F.rand() * salt_factor)
  • Distribute skewed keys evenly across parallel partitions with salted keys
  • Execute two-stage distributed aggregations to compute exact totals without stragglers

What you need first

Data Skew & Task StragglersData skew & salting techniques

Data skew occurs when one key contains vastly more records than others, causing a single executor to bottleneck the entire stage.

PythonMini example
from pyspark.sql import SparkSession
spark = SparkSession.builder.master('local[1]').getOrCreate()
df = spark.createDataFrame([('hot',)] * 10 + [('cold',)], ['k'])
print(df.filter("k = 'hot'").count())
10
Watch out

Increasing cluster memory to solve skew instead of addressing key distribution, which wastes cloud budget.

Key Salting ArchitecturePartitioning & repartition strategies

Key salting appends random integers (0 to N-1) to disperse a hot key across N parallel partitions during grouping.

PythonMini example
from pyspark.sql import SparkSession
from pyspark.sql import functions as F
spark = SparkSession.builder.master('local[1]').getOrCreate()
df = spark.createDataFrame([('A',)], ['k'])
print(df.withColumn('salt', F.lit(1)).select(F.concat_ws('_', 'k', 'salt')).first()[0])
A_1
Watch out

Forgetting to perform the second-stage aggregation to combine salted sub-totals back into the true key total.

Architecture and data flow

Severe Partition Skew Bottleneck
Severe Partition Skew Bottleneck. A single viral video key routes millions of records to one executor task, stalling the entire stage.
Two-Stage Salted Aggregation DAG
Two-Stage Salted Aggregation DAG. Random salt integers disperse hot keys across parallel tasks before stage-2 rolls up exact totals.
Balanced Final Playback Mart
Balanced Final Playback Mart. The final mart reports exact total plays and watch duration without task stragglers.

The data

A playback event dataset where a viral video key accounts for 80 percent of all rows.

video_plays: Streaming video playback start events with severe key skew.

play_idvideo_idwatch_seconds
p1video_hot120
p2video_hot95
p3video_hot240
p4video_hot180
p5video_hot300
p6video_hot150
p7video_hot80
p8video_hot210
p9video_cold60
p10video_norm180

Step by step

Read each step, run its code, then check your output against the one shown.

  1. Initialize skewed playback DataFrame and inspect key distribution

    We initialize the playback events DataFrame containing 10 events: 8 events belong to a single viral video ('video_hot'), while only 1 event belongs to 'video_cold' and 1 to 'video_norm'. We inspect the record distribution across video IDs. In production video streaming platforms, viral live premieres create acute 80-20 Pareto skew where a single content identifier concentrates millions of events, turning standard distributed hash partitioning into a single-executor bottleneck. We display the key frequencies to observe the severe skew before applying salting.

    PythonStep 1
    from pyspark.sql import SparkSession
from pyspark.sql import functions as F

spark = SparkSession.builder.master("local[1]").appName("skew_salting").getOrCreate()
plays = [
    ("p1", "video_hot", 120),
    ("p2", "video_hot", 95),
    ("p3", "video_hot", 240),
    ("p4", "video_hot", 180),
    ("p5", "video_hot", 300),
    ("p6", "video_hot", 150),
    ("p7", "video_hot", 80),
    ("p8", "video_hot", 210),
    ("p9", "video_cold", 60),
    ("p10", "video_norm", 180),
]
df = spark.createDataFrame(plays, ["play_id", "video_id", "watch_seconds"])
df.groupBy("video_id").count().orderBy(F.col("count").desc()).show()
    +----------+-----+
|  video_id|count|
+----------+-----+
| video_hot|    8|
|video_norm|    1|
|video_cold|    1|
+----------+-----+
    Common mistake

    Failing to check key cardinality and distribution before designing distributed group-by aggregations.

    How you will know

    You will see video_hot with count 8 and cold/norm videos with count 1.

  2. Demonstrate partition straggler with naive un-salted aggregation

    Wrong turn first: this step shows the mistake before the fix

    Here is the naive approach developers often attempt first: running a standard groupBy('video_id').count(). In a distributed cluster, Spark hashes 'video_hot' to a single executor partition. That single partition must compute all 8 records while other partition tasks finish early. At production scale of 15 million rows, this naive aggregation forces one executor to process millions of records while the remaining 99 cluster nodes sit idle, producing the classic 99% progress bar stall. We execute this un-salted grouping to highlight the root cause of task stragglers.

    PythonStep 2
    from pyspark.sql import SparkSession
from pyspark.sql import functions as F

spark = SparkSession.builder.master("local[1]").appName("skew_salting").getOrCreate()
plays = [
    ("p1", "video_hot", 120),
    ("p2", "video_hot", 95),
    ("p3", "video_hot", 240),
    ("p4", "video_hot", 180),
    ("p5", "video_hot", 300),
    ("p6", "video_hot", 150),
    ("p7", "video_hot", 80),
    ("p8", "video_hot", 210),
    ("p9", "video_cold", 60),
    ("p10", "video_norm", 180),
]
df = spark.createDataFrame(plays, ["play_id", "video_id", "watch_seconds"])

# Naive un-salted grouping routes all hot keys to one executor task
naive = df.groupBy("video_id").agg(F.count("*").alias("plays"), F.sum("watch_seconds").alias("total_sec"))
naive.show()
    +----------+-----+---------+
|  video_id|plays|total_sec|
+----------+-----+---------+
|video_norm|    1|      180|
| video_hot|    8|     1375|
|video_cold|    1|       60|
+----------+-----+---------+
    Common mistake

    Relying on standard groupBy on skewed keys, forcing one executor to process millions of rows alone.

    How you will know

    Observe that all 8 hot records are aggregated in a single un-isolated key group.

  3. Append random salt integers to disperse hot key across partitions

    We mitigate key skew by salting the video_id column. We define a salt range of 4 (integers 0, 1, 2, 3) using F.floor(F.rand(seed=42) * 4). We concatenate the original video_id with the salt integer using F.concat_ws('_', 'video_id', 'salt') to generate a composite salted key: 'video_id_salted'. This randomly disperses the 8 records of 'video_hot' across four distinct hash values: 'video_hot_0', 'video_hot_1', 'video_hot_2', and 'video_hot_3'. Salting distributes hot keys across separate partitions, unlocking parallel execution across cluster workers without creating executor bottlenecks.

    PythonStep 3
    from pyspark.sql import SparkSession
from pyspark.sql import functions as F

spark = SparkSession.builder.master("local[1]").appName("skew_salting").getOrCreate()
plays = [
    ("p1", "video_hot", 120),
    ("p2", "video_hot", 95),
    ("p3", "video_hot", 240),
    ("p4", "video_hot", 180),
    ("p5", "video_hot", 300),
    ("p6", "video_hot", 150),
    ("p7", "video_hot", 80),
    ("p8", "video_hot", 210),
    ("p9", "video_cold", 60),
    ("p10", "video_norm", 180),
]
df = spark.createDataFrame(plays, ["play_id", "video_id", "watch_seconds"])

salt_factor = 4
salted_df = (
    df
    .withColumn("salt", F.floor(F.rand(seed=42) * salt_factor))
    .withColumn("video_id_salted", F.concat_ws("_", F.col("video_id"), F.col("salt")))
)
salted_df.select("play_id", "video_id", "salt", "video_id_salted").show(5, truncate=False)
    +-------+---------+----+---------------+
|play_id|video_id |salt|video_id_salted|
+-------+---------+----+---------------+
|p1     |video_hot|2   |video_hot_2    |
|p2     |video_hot|2   |video_hot_2    |
|p3     |video_hot|3   |video_hot_3    |
|p4     |video_hot|1   |video_hot_1    |
|p5     |video_hot|2   |video_hot_2    |
+-------+---------+----+---------------+
only showing top 5 rows
    Common mistake

    Using too large a salt factor on small clusters, which creates excessive partition fragmentation.

    How you will know

    Output shows salted keys like video_hot_0, video_hot_1, dispersing the single hot key.

  4. Execute Stage-1 partial aggregation over salted keys

    We execute Stage 1 of the two-stage aggregation by grouping by the composite 'video_id_salted' key. In a distributed cluster, each salted key hashes to a different executor partition. The 8 hot records are now processed in parallel across 4 tasks rather than bottlenecking a single task. We compute the partial count and partial sum of watch seconds per salted key, distributing memory pressure and computational load evenly across the cluster. Stage 1 collapses millions of raw records into a compact set of pre-aggregated sub-totals.

    PythonStep 4
    from pyspark.sql import SparkSession
from pyspark.sql import functions as F

spark = SparkSession.builder.master("local[1]").appName("skew_salting").getOrCreate()
plays = [
    ("p1", "video_hot", 120),
    ("p2", "video_hot", 95),
    ("p3", "video_hot", 240),
    ("p4", "video_hot", 180),
    ("p5", "video_hot", 300),
    ("p6", "video_hot", 150),
    ("p7", "video_hot", 80),
    ("p8", "video_hot", 210),
    ("p9", "video_cold", 60),
    ("p10", "video_norm", 180),
]
df = spark.createDataFrame(plays, ["play_id", "video_id", "watch_seconds"])
salt_factor = 4
salted_df = (
    df
    .withColumn("salt", F.floor(F.rand(seed=42) * salt_factor))
    .withColumn("video_id_salted", F.concat_ws("_", F.col("video_id"), F.col("salt")))
)

stage1 = (
    salted_df
    .groupBy("video_id_salted", "video_id")
    .agg(
        F.count("*").alias("sub_count"),
        F.sum("watch_seconds").alias("sub_seconds")
    )
)
stage1.show(truncate=False)
    +---------------+----------+---------+-----------+
|video_id_salted|video_id  |sub_count|sub_seconds|
+---------------+----------+---------+-----------+
|video_hot_1    |video_hot |1        |180        |
|video_norm_3   |video_norm|1        |180        |
|video_hot_2    |video_hot |4        |665        |
|video_cold_3   |video_cold|1        |60         |
|video_hot_0    |video_hot |1        |210        |
|video_hot_3    |video_hot |2        |320        |
+---------------+----------+---------+-----------+
    Common mistake

    Dropping the original unsalted video_id during stage 1, preventing stage 2 re-aggregation.

    How you will know

    Output shows sub_count and sub_seconds distributed across multiple salted sub-keys.

  5. Execute Stage-2 global aggregation to compute exact final totals

    In the final step, we execute Stage 2 global aggregation. We group by the original 'video_id' column, sum the sub_counts to get the true total count, and sum the sub_seconds to get total watch time. Because Stage 1 collapsed millions of rows into just 4 sub-total rows per salted key, Stage 2 processes only a handful of records and completes in milliseconds. The final output verifies exact totals: 'video_hot' has 8 plays and 1,475 total seconds without any task stragglers or executor memory exhaustion.

    PythonStep 5
    from pyspark.sql import SparkSession
from pyspark.sql import functions as F

spark = SparkSession.builder.master("local[1]").appName("skew_salting").getOrCreate()
plays = [
    ("p1", "video_hot", 120),
    ("p2", "video_hot", 95),
    ("p3", "video_hot", 240),
    ("p4", "video_hot", 180),
    ("p5", "video_hot", 300),
    ("p6", "video_hot", 150),
    ("p7", "video_hot", 80),
    ("p8", "video_hot", 210),
    ("p9", "video_cold", 60),
    ("p10", "video_norm", 180),
]
df = spark.createDataFrame(plays, ["play_id", "video_id", "watch_seconds"])
salt_factor = 4
salted_df = (
    df
    .withColumn("salt", F.floor(F.rand(seed=42) * salt_factor))
    .withColumn("video_id_salted", F.concat_ws("_", F.col("video_id"), F.col("salt")))
)

stage1 = (
    salted_df
    .groupBy("video_id_salted", "video_id")
    .agg(
        F.count("*").alias("sub_count"),
        F.sum("watch_seconds").alias("sub_seconds")
    )
)

final_mart = (
    stage1
    .groupBy("video_id")
    .agg(
        F.sum("sub_count").alias("total_plays"),
        F.sum("sub_seconds").alias("total_watch_seconds")
    )
    .orderBy(F.col("total_plays").desc())
)
final_mart.show(truncate=False)
    +----------+-----------+-------------------+
|video_id  |total_plays|total_watch_seconds|
+----------+-----------+-------------------+
|video_hot |8          |1375               |
|video_norm|1          |180                |
|video_cold|1          |60                 |
+----------+-----------+-------------------+
    Common mistake

    Failing to sum partial counts during stage 2, mistakenly counting sub-groups instead of summing sub_counts.

    How you will know

    Output prints final mart showing video_hot with total_plays 8 and total_watch_seconds 1375.

Build it yourself

Write your solution below, run it in your browser, and pass every check.

PySpark (Catalyst Engine)Ready to run
Loading editor…
Press Run Code to evaluate
Script executes cleanly without PySpark syntax or analysis errors
Uses F.rand to generate random salt values for skewed keys
Correctly aggregates skewed video_hot records
Evaluates exact total watch duration of 1375 seconds for video_hot

What the result means

Interpretation

Key salting eliminates distributed stragglers caused by severe key skew. By appending random salt integers, the 15 million records of a viral video are dispersed across multiple parallel executor partitions for stage-1 pre-aggregation.

The message you would send

Streaming metrics aggregation optimization complete: we eliminated partition straggler bottlenecks using distributed two-stage key salting. The hourly playback mart now finishes in 3 minutes instead of 45 minutes, eliminating executor OOM crashes and saving 120 EC2 worker hours daily.

Put it on your CV

Resume or LinkedIn bullet

OutputCopy this line
Engineered distributed two-stage key salting pipelines in PySpark eliminating task stragglers and cutting execution time by 90%.

README for GitHub

OutputREADME.md
# Data Skew Mitigation with Key Salting in PySpark

## Problem
A viral video premiere concentrated 15 million events on a single key, stalling 99% of tasks and crashing executor nodes.

## Solution Architecture
- Appended random salt integers (0..3) using F.floor(F.rand() * 4).
- Stage 1: Dispersed hot records across 4 parallel partition tasks for pre-aggregation.
- Stage 2: Aggregated sub-totals back to original video keys in milliseconds.

## Results
Reduced stage turnaround from 45 minutes to 3 minutes and saved 120 EC2 worker hours daily.
Stretch goal

Implement dynamic salting that applies salt factors only to keys identified as skewed by empirical sampling.

Keep going