Aug 25, 2026

PySpark Memory Management: How to Fix OutOfMemory Errors in Production Jobs

SEO Meta Description: Learn how to fix PySpark OutOfMemory (OOM) errors in production. Master executor memory architecture, data skew salting, PyArrow tuning, and Kubernetes OOMKilled fixes.

PySpark Memory Management: How to Fix OutOfMemory Errors in Production Jobs

It is 2:00 AM. Your phone triggers a PagerDuty alert: Data Pipeline Execution Failed. You open your monitoring dashboard and see the dreaded signature of a failed PySpark job:

java.lang.OutOfMemoryError: Java heap space or a silent Command exited with code 137.

PySpark memory management is one of the most misunderstood areas of cloud data engineering. Engineers often respond to OutOfMemory (OOM) failures by simply doubling the executor memory or blindly increasing instance sizes. While this "vertical scaling patch" might temporarily green the pipeline status, it spikes cloud infrastructure costs and fails when data volume grows by another 20%.

To build resilient, cost-effective data pipelines, you must understand how PySpark manages memory across both the Java Virtual Machine (JVM) and the Python worker process, how data frames flow through execution contexts, and how to systematically diagnose memory leaks.

This guide breaks down PySpark memory architecture, explores root causes of OOM errors, and provides actionable code patterns, configuration parameters, and infrastructure adjustments to stabilize your production jobs.

Table of Contents

  1. Understanding the PySpark Memory Model
  2. Diagnosing OOM Errors: Driver vs. Executor
  3. Root Causes of PySpark OOM Errors and Fixes
  4. Essential PySpark Memory Configurations
  5. Kubernetes (Spark-on-k8s) & Infrastructure Debugging
  6. Common Production Mistakes
  7. Production Best Practices Checklist
  8. Interview Questions & Answers
  9. Frequently Asked Questions (FAQ)
  10. Conclusion

1. Understanding the PySpark Memory Model

PySpark runs on a dual-runtime model. A PySpark application consists of a Driver process and one or more Executor processes. Because Spark core is written in Scala/Java, execution occurs within a Java Virtual Machine (JVM). However, PySpark code runs inside a Python worker process.

Understanding how memory is partitioned within a PySpark Executor container is critical for preventing OOM exceptions.

+-----------------------------------------------------------------------------------+
| Kubernetes Pod / YARN Node Manager Container Boundary                             |
|                                                                                   |
|  +-----------------------------------------------------------------------------+  |
|  | JVM Heap Memory (spark.executor.memory)                                     |  |
|  |                                                                             |  |
|  |  +-----------------------------------------------------------------------+  |  |
|  |  | Spark Unified Memory (spark.memory.fraction = 0.6)                     |  |  |
|  |  |                                                                       |  |  |
|  |  |  +--------------------------------+--------------------------------+  |  |  |
|  |  |  | Execution Memory               | Storage Memory                 |  |  |  |
|  |  |  | (Joins, Aggregations, Shuffles)| (Cached DataFrames, Broadcasts)|  |  |  |
|  |  |  +--------------------------------+--------------------------------+  |  |  |
|  |  +-----------------------------------------------------------------------+  |  |
|  |  | User Memory (0.4 * (Heap - Reserved))                                  |  |  |
|  |  | (Custom Data Structures, Internal Spark Metadata)                     |  |  |
|  |  +-----------------------------------------------------------------------+  |  |
|  |  | Reserved Memory (300 MB fixed)                                        |  |  |
|  |  +-----------------------------------------------------------------------+  |  |
|  +-----------------------------------------------------------------------------+  |
|                                                                                   |
|  +-----------------------------------------------------------------------------+  |
|  | PySpark Python Worker Process (spark.python.worker.memory)                  |  |
|  | (PyArrow buffers, Python objects, C-extensions like NumPy/Pandas)            |  |
|  +-----------------------------------------------------------------------------+  |
|                                                                                   |
|  +-----------------------------------------------------------------------------+  |
|  | Overhead & Off-Heap Memory (spark.executor.memoryOverhead)                 |  |
|  | (Netty frame buffers, Metaspace, Off-Heap Execution Memory)                 |  |
|  +-----------------------------------------------------------------------------+  |
+-----------------------------------------------------------------------------------+

