95% of My PySpark Job Finished in 4 Minutes. The Last Task Took 40. Here's Why.

I executed a PySpark job that joined a massive transactions table with a customer dimension table. On the surface, it seemed like it should seamlessly scale across the entire cluster. However, the job behaved quite differently in practice. The Spark User Interface (UI) revealed that while almost every task completed in a matter of minutes, a select few of them took over 40 minutes to finish.

During this time, their respective executors used up 80-90% of the CPU, while the remaining cluster resources sat mostly idle, eagerly waiting for those slow tasks to complete.

This peculiar behavior brought me to the concept of data skew. Data skew occurs when a specific key or a few keys hold an outsized share of the data. When Spark distributes work across partitions using a hashing technique, all rows associated with a particular key end up in the same partition. If one key contains millions of rows while the rest have only a few thousand, that single partition and the task processing it will inevitably perform substantially more work than all the other partitions combined.

As a result, the job's total execution time is dictated by the slowest task, not the average task, meaning a single skewed key can dictate the entire runtime, even if it represents a minuscule portion of the overall row count.

To confirm my suspicion, I first examined the Spark UI's stage view. The stage's task duration chart displayed almost uniform short tasks, with one or two tasks extending far beyond the rest. This visual pattern — a uniform distribution of short tasks juxtaposed with one or two extreme outliers — is a telltale sign of data skew. To substantiate my hypothesis, I ran a straightforward aggregation on the join key before initiating any further operations:

```python

from pyspark.sql import functions as F

df.groupBy("customer_id").count().orderBy(F.desc("count")).show(20)

```

The output of this query conclusively demonstrated that a handful of customer IDs boasted row counts that dwarfed the rest of the dataset. These high-volume accounts were responsible for a disproportionate share of all transactions, causing Spark's default hash partitioning to place every one of these rows into a single partition.

To rectify this issue, I employed a technique known as salting. The core idea behind salting is to append a random, unique "salt" value to the skewed join key, transforming a single hot key into several smaller synthetic keys. This process essentially fragments the hot key into multiple smaller keys, thereby distributing the associated data across multiple partitions.

In essence, instead of all rows pertaining to the hot customer ending up in one massive partition, they are spread across several smaller partitions, each handled by a different executor in parallel. This modification eliminates the bottleneck caused by the skewed key.

The implementation of salting involved the following steps:

1. Generate a specified number of salt buckets (e.g., 10).

2. Add a new "salt" column to the original DataFrame, assigning a random salt value to each row.

3. Create a new "join_key" by concatenating the original join key and the salt value.

4. Expand the smaller DataFrame by performing a cross join with the salt values DataFrame, which enables the creation of a joined DataFrame where the salted keys are evenly distributed across the partitions.

5. Join the expanded smaller DataFrame with the original large DataFrame using the newly created salted join key.

6. Finally, drop the unnecessary "salt" and "join_key" columns from the resulting DataFrame to finalize the process.

By employing this salting method, I effectively distributed the workload across multiple partitions instead of concentrating it in one single partition. Consequently, the slowest task in the job drastically reduced in execution time, thereby improving the overall job performance.

One noteworthy observation was that Spark does not proactively warn about data skew. Without actively monitoring task-level timings in the Spark UI, one might not even notice the job's uneven performance. Furthermore, while salting is an effective solution, it does come with a performance cost due to increased data shuffling and joining.

It is advisable to apply salting selectively to the specific hot keys causing the issue, rather than salting the entire dataset, as over-salting could potentially introduce more overhead than it resolves.

Before resorting to salting, it is worth exploring alternative strategies. Broadcast joins can be employed if one side of the join is small enough to fit comfortably in executor memory, typically under a threshold of around 500MB, depending on the cluster's settings. This approach circumvents the shuffle process entirely, thereby eliminating the skew issue. Broadcast joins should be considered first if the data size is feasible for broadcasting.

Additionally, data skew is not limited to join operations; it can also manifest in groupBy aggregations, window functions partitioned by a skewed column, and even repartitioning by a skewed key. The same detection and resolution techniques, namely row count analysis per key and the application of salting or broadcast joins, can be generalized to tackle skew in these scenarios as well.

In conclusion, identifying and addressing data skew is crucial for optimizing Spark job performance. By detecting skewed key distributions early on, employing salting or broadcast joins, and continuously monitoring Spark's task-level timings, one can mitigate the adverse effects of data skew and ensure efficient resource utilization within the cluster.