← Back to Data Engineering
Data Engineering

Staff-Level Guide to Scaling Apache Spark: Tackling Data Skew and OOMs

As a Staff Data Engineer, one of the most common issues you’ll be paged for is a critical Spark job failing with an OutOfMemoryError (OOM) or grinding to a crawl due to data skew. Standard suggestions like “just add more executors” are expensive and often fail to fix the root cause.

In this post, we’ll dive deep into Spark memory architecture, diagnosing bottlenecks, and specific code-level patterns to make your Spark pipelines bulletproof.

1. Demystifying Spark Memory OOMs

OOM errors fall into two main categories:

  1. Driver OOM: Usually caused by calling .collect() on large datasets, excessive broadcast joins, or driver-side data aggregation.
  2. Executor OOM: Occurs when data partitions are too large to fit in executor memory, or user code uses objects that aren’t garbage-collected.

The Fix: Dynamic Partition Sizing

Aim for an ideal partition size of 100MB to 200MB in memory. If your partitions are too small, Spark suffers from task scheduling overhead. If they are too large, you risk executor OOMs.

Adjust this with:

# Adjust default shuffle partitions based on data size
spark.conf.set("spark.sql.shuffle.partitions", "2000")

2. Dealing with Data Skew

Data skew occurs when a join or aggregation key is unevenly distributed, forcing a single executor to process 90% of the data while the other executors sit idle.

Pattern 1: Salting the Join Key

If you are joining two datasets on a skewed key (e.g., null values or a generic organization_id), you can append a random “salt” integer to the key in the left table, and replicate the keys in the right table to match.

import pyspark.sql.functions as F

# Adding a salt column to the skewed dataframe
salted_df = df_left.withColumn("salt", F.concat(F.col("skewed_key"), F.lit("_"), F.randint(0, 9)))

# Replicating the right lookup dataframe to match the salts
lookup_df_salted = df_right.withColumn("salt_array", F.array([F.lit(f"{i}") for i in range(10)])) \
                           .withColumn("salt_elem", F.explode("salt_array")) \
                           .withColumn("salt", F.concat(F.col("skewed_key"), F.lit("_"), F.col("salt_elem")))

# Perform the join on the salted key
joined_df = salted_df.join(lookup_df_salted, on="salt", how="inner")

3. Adapting to Adaptive Query Execution (AQE)

Starting with Spark 3.0, Adaptive Query Execution (AQE) is enabled by default. AQE is your best friend for auto-tuning:

# Ensure AQE is enabled
spark.conf.set("spark.sql.adaptive.enabled", "true")

# Auto-coalesce small post-shuffle partitions
spark.conf.set("spark.sql.adaptive.coalescePartitions.enabled", "true")

# Automatically handle skewed joins dynamically
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "true")

AQE detects skew at runtime, automatically splits skewed partitions into smaller sub-partitions, and joins them using union operations.

4. Best Practices Summary

  • Broadcast Joins: Only broadcast tables < 10MB (default). If you are sure your cluster can handle it, raise the limit, but monitor the driver memory carefully.
  • Avoid .collect(): Use .take(n) or write output directly to storage (S3/Adls/GCS) in parquet format.
  • Cache Responsibly: Only cache when a DataFrame is reused multiple times in the same session. Always .unpersist() when you are done.

With these techniques, you’ll save cloud infrastructure costs and make your pipelines significantly faster.