Breakdown of Executor Components

  1. JVM Heap Memory (spark.executor.memory):
  2. Reserved Memory (300MB): Hardcoded JVM memory set aside for Spark system processes.
  3. User Memory: Calculated as (Heap Memory - 300MB) * (1.0 - spark.memory.fraction). Stores custom RDD transformations, user-defined data structures, and hash tables.

  4. Unified Memory: Calculated as (Heap Memory - 300MB) * spark.memory.fraction (default: 0.6). Shared dynamically between:

    • Execution Memory: Used for shuffles, joins, aggregations, and sorts.
    • Storage Memory: Used for cached data (.cache(), .persist()) and broadcast variables. Controlled by spark.memory.storageFraction (default: 0.5).

  5. PySpark Python Worker Memory (spark.python.worker.memory):

  6. Python processes execute UDFs (User-Defined Functions) and interact with PyArrow buffers. This memory sits outside the JVM Heap, but within the physical container boundary.

  7. Memory Overhead (spark.executor.memoryOverhead):

  8. Allocates memory for JVM Metaspace, thread stacks, native C/C++ allocations (e.g., PyArrow), and I/O buffers.
  9. Default value: max(384MB, 0.10 * spark.executor.memory).

2. Diagnosing OOM Errors: Driver vs. Executor

When PySpark jobs crash with memory exceptions, pinpointing where the failure occurred is essential.

                  +-----------------------------------+
                  |      PySpark OOM Failure          |
                  +-----------------------------------+
                                    |
            +-----------------------+-----------------------+
            |                                               |
  [ Driver Logs / UI ]                           [ Executor Logs / UI ]
            |                                               |
    +-------+-------+                               +-------+-------+
    |               |                               |               |
 Driver Heap     Driver Overhead               Executor Heap    Executor Overhead / Container
(toPandas,      (Huge metadata,               (Data Skew, GC    (Container OOMKilled - 137,
 collect)       Py4J Socket leak)             Spill Failure)     Python Memory breach)

Driver OOM Symptoms

  • Error Pattern: java.lang.OutOfMemoryError: Java heap space appearing in Driver logs.
  • Cause: Pulling too much data back to the driver using .collect(), .toPandas(), or broadcasting massive tables using broadcast().
  • Impact: The master controller crashes; the entire application terminates immediately.

Executor OOM Symptoms

  • Error Pattern 1: java.lang.OutOfMemoryError: Java heap space inside Task Executor logs.
  • Error Pattern 2: Container killed by YARN for exceeding memory limits or Kubernetes OOMKilled (Exit Code 137).
  • Cause: High partition skew, memory-intensive transformations (e.g., wide joins, non-aggregated groupBy), or insufficient memoryOverhead for PyArrow/Python operations.

3. Root Causes of PySpark OOM Errors and Fixes

Data Skew and Unbalanced Partitions

The Problem: Data skew happens when data is unevenly distributed across partitions. A single key (e.g., a NULL key or a popular category) contains millions of records, while other partitions contain thousands. The executor handling the bloated partition runs out of JVM heap space.

Unbalanced Partitions (Skewed):
Executor 1: [Key A: 100 rows]
Executor 2: [Key B: 100 rows]
Executor 3: [NULL/Key C: 10,000,000 rows]  <-- OOM Crash Happens Here!

The Solution: Apply a Salting Technique to distribute the hot key evenly across multiple partitions.

Python Code: Resolving Data Skew with Salting

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

spark = SparkSession.builder \
    .appName("FixDataSkewWithSalting") \
    .getOrCreate()

