Spark Optimizations in Depth
While the previous section covered debugging and common performance killers, this section covers specific optimization techniques like caching, broadcasting, and partition tuning.
1. Caching and Persisting
If you reuse a DataFrame multiple times, Spark will recompute it from the source every time an action is called. Prevent this by caching the DataFrame in memory.
# Cache in memory (default)
df.cache()
# Persist with a specific storage level
from pyspark import StorageLevel
df.persist(StorageLevel.MEMORY_AND_DISK)
# Unpersist when done
df.unpersist()2. Broadcast Joins
When joining a large table with a small table (like a lookup table), you can broadcast the small table to all executors, completely avoiding the shuffle phase.
from pyspark.sql.functions import broadcast
large_df = spark.read.parquet("hdfs://...")
small_df = spark.read.parquet("hdfs://...")
# Force a broadcast join
result = large_df.join(broadcast(small_df), "key")
Spark does this automatically if a table is smaller than spark.sql.autoBroadcastJoinThreshold (default 10MB).
3. Managing Partitions
Choosing the right number of partitions is crucial for parallelism and memory management.
| Method | When to use | Notes |
|---|---|---|
df.repartition(n) |
Increasing partitions or balancing skewed data | Triggers a full shuffle. Expensive but distributes data evenly. |
df.coalesce(n) |
Decreasing partitions, especially before writing to disk | Avoids a full shuffle by combining existing partitions. Cheaper than repartition. |
4. Salting for Data Skew
If you are joining on a key where one value is incredibly common (e.g., `user_id = null` or a default ID), one executor will get overwhelmed. Salting involves adding a random number to the key to distribute the large group across multiple executors. Note that with Spark 3+, Adaptive Query Execution (AQE) handles a lot of skew automatically if enabled.
5. Long-Running Jobs & Memory Management
When running streaming jobs or keeping a Spark cluster alive for a long period to serve multiple jobs, memory leaks can become a silent killer.
- Unpersist Data: Always
unpersist()cached DataFrames when you are done with them. Lingering cached data will slowly consume cluster memory until it crashes with an OOM (Out Of Memory) error. - Clear JVM Memory: In PySpark, variables kept in the Python driver can sometimes hold references to large JVM objects. Explicitly deleting large Python objects (
del df) and letting Python's garbage collector run can help the JVM clean up faster. - Spark Context Restart: If a cluster is meant to run for days or weeks (e.g., a notebook cluster), it is sometimes necessary to periodically restart the SparkContext or the cluster itself to get a completely fresh JVM heap.