# Simulating skewed dataset (df_large) and lookup dataset (df_small)
# Assume 'user_id' has a massive skew with millions of nulls or default IDs

SALT_FACTOR = 8

# Step 1: Add a random salt key to the large skewed DataFrame
df_large_salted = df_large.withColumn(
    "salt_key", 
    F.concat(F.col("user_id"), F.lit("_"), F.floor(F.rand() * SALT_FACTOR))
)

# Step 2: Explode the lookup DataFrame so it matches all potential salt combinations
df_small_exploded = df_small.withColumn(
    "salt_array", 
    F.array([F.lit(i) for i in range(SALT_FACTOR)])
).withColumn(
    "exploded_salt", 
    F.explode("salt_array")
).withColumn(
    "salt_key", 
    F.concat(F.col("user_id"), F.lit("_"), F.col("exploded_salt"))
).drop("salt_array", "exploded_salt")

# Step 3: Perform join on the salted key
df_joined = df_large_salted.join(
    df_small_exploded,
    on="salt_key",
    how="inner"
).drop("salt_key")

# Step 4: Write output without executor OOM
df_joined.write.mode("overwrite").parquet("s3a://production-bucket/processed_data/")

Improper Broadcast Joins

The Problem: Broadcast joins bypass standard Spark shuffles by sending an entire DataFrame to every executor node. However, if the DataFrame exceeds the available driver memory or executor storage memory, it crashes the process.

The Solution: Adjust the auto-broadcast threshold or disable implicit broadcasting when working with dynamic datasets.

# Check current threshold (default is 10MB)
print(spark.conf.get("spark.sql.autoBroadcastJoinThreshold"))

# Option A: Disable broadcast joins completely for large datasets
spark.conf.set("spark.sql.autoBroadcastJoinThreshold", -1)

# Option B: Safely increase broadcast threshold to 50MB (if memory permits)
spark.conf.set("spark.sql.autoBroadcastJoinThreshold", 50 * 1024 * 1024)

# Option C: Use explicit Shuffle Hash Join hint instead of Broadcast
df_result = df_large.join(df_medium.hint("SHUFFLE_HASH"), "account_id")

PySpark Python-JVM Memory Overhead

The Problem: Converting large PySpark DataFrames to Pandas using .toPandas() pulls all data into the Driver JVM heap, serializes it across Py4J sockets, and deserializes it into C-structures in Python. This frequently triggers an unrecoverable Driver OOM.

PySpark DataFrame ---> Py4J Bridge ---> JVM Heap ---> Python Socket ---> C-Pandas Alloc (OOM!)

The Solution: Enable PyArrow optimization to enable zero-copy/fast binary serialization, or use PySpark's native Pandas API.

# Enable Apache Arrow optimization in PySpark
spark.conf.set("spark.sql.execution.arrow.pyspark.enabled", "true")
spark.conf.set("spark.sql.execution.arrow.pyspark.fallback.enabled", "false")

# BAD: Pulling entire dataset to Driver Pandas
# pandas_df = spark_df.toPandas() 

# GOOD: Use PySpark Pandas API (formerly Koalas) for distributed computation
import pyspark.pandas as ps

ps_df = spark_df.pandas_api()
summary = ps_df.groupby("category").mean()

# GOOD: Process in batches using Vectorized UDFs (Pandas UDF)
from pyspark.sql.functions import pandas_udf
import pandas as pd

@pandas_udf("double")
def calculate_tax_udf(v: pd.Series) -> pd.Series:
    return v * 0.20

df_transformed = spark_df.withColumn("tax", calculate_tax_udf("amount"))

Garbage Collection (GC) Pauses

The Problem: Long Java Garbage Collection pauses occur when millions of short-lived objects are instantiated during wide transformations. If GC takes longer than spark.network.timeout (default 120s), the driver assumes the executor is dead, triggering executor loss and task retries.

The Solution: Switch to the G1GC (Garbage-First Garbage Collector) and tune parameters.

Include these configurations in your spark-submit script:

spark-submit \
  --master k8s://https://10.0.0.1:6443 \
  --conf "spark.executor.extraJavaOptions=-XX:+UseG1GC -XX:InitiatingHeapOccupancyPercent=35 -XX:G1ReservePercent=15" \
  --conf "spark.driver.extraJavaOptions=-XX:+UseG1GC -XX:InitiatingHeapOccupancyPercent=35" \
  --class com.bitcodematrix.pipeline.Main \
  local:///opt/spark/work-dir/app.jar
  • -XX:+UseG1GC: Enables Garbage-First GC, optimal for high-memory heaps (>4GB).
  • -XX:InitiatingHeapOccupancyPercent=35: Starts GC early (at 35% heap occupancy) to prevent catastrophic full GC pauses.
  • -XX:G1ReservePercent=15: Keeps a 15% free buffer to prevent allocation failures during concurrent cycles.

4. Essential PySpark Memory Configurations

The table below outlines key parameters for tuning PySpark memory usage:

Configuration Parameter Default Value Recommended Tuning Strategy
spark.executor.memory 1g Allocate 4GB–16GB per executor based on worker core counts.
spark.driver.memory 1g Set to 4GB+ for large driver-side aggregation or broadcast operations.
spark.executor.memoryOverhead max(384MB, 10%) Increase to 0.20–0.30 (20-30%) when using heavy PyArrow, Pandas, or C-libraries.
spark.memory.fraction 0.6 Reduce to 0.4 if memory is needed for user objects; increase to 0.8 for heavy joins.
spark.memory.storageFraction 0.5 Lower to 0.2 if cached DataFrames are minimal, freeing space for joins.
spark.python.worker.memory 512m Set explicit upper bounds for Python process heap allocations.
spark.sql.shuffle.partitions 200 Adjust dynamically: set to 2 * (Total CPU Cores) or target 100MB–200MB per partition.
spark.sql.adaptive.enabled true Keep enabled! Allows Spark to coalesce small partitions and handle skew at runtime.

5. Kubernetes (Spark-on-k8s) & Infrastructure Debugging

When running PySpark on Kubernetes, the Kubernetes OOM Killer terminates containers whose total memory consumption (JVM Heap + Python Workers + Native Overhead) exceeds the Pod memory limit.

This shows up as an Exit Code 137.

Pod Status: OOMKilled
Exit Code: 137
Container: spark-kubernetes-executor

Manifest Fix: SparkApplication Custom Resource (Kubernetes Operator)

apiVersion: "sparkoperator.k8s.io/v1beta2"
kind: SparkApplication
metadata:
  name: pyspark-memory-optimized-job
  namespace: data-workloads
spec:
  type: Python
  mode: cluster
  image: "docker.io/bitcodematrix/spark-py:v3.5.0"
  imagePullPolicy: Always
  mainApplicationFile: "local:///opt/spark/jobs/production_pipeline.py"
  sparkVersion: "3.5.0"
  driver:
    cores: 2
    coreLimit: "2100m"
    memory: "4g"
    memoryOverhead: "1024m"
    labels:
      version: 3.5.0
    serviceAccount: spark-sa
  executor:
    cores: 4
    coreLimit: "4300m"
    instances: 10
    memory: "8g"
    # Essential: Account for PyArrow and Python native code memory overhead!
    memoryOverhead: "3072m" 
  hadoopConf:
    "fs.s3a.impl": "org.apache.hadoop.fs.s3a.S3AFileSystem"
  sparkConf:
    "spark.sql.adaptive.enabled": "true"
    "spark.sql.adaptive.skewJoin.enabled": "true"
    "spark.executor.extraJavaOptions": "-XX:+UseG1GC -XX:InitiatingHeapOccupancyPercent=35"

Linux Infrastructure Commands for Debugging OOM Pods

When an executor Pod crashes, inspect node-level cgroup metrics and event logs using these Linux CLI commands:

# 1. Fetch Pod logs for JVM error traces
kubectl logs -n data-workloads spark-executor-pod-1 --tail=500 | grep -i "OutOfMemory"

# 2. Inspect Pod termination reason (Verify if code 137 occurred)
kubectl describe pod -n data-workloads spark-executor-pod-1 | grep -A 5 "State:"

# 3. Stream real-time resource utilization (RAM/CPU) across active executors
kubectl top pods -n data-workloads -l spark-role=executor

# 4. SSH/Exec into a running executor container to monitor active CGroup limits
kubectl exec -it -n data-workloads spark-executor-pod-1 -- bash

# [Inside Container] Check current CGroup Memory usage vs Hard Limits (CGroup v2)
cat /sys/fs/cgroup/memory.current
cat /sys/fs/cgroup/memory.max

# [Inside Container] Inspect top Python worker processes by memory utilization
ps aux --sort=-%mem | head -n 10

6. Common Production Mistakes

  1. Calling .collect() on Production DataFrames: Calling .collect() pulls all distributed data into driver memory. Use .take(n), .limit(n), or write directly to storage targets (S3, GCS, Snowflake).
  2. Ignoring Data Spill to Disk: When Execution Memory fills up, Spark spills excess data to local disk. While this prevents immediate OOMs, it drastically slows down jobs. Check the Spark UI Spill (Memory) vs. Spill (Disk) metrics.
  3. Persisting Data Without Explicit Unpersisting: Calling .cache() or .persist() keeps DataFrames in Storage Memory indefinitely. Always call .unpersist() once downstream transformations finish.
  4. Setting Partition Counts Too Low: Leaving spark.sql.shuffle.partitions at its default value (200) when processing terabytes of data causes individual tasks to handle tens of gigabytes, leading to JVM heap failure.
  5. Misunderstanding memoryOverhead in PySpark: Assuming spark.executor.memory controls all memory usage. Python workers and PyArrow rely heavily on memoryOverhead.

7. Production Best Practices Checklist

  • [ ] Enable Adaptive Query Execution (AQE): Set spark.sql.adaptive.enabled = true to allow dynamic partition coalescing and runtime skew join handling.
  • [ ] Tune Memory Overhead: Set spark.executor.memoryOverhead to at least 20–30% of executor memory for PySpark workflows using PyArrow/Pandas.
  • [ ] Enable Vectorized Engine (Arrow): Ensure spark.sql.execution.arrow.pyspark.enabled = true is explicitly configured.
  • [ ] Monitor GC Overhead: Switch JVM execution flag to -XX:+UseG1GC.
  • [ ] Standardize Disk Partition Sizes: Aim for shuffled output partitions sized between 100 MB and 200 MB.
  • [ ] Audit Broadcast Join Sizes: Keep spark.sql.autoBroadcastJoinThreshold conservative (< 100MB) to protect the Driver node.
  • [ ] Always Cleanup Cache: Clear intermediate cache entries via df.unpersist() in a try...finally block.

8. Interview Questions & Answers

Q1: What is the difference between Execution Memory and Storage Memory in Spark?

Answer: Both sit within Spark's Unified Memory space (spark.memory.fraction). Execution Memory hosts transient data structures needed during joins, shuffles, and aggregations (e.g., hash join tables). Storage Memory stores cached data frames (.cache(), .persist()) and broadcast variables.

Execution memory takes priority over storage memory; if execution needs space, it can evict cached data from storage until the limit set by spark.memory.storageFraction is reached.

Q2: What causes a Container killed by YARN / Kubernetes (Exit Code 137) error in PySpark?

Answer: Exit Code 137 occurs when the OS/CGroup memory manager terminates the Spark container for exceeding its allocated hard memory limit.

In PySpark, this is usually caused by off-heap allocations, native memory leaks, or heavy Python process memory usage (e.g., NumPy/Pandas/PyArrow arrays) that exceeds the designated spark.executor.memoryOverhead boundary.

Q3: How do you identify and resolve Data Skew in PySpark?

Answer: Identify data skew using the Spark UI by observing if a few tasks take significantly longer to complete than others, or if single tasks fail with Heap OOMs while others finish instantly.

To resolve skew: 1. Enable Adaptive Query Execution (spark.sql.adaptive.skewJoin.enabled = true). 2. Apply a Salting Technique: append isolated random integers to hot keys to distribute them across distinct execution partitions, perform the join, and aggregate the results.

Q4: Why can spark.executor.memory be set to 8GB while the Kubernetes Executor Pod still gets killed for using 12GB?

Answer: spark.executor.memory only controls the JVM Heap. The total memory footprint of a Kubernetes Pod is:

$$\text{Total Pod Memory} = \text{JVM Heap} + \text{Memory Overhead} + \text{Python Worker Memory}$$

If native C-libraries, PyArrow buffers, or Python UDFs run inside the container, off-heap memory growth can easily push total usage past 12GB. The host node's Linux kernel then terminates the Pod via the OOM Killer.

Q5: What is the risk of using .toPandas() in PySpark, and how should you handle large data exports?

Answer: .toPandas() collects all distributed dataset partitions onto the single Driver JVM process, converts them across Py4J, and builds a single monolithic Pandas DataFrame in memory. If the dataset exceeds Driver RAM, the driver process crashes immediately.

To safely handle large exports: * Write data in parallel to storage (e.g., Parquet/Delta on S3/GCS). * Use the distributed PySpark Pandas API (import pyspark.pandas as ps). * Use .toLocalIterator() to stream partitions sequentially instead of loading everything at once.

9. Frequently Asked Questions (FAQ)

1. Why is my PySpark job failing with OOM even though my total cluster RAM is larger than the input data size?

Input file size on disk does not reflect memory footprint in RAM. Compressed formats like Parquet/ORC expand up to 5x–10x when deserialized into uncompressed memory objects. Wide joins and shuffles also generate intermediate datasets that multiply actual memory consumption.

2. Should I use dynamic allocation (spark.dynamicAllocation.enabled) to solve OOM errors?

Dynamic allocation scales the number of executor instances up or down based on queue workload. However, it does not increase the memory allocated to individual executors. If a single task fails due to a bloated partition, scaling out to more executors won't fix the issue; you need to increase executor memory resources or rebalance your partitions.

3. What is the optimal number of cores per executor (spark.executor.cores)?

A common operational standard is 4 to 5 cores per executor. Running with more than 5 cores can lead to high Garbage Collection contention, while running with fewer than 2 cores reduces parallel processing efficiency within the same JVM instance.

4. Does PySpark use off-heap memory by default?

Spark uses off-heap memory for certain internal tasks (Netty buffers, Metaspace), but project Tungsten off-heap execution (spark.memory.offHeap.enabled) is disabled by default. Enabling explicit off-heap memory allocation can improve performance by bypassing Java GC overhead for large allocations.

5. How do I clear cached DataFrames from memory?

Call df.unpersist() to free memory used by cached DataFrames. You can also run spark.catalog.clearCache() to remove all cached DataFrames from memory at once.

10. Conclusion

Fixing PySpark OutOfMemory errors comes down to understanding the dual JVM-Python runtime, managing memory allocation boundaries, and distributing partition workloads evenly.

Instead of arbitrarily scaling cluster hardware, systematically diagnose your job failures: * Determine whether the driver or executor failed. * Check if high memory usage stems from off-heap overhead or skewed data partitions. * Use salting to rebalance skewed data, enable G1GC for faster garbage collection, tune PyArrow options, and configure container memory limits appropriately.

By applying these memory management techniques, you can transform unstable, high-cost data jobs into predictable, cost-efficient production pipelines.

By at

No comments:

Post a Comment