Spark Performance Masterclass

Chapter 1

Chapter 1: The Spark Execution Model

Understanding how Spark executes your code is the foundation of all performance tuning.
Without this mental model, every optimization is just guesswork.

1. The Big Picture

2. Driver and Executors

3. From Code to Execution: The Query Planning Pipeline

4. Jobs, Stages, and Tasks

5. Narrow vs Wide Transformations

6. The Shuffle: Spark's Most Expensive Operation

7. Partitions: The Unit of Parallelism

8. Lazy Evaluation and the DAG

9. Configuration Settings

10. Code Examples with Before/After Metrics

The Big Picture

When you write df.groupBy("region").agg(sum("revenue")), Spark does not immediately

execute that computation. Instead, it builds a Directed Acyclic Graph (DAG) of

transformations. Only when you trigger an action (like .show(), .write(), or

.count()) does Spark compile this DAG into an optimized execution plan and distribute

the work across a cluster.

Your Code (PySpark / SQL)
        |
        v
Unresolved Logical Plan
        |   <- Catalog lookup (resolve table/column names)
        v
Analyzed Logical Plan
        |   <- Catalyst rule-based optimization
        v
Optimized Logical Plan
        |   <- Physical planning + cost-based optimization
        v
Selected Physical Plan
        |   <- Whole-stage code generation
        v
RDDs (internal execution)
        |
        v
Tasks distributed to Executors

This pipeline is important because **most performance problems originate in how the

physical plan is constructed**, not in your code itself.

Driver and Executors

The Driver

The driver is a single JVM process responsible for:

Parsing your code into a logical plan
Optimizing the plan via the Catalyst optimizer
Scheduling tasks onto executors
Collecting results from actions like collect() or show()
Maintaining the SparkContext and all metadata

Common driver bottlenecks:

Issue	Cause	Impact
OOM on `collect()`	Pulling large dataset to driver	Driver crash
Slow scheduling	Too many small tasks (> 100K)	High latency
Broadcast timeout	Broadcasting table > driver memory	Job failure
Metadata overhead	Too many partitions in Delta table	Slow planning

Driver configuration:

python

# Databricks cluster config (set in cluster configuration, not in code)
# spark.driver.memory: 8g          # Default is often sufficient
# spark.driver.maxResultSize: 4g   # Max size of serialized results

# In your notebook, you can check:
print(f"Driver memory: {spark.conf.get('spark.driver.memory')}")
print(f"Max result size: {spark.conf.get('spark.driver.maxResultSize')}")

Executors

Executors are JVM processes on worker nodes that:

Execute tasks assigned by the driver
Store data in memory or on disk (cache, shuffle data)
Report status back to the driver

Each executor has:

A fixed number of cores (slots for concurrent tasks)
A fixed amount of memory (split between execution, storage, and overhead)

Executor configuration:

python

# Databricks typically manages these, but you can override:
# spark.executor.memory: 8g        # Heap memory per executor
# spark.executor.cores: 4          # Concurrent tasks per executor
# spark.executor.memoryOverhead: 1g # Off-heap memory (PySpark needs more)

# Check current settings:
print(f"Executor memory: {spark.conf.get('spark.executor.memory')}")
print(f"Executor cores: {spark.conf.get('spark.executor.cores')}")

The Relationship

                    +------------------+
                    |     DRIVER       |
                    | - SparkContext   |
                    | - DAG Scheduler  |
                    | - Task Scheduler |
                    +--------+---------+
                             |
              +--------------+--------------+
              |              |              |
    +---------+----+ +------+-------+ +----+---------+
    | EXECUTOR #1  | | EXECUTOR #2  | | EXECUTOR #3  |
    | Core1 Core2  | | Core1 Core2  | | Core1 Core2  |
    | Core3 Core4  | | Core3 Core4  | | Core3 Core4  |
    | [Cache]      | | [Cache]      | | [Cache]      |
    | [Shuffle]    | | [Shuffle]    | | [Shuffle]    |
    +--------------+ +--------------+ +--------------+

Key insight: The driver is a single point of failure and a potential bottleneck.

Never send large amounts of data to it. The executors do the real work.

From Code to Execution

Step 1: Unresolved Logical Plan

When you write DataFrame operations, Spark builds a tree of operations:

python

# This code creates a logical plan, NOT immediate execution
df = (
    spark.read.table("sales.transactions")
    .filter(col("year") == 2025)
    .groupBy("region")
    .agg(
        sum("revenue").alias("total_revenue"),
        count("*").alias("num_transactions")
    )
)
# Nothing has executed yet!

The unresolved plan has references to table and column names that haven't been validated.

Step 2: Analysis (Catalog Resolution)

Spark resolves table names against the catalog (Unity Catalog in Databricks), validates

column names exist, and resolves data types.

python

# View the analyzed plan:
df.explain(mode="formatted")

# Or in Databricks, use:
# display(df)  # triggers execution and shows plan in the UI

Step 3: Logical Optimization (Catalyst)

The Catalyst optimizer applies rule-based optimizations:

Rule	What It Does	Example
Predicate Pushdown	Pushes filters closer to the data source	Filter before join
Column Pruning	Drops unused columns early	Only read needed cols
Constant Folding	Pre-computes constant expressions	`1 + 2` becomes `3`
Boolean Simplification	Simplifies boolean logic	`NOT (NOT x)` becomes `x`
Combine Filters	Merges adjacent filters	Two `.filter()` become one

python

# See the optimized logical plan:
df.explain(mode="extended")

Step 4: Physical Planning

The optimizer generates candidate physical plans and selects the best one using

cost-based optimization (CBO):

python

# Enable cost-based optimization (on by default in Databricks):
spark.conf.set("spark.sql.cbo.enabled", "true")
spark.conf.set("spark.sql.cbo.joinReorder.enabled", "true")

# CBO requires table statistics. Compute them:
spark.sql("ANALYZE TABLE sales.transactions COMPUTE STATISTICS FOR ALL COLUMNS")

Step 5: Code Generation

Spark uses whole-stage code generation to compile the physical plan into optimized

Java bytecode, eliminating virtual function dispatch overhead:

python

# Verify codegen is enabled (default: true):
print(spark.conf.get("spark.sql.codegen.wholeStage"))

# View the generated code:
df.explain(mode="codegen")

Jobs, Stages, and Tasks

The Hierarchy

Action (e.g., .count(), .write())
  └── Job (one per action)
        ├── Stage 0 (all narrow transforms before first shuffle)
        │     ├── Task 0 (partition 0)
        │     ├── Task 1 (partition 1)
        │     └── Task N (partition N)
        ├── Stage 1 (after shuffle, before next shuffle)
        │     ├── Task 0
        │     └── Task M
        └── Stage 2 (final stage)
              ├── Task 0
              └── Task K

Jobs

Every action creates a new job. Common actions:

python

df.count()           # Job 1
df.show()            # Job 2
df.write.save(...)   # Job 3
df.collect()         # Job 4 (dangerous for large data!)
df.toPandas()        # Job 5 (also dangerous - collects to driver)

Stages

Stages are bounded by shuffle operations (wide transformations). Within a stage,

all operations can be pipelined within a single task without exchanging data between

executors.

python

# This creates 3 stages:
result = (
    spark.read.parquet("/data/input")         # Stage 0: Read
    .filter(col("status") == "active")        # Stage 0: Filter (narrow)
    .groupBy("department")                    # --- SHUFFLE BOUNDARY ---
    .agg(sum("salary"))                       # Stage 1: Aggregate
    .join(department_lookup, "department")     # --- SHUFFLE BOUNDARY ---
    .select("department", "total_salary",     # Stage 2: Project
            "department_name")
)
result.write.parquet("/data/output")          # Triggers execution

Tasks

One task processes one partition. The total parallelism of a stage equals the number

of partitions in that stage:

python

# Check partition count:
print(f"Number of partitions: {df.rdd.getNumPartitions()}")

# Repartition to control parallelism:
df_repartitioned = df.repartition(200)   # Shuffle repartition
df_coalesced = df.coalesce(50)           # No-shuffle reduction only

Critical rule: Aim for tasks that run 30 seconds to 3 minutes each. Shorter

tasks have too much scheduling overhead. Longer tasks risk failures and poor utilization.

Narrow vs Wide Transformations

This distinction is the single most important concept for performance:

Narrow Transformations (No Shuffle)

Data stays on the same executor. One input partition maps to one output partition.

python

# All of these are narrow (no shuffle):
df.filter(col("age") > 30)              # Filter
df.select("name", "age")                # Select
df.withColumn("age_plus", col("age")+1) # WithColumn
df.drop("temp_col")                     # Drop
df.limit(100)                           # Limit (mostly narrow)
df.union(other_df)                      # Union

Wide Transformations (Shuffle Required)

Data must be exchanged between executors. These create stage boundaries.

python

# All of these trigger a shuffle:
df.groupBy("key").agg(sum("value"))     # Aggregation
df.join(other_df, "key")                # Join (unless broadcast)
df.repartition(100)                     # Explicit repartition
df.distinct()                           # Distinct
df.orderBy("timestamp")                 # Sort
df.rollup("a", "b").agg(sum("c"))      # Rollup

Performance Impact

Narrow transforms:    O(data_size) - linear scan, no network I/O
Wide transforms:      O(data_size) + O(shuffle_size) + O(network_IO) + O(disk_IO)

Before (naive approach):

python

# Multiple unnecessary shuffles
result = (
    df.repartition(200)                      # Shuffle #1 (unnecessary)
    .groupBy("region").agg(sum("revenue"))   # Shuffle #2
    .orderBy("total_revenue")                # Shuffle #3
)
# Total: 3 shuffles
# Execution time: 180 seconds
# Shuffle write: 12 GB

After (optimized):

python

# Minimize shuffles
result = (
    df.groupBy("region").agg(sum("revenue").alias("total_revenue"))  # Shuffle #1
    .orderBy(desc("total_revenue"))                                   # Shuffle #2
)
# Total: 2 shuffles (repartition removed, sort still needed)
# Execution time: 95 seconds
# Shuffle write: 0.5 GB (aggregation reduces data before sort)

The Shuffle: Spark's Most Expensive Operation

What Happens During a Shuffle

STAGE 1 (Map Side)                    STAGE 2 (Reduce Side)
+-------------------+                +-------------------+
| Executor 1        |   Network      | Executor 1        |
| Partition A ------+---+----------->| Reduce Task 1     |
| Partition B ------+---+--+         | (key range 0-99)  |
+-------------------+   |  |         +-------------------+
                        |  |
+-------------------+   |  |         +-------------------+
| Executor 2        |   |  +-------->| Executor 2        |
| Partition C ------+---+----------->| Reduce Task 2     |
| Partition D ------+---+            | (key range 100+)  |
+-------------------+                +-------------------+

Steps:
1. Map tasks write shuffle files (sorted by partition key)
2. Shuffle service transfers files across the network
3. Reduce tasks read their relevant portions

Shuffle Metrics to Monitor

Metric	What It Means	Healthy Range
Shuffle Write	Data written by map side	< input size
Shuffle Read	Data read by reduce side	≈ shuffle write
Shuffle Spill (Memory)	Data spilled before write	0 bytes
Shuffle Spill (Disk)	Data spilled to disk	0 bytes
Records Read	Total records shuffled	Depends on use case
Fetch Wait Time	Time waiting for shuffle data	< 10% of task time

Configuration

python

# Number of partitions after a shuffle (critical setting!)
spark.conf.set("spark.sql.shuffle.partitions", "200")
# Rule of thumb: target 128-256 MB per partition after shuffle

# Shuffle compression (always enable)
spark.conf.set("spark.shuffle.compress", "true")
spark.conf.set("spark.io.compression.codec", "zstd")  # Better than lz4 for shuffle

# External shuffle service (enabled by default on Databricks)
spark.conf.set("spark.shuffle.service.enabled", "true")

Partitions: The Unit of Parallelism

Input Partitions

Determined by the data source:

python

# Parquet/Delta: one partition per file (roughly)
df = spark.read.parquet("/data/large_dataset/")
print(f"Input partitions: {df.rdd.getNumPartitions()}")

# You can control max partition size:
spark.conf.set("spark.sql.files.maxPartitionBytes", "128m")  # Default: 128MB
spark.conf.set("spark.sql.files.openCostInBytes", "4m")      # Cost to open a file

Shuffle Partitions

After any wide transformation:

python

# Global setting for all shuffles:
spark.conf.set("spark.sql.shuffle.partitions", "200")  # Default in OSS Spark
# Databricks default is "auto" with AQE

# The right number depends on data size after the shuffle:
# Target: 128-256 MB per partition
# Formula: shuffle_data_size / target_partition_size
# Example: 100 GB shuffle / 200 MB target = 500 partitions

Repartitioning Strategies

python

# Hash repartition (shuffle): distributes evenly by hash of key
df.repartition(200, "customer_id")

# Range repartition: useful before writes for partition pruning
df.repartitionByRange(200, "date")

# Coalesce (no shuffle): only reduces partitions
df.coalesce(50)  # Combine partitions on same executor

# IMPORTANT: coalesce can create uneven partitions if starting
# partitions have very different sizes!

Before/After: Partition Sizing

Before (too many small partitions):

python

# 10,000 partitions of 10 MB each = 100 GB
# Each partition creates overhead: task scheduling, memory management
# Metrics:
#   - Tasks: 10,000
#   - Median task time: 0.3 seconds
#   - Scheduler delay: 45% of total time
#   - Total time: 420 seconds

After (right-sized partitions):

python

# 500 partitions of 200 MB each = 100 GB
spark.conf.set("spark.sql.shuffle.partitions", "500")
# Metrics:
#   - Tasks: 500
#   - Median task time: 12 seconds
#   - Scheduler delay: 2% of total time
#   - Total time: 85 seconds

Improvement: 5x faster by simply adjusting partition count.

Lazy Evaluation and the DAG

How Lazy Evaluation Works

Spark transformations are lazy: they build a plan but don't execute it.

python

# None of these lines trigger execution:
df1 = spark.read.table("events")
df2 = df1.filter(col("event_type") == "purchase")
df3 = df2.groupBy("user_id").count()
df4 = df3.filter(col("count") > 5)

# THIS triggers execution:
df4.show()

Why This Matters for Performance

Lazy evaluation enables the Catalyst optimizer to see the entire pipeline before

executing anything. This means it can:

1. Push filters down before joins (even if you wrote the filter after the join)

2. Eliminate unnecessary columns early

3. Choose optimal join strategies based on the full plan

4. Combine operations into single stages

python

# You write this (filter after join):
result = (
    big_table.join(small_table, "key")
    .filter(col("status") == "active")
)

# Catalyst optimizes it to (filter before join):
# 1. Filter big_table where status = 'active'
# 2. Join filtered big_table with small_table
# Much less data to shuffle!

Breaking Lazy Evaluation (Performance Anti-Patterns)

python

# ANTI-PATTERN 1: Unnecessary actions
count = df.count()  # Forces execution of entire DAG
if count > 0:
    df.write.save(...)  # DAG executes AGAIN from scratch!

# FIX: Use caching or restructure
df.cache()
count = df.count()
if count > 0:
    df.write.save(...)
df.unpersist()

# ANTI-PATTERN 2: Collecting to driver for logic
rows = df.collect()  # Pulls ALL data to driver
filtered = [r for r in rows if r["value"] > threshold]  # Python loop

# FIX: Keep it in Spark
filtered_df = df.filter(col("value") > threshold)

Configuration Settings

Essential Settings for Execution Model Tuning

python

# === Core Execution Settings ===

# Shuffle partitions - the most impactful single setting
spark.conf.set("spark.sql.shuffle.partitions", "auto")  # Let AQE handle it

# Max partition bytes for reads
spark.conf.set("spark.sql.files.maxPartitionBytes", "134217728")  # 128 MB

# Whole-stage code generation
spark.conf.set("spark.sql.codegen.wholeStage", "true")  # Default: true

# === Parallelism Settings ===

# Default parallelism for RDD operations
spark.conf.set("spark.default.parallelism", "200")

# Max concurrent tasks per executor
# This is controlled by spark.executor.cores

# === Timeout and Retry Settings ===

# Task retry on failure
spark.conf.set("spark.task.maxFailures", "4")  # Default: 4

# Network timeout for shuffle
spark.conf.set("spark.network.timeout", "300s")

# Shuffle fetch retry
spark.conf.set("spark.shuffle.io.maxRetries", "5")
spark.conf.set("spark.shuffle.io.retryWait", "10s")

# === Cost-Based Optimizer ===
spark.conf.set("spark.sql.cbo.enabled", "true")
spark.conf.set("spark.sql.cbo.joinReorder.enabled", "true")
spark.conf.set("spark.sql.statistics.histogram.enabled", "true")

Code Examples

Example 1: Understanding Your Execution Plan

python

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, sum, count, desc

spark = SparkSession.builder.getOrCreate()

# Build a multi-stage pipeline
orders = spark.read.table("catalog.schema.orders")
customers = spark.read.table("catalog.schema.customers")

result = (
    orders
    .filter(col("order_date") >= "2025-01-01")        # Narrow: filter
    .join(customers, "customer_id")                     # Wide: shuffle/broadcast
    .groupBy("region", "product_category")              # Wide: shuffle
    .agg(
        sum("amount").alias("total_amount"),
        count("*").alias("order_count")
    )
    .orderBy(desc("total_amount"))                      # Wide: sort
)

# Examine the plan BEFORE executing
result.explain(mode="formatted")
# Look for:
#   - Number of exchanges (shuffles)
#   - Join strategy (BroadcastHashJoin vs SortMergeJoin)
#   - Filter placement (pushed down?)
#   - Scan columns (pruned?)

Example 2: Monitoring Execution Metrics Programmatically

python

from pyspark.sql import SparkSession
from typing import Any

def get_stage_metrics(spark: SparkSession) -> list[dict[str, Any]]:
    """Extract stage metrics from the most recent job.

    Returns a list of stage dictionaries with key metrics.
    """
    sc = spark.sparkContext
    status_tracker = sc.statusTracker()

    job_ids = status_tracker.getJobIdsForGroup()
    if not job_ids:
        return []

    latest_job = max(job_ids)
    job_info = status_tracker.getJobInfo(latest_job)
    if job_info is None:
        return []

    stages = []
    for stage_id in job_info.stageIds:
        stage_info = status_tracker.getStageInfo(stage_id)
        if stage_info is not None:
            stages.append({
                "stage_id": stage_id,
                "num_tasks": stage_info.numTasks,
                "num_active_tasks": stage_info.numActiveTasks,
                "num_completed_tasks": stage_info.numCompletedTasks,
            })
    return stages


# Usage:
spark.sparkContext.setJobGroup("perf_test", "Performance test job")
result = df.groupBy("key").count()
result.write.format("noop").mode("overwrite").save()

metrics = get_stage_metrics(spark)
for stage in metrics:
    print(f"Stage {stage['stage_id']}: "
          f"{stage['num_tasks']} tasks, "
          f"{stage['num_completed_tasks']} completed")

Example 3: Minimizing Stage Count

python

# BEFORE: 5 stages, 4 shuffles
# Step 1: Read -> filter (Stage 0)
# Step 2: Repartition (UNNECESSARY shuffle - Stage 1)
# Step 3: Join with lookup (shuffle - Stage 2)
# Step 4: GroupBy aggregate (shuffle - Stage 3)
# Step 5: OrderBy (shuffle - Stage 4)

result_before = (
    spark.read.table("events")
    .filter(col("event_date") >= "2025-01-01")
    .repartition(200)                                   # Unnecessary!
    .join(lookup_table, "event_type")
    .groupBy("category").agg(sum("value").alias("total"))
    .orderBy(desc("total"))
)
# Execution time: 240s | Shuffle: 45 GB | Stages: 5

# AFTER: 4 stages, 3 shuffles (removed unnecessary repartition)
# Also: broadcast the small lookup table to eliminate one shuffle

from pyspark.sql.functions import broadcast

result_after = (
    spark.read.table("events")
    .filter(col("event_date") >= "2025-01-01")
    .join(broadcast(lookup_table), "event_type")        # Broadcast = no shuffle
    .groupBy("category").agg(sum("value").alias("total"))
    .orderBy(desc("total"))
)
# Execution time: 85s | Shuffle: 8 GB | Stages: 3

Key Takeaways

1. The driver is a coordinator, not a worker. Never send large data to it.

2. Catalyst sees your entire plan. Trust it - but verify with .explain().

3. Stages are bounded by shuffles. Fewer shuffles = fewer stages = faster jobs.

4. Every task processes one partition. Size partitions for 30s-3min task duration.

5. Lazy evaluation is your friend. Don't break it with unnecessary actions.

6. Wide transformations are expensive. Minimize them; use broadcast when possible.

7. Read the query plan before running. Most problems are visible in the plan.

Next Chapter

Chapter 2: Adaptive Query Execution Deep Dive - Learn how Spark

can dynamically optimize execution plans at runtime, automatically handling many of the

issues discussed in this chapter.

Chapter 2

Chapter 2: Adaptive Query Execution (AQE) Deep Dive

AQE is the single most impactful Spark feature since the introduction of DataFrames.
Understanding it means understanding why your modern Spark jobs are already faster than
they would have been five years ago — and how to make them faster still.

1. What Is AQE?

2. How AQE Works Under the Hood

3. Feature 1: Coalescing Post-Shuffle Partitions

4. Feature 2: Skew Join Optimization

5. Feature 3: Converting Sort-Merge to Broadcast Joins

6. Feature 4: Optimizing Skewed Aggregations

7. AQE Configuration Reference

8. When AQE Cannot Help

9. Debugging AQE Decisions

10. Code Examples with Before/After

What Is AQE?

Adaptive Query Execution (AQE) is a runtime optimization framework introduced in Spark

3.0 and significantly enhanced in subsequent releases. Unlike the Catalyst optimizer,

which makes all decisions at plan time (before any data is processed), AQE makes

optimization decisions at runtime based on actual data statistics collected during

shuffle operations.

The Core Insight

Static optimization faces a fundamental problem: **accurate statistics are often

unavailable**. Table statistics may be stale, column distributions unknown, and join

selectivity impossible to predict. AQE solves this by re-optimizing the query plan at

stage boundaries (shuffles), where it can observe the actual data.

Static Optimization (Catalyst):
  Plan -> Execute -> Done
  (decisions based on estimated statistics)

Adaptive Optimization (AQE):
  Plan -> Execute Stage 1 -> Observe Stats -> Re-Plan -> Execute Stage 2 -> ...
  (decisions based on real data at each stage boundary)

AQE in Databricks vs Open Source

Feature	OSS Spark 3.4+	Databricks Runtime 14+
Coalesce partitions	Yes	Yes (enhanced)
Skew join	Yes	Yes (enhanced)
Convert to broadcast	Yes	Yes (lower thresholds)
Skewed aggregation	Limited	Yes (full support)
Custom shuffle reader	Yes	Enhanced
Auto-tune partitions	Limited	`auto` setting

How AQE Works Under the Hood

The Re-Optimization Loop

1. Catalyst generates initial physical plan
2. Spark submits first set of stages (up to first shuffle)
3. Shuffle writes complete -> AQE collects statistics:
   - Partition sizes (bytes)
   - Partition row counts
   - Data distribution across partitions
4. AQE re-optimizes remaining plan using real statistics
5. Next stages execute with optimized plan
6. Repeat at each shuffle boundary

What Statistics AQE Collects

At each shuffle boundary, AQE knows:

Exact size (bytes) of each shuffle partition
Number of rows in each partition
Total data volume per shuffle
Distribution skew (max partition size / median partition size)

This information drives all four AQE optimizations.

Enabling AQE

python

# AQE is enabled by default in Databricks Runtime 12+
# For OSS Spark or older runtimes:
spark.conf.set("spark.sql.adaptive.enabled", "true")

# Verify it's active:
print(f"AQE enabled: {spark.conf.get('spark.sql.adaptive.enabled')}")

Feature 1: Coalescing Post-Shuffle Partitions

The Problem

spark.sql.shuffle.partitions is a static setting applied to ALL shuffles. But

different shuffles in the same query may produce vastly different amounts of data:

python

# shuffle.partitions = 200
# Stage 1 shuffle: 100 GB -> 200 partitions of 500 MB each (too large!)
# Stage 2 shuffle: 500 MB -> 200 partitions of 2.5 MB each (too small!)

How AQE Solves It

After a shuffle completes, AQE examines the actual partition sizes and **coalesces

adjacent small partitions** into larger ones:

Before AQE coalesce (200 partitions):
  [2MB][3MB][1MB][4MB][2MB][1MB][3MB][2MB]...  (200 tiny partitions)

After AQE coalesce (15 partitions):
  [48MB][52MB][45MB][55MB][50MB]...  (15 right-sized partitions)

Configuration

python

# Enable partition coalescing (default: true when AQE is on)
spark.conf.set("spark.sql.adaptive.coalescePartitions.enabled", "true")

# Target size for coalesced partitions
spark.conf.set("spark.sql.adaptive.advisoryPartitionSizeInBytes", "128m")
# Default: 64MB. Increase for large shuffles, decrease for streaming.

# Minimum number of partitions after coalescing
spark.conf.set("spark.sql.adaptive.coalescePartitions.minPartitionNum", "1")
# Set higher if you want to maintain parallelism

# Initial shuffle partition count (starting point for coalescing)
spark.conf.set("spark.sql.adaptive.coalescePartitions.initialPartitionNum", "4096")
# AQE will set shuffle.partitions to this value, then coalesce down.
# Set high to give AQE room to work.

# In Databricks, you can use "auto":
spark.conf.set("spark.sql.shuffle.partitions", "auto")
# This sets initial partitions high and lets AQE coalesce optimally.

Before/After Metrics

Before (static partitions = 200, small aggregate output):

Shuffle output: 800 MB across 200 partitions
Average partition: 4 MB
Tasks: 200
Median task time: 0.5s
Total stage time: 45s (dominated by task scheduling overhead)

After (AQE coalesces to 7 partitions):

Shuffle output: 800 MB across 7 partitions
Average partition: 114 MB
Tasks: 7
Median task time: 8s
Total stage time: 12s

Improvement: 3.75x faster with zero code changes.

Feature 2: Skew Join Optimization

The Problem

Data skew in joins is one of the most common and devastating performance problems.

When one key has disproportionately more rows than others, one task processes vastly

more data while other tasks sit idle:

Partition distribution by join key:
  Key A: 1,000 rows     -> Task finishes in 2s
  Key B: 500 rows        -> Task finishes in 1s
  Key C: 50,000,000 rows -> Task finishes in 3 hours (!)
  Key D: 800 rows        -> Task finishes in 1.5s

How AQE Solves It

AQE detects skewed partitions after the shuffle write and **splits them into multiple

smaller partitions**, duplicating the corresponding partition from the other side of

the join:

Without AQE skew optimization:
  Left side:  [Small][Small][HUGE_SKEWED][Small]
  Right side: [Small][Small][Normal     ][Small]
  -> One task processes HUGE_SKEWED x Normal = very slow

With AQE skew optimization:
  Left side:  [Small][Small][Skew_1][Skew_2][Skew_3][Small]
  Right side: [Small][Small][Dup_1 ][Dup_2 ][Dup_3 ][Small]
  -> Three tasks process Skew_1xDup_1, Skew_2xDup_2, Skew_3xDup_3
  -> Each task is ~1/3 the work

Configuration

python

# Enable skew join optimization (default: true when AQE is on)
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "true")

# A partition is considered skewed if:
# 1. Its size > skewedPartitionFactor * median partition size
# 2. AND its size > skewedPartitionThresholdInBytes

# Skew factor (how many times larger than median)
spark.conf.set("spark.sql.adaptive.skewJoin.skewedPartitionFactor", "5")
# Default: 5. Lower this for more aggressive skew handling.

# Minimum absolute size to be considered skewed
spark.conf.set("spark.sql.adaptive.skewJoin.skewedPartitionThresholdInBytes", "256m")
# Default: 256MB. Lower this for smaller datasets.

# For severe skew, consider more aggressive settings:
spark.conf.set("spark.sql.adaptive.skewJoin.skewedPartitionFactor", "2")
spark.conf.set("spark.sql.adaptive.skewJoin.skewedPartitionThresholdInBytes", "64m")

Before/After Metrics

Before (no skew handling):

Join: orders (500M rows) x customers (10M rows) on customer_id
Skew: Top customer has 50M orders (10% of all data)
Total tasks: 200
Slowest task: 2400s (processing 50M rows)
Median task time: 12s
Total wall clock: 2400s (bottlenecked on one task)
Shuffle write: 45 GB

After (AQE skew join):

Same join, AQE enabled
Total tasks: 245 (AQE split the skewed partition into ~45 sub-tasks)
Slowest task: 55s
Median task time: 14s
Total wall clock: 55s
Shuffle write: 47 GB (slightly more due to duplication)

Improvement: 43x faster with zero code changes.

Feature 3: Converting Sort-Merge to Broadcast Joins

The Problem

At plan time, Spark estimates table sizes to decide between broadcast and sort-merge

joins. If statistics are missing or stale, it may choose sort-merge when broadcast

would be much faster:

python

# Table A: estimated 500 MB (actually 500 MB after filtering)
# Table B: estimated 2 GB (but filter reduces it to 50 MB!)
# Catalyst chose SortMergeJoin based on estimated 2 GB
# But after filtering, B is only 50 MB - should have been broadcast!

How AQE Solves It

After the filter stage completes, AQE sees the actual output size. If one side of a

join is small enough, AQE converts the sort-merge join to a broadcast hash join:

Original plan (estimated):     AQE-optimized plan (actual):
SortMergeJoin                  BroadcastHashJoin
├── Scan A (500 MB est.)       ├── Scan A (500 MB actual)
└── Filter+Scan B (2 GB est.)  └── Broadcast(Filter+Scan B) (50 MB actual)

Configuration

python

# Broadcast join threshold for AQE runtime conversion
spark.conf.set("spark.sql.adaptive.autoBroadcastJoinThreshold", "100m")
# This is SEPARATE from the static broadcast threshold.
# AQE uses this threshold when deciding to convert at runtime.
# Can be set higher than the static threshold because it uses real sizes.

# Static broadcast threshold (used at plan time)
spark.conf.set("spark.sql.autoBroadcastJoinThreshold", "100m")
# Default: 10MB. Increase to 100-200MB for most workloads.

# Disable broadcast conversion if needed (rare):
# spark.conf.set("spark.sql.adaptive.autoBroadcastJoinThreshold", "-1")

Before/After Metrics

Before (SortMergeJoin on filtered data):

Join: large_table (10 GB) x filtered_lookup (50 MB after filter)
Strategy: SortMergeJoin (Catalyst estimated lookup at 2 GB pre-filter)
Shuffle write: 10.05 GB (both sides shuffled)
Stages: 3 (scan+filter, shuffle, join)
Total time: 180s

After (AQE converts to BroadcastHashJoin):

Same join, AQE converts at runtime
Strategy: BroadcastHashJoin
Shuffle write: 0 (broadcast eliminates shuffle!)
Stages: 2 (scan+filter, broadcast+join)
Total time: 25s

Improvement: 7.2x faster, shuffle eliminated entirely.

Feature 4: Optimizing Skewed Aggregations

The Problem

Skewed data affects aggregations too. If one group has far more rows than others:

python

# Group by country, sum revenue
# United States: 500M rows -> one task processes all of them
# Luxembourg: 10K rows -> task finishes instantly

How AQE Handles It (Databricks Enhanced)

Databricks extends AQE with partial aggregation optimization for skewed groups. The

skewed group is split across multiple tasks for partial aggregation, then combined:

Without optimization:
  Task 1: Aggregate US (500M rows) -> 45 minutes
  Task 2: Aggregate LU (10K rows) -> 0.5 seconds

With AQE skewed aggregation:
  Task 1a: Partial agg US rows 0-100M -> 9 minutes
  Task 1b: Partial agg US rows 100M-200M -> 9 minutes
  ...
  Task 1e: Partial agg US rows 400M-500M -> 9 minutes
  Task 2: Aggregate LU (10K rows) -> 0.5 seconds
  Final: Combine partial US aggregates -> 1 second
  Total: 9 minutes (instead of 45)

Configuration

python

# For Databricks Runtime, this is part of enhanced AQE
# Ensure AQE and skew handling are enabled:
spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "true")

# For additional skew handling in aggregations, consider:
spark.conf.set("spark.sql.adaptive.advisoryPartitionSizeInBytes", "64m")
# Smaller advisory size = more granular splitting of skewed partitions

AQE Configuration Reference

Complete Configuration Table

Setting	Default	Recommended	Purpose
`spark.sql.adaptive.enabled`	`true` (DBR)	`true`	Master switch
`spark.sql.adaptive.coalescePartitions.enabled`	`true`	`true`	Coalesce small partitions
`spark.sql.adaptive.coalescePartitions.initialPartitionNum`	(none)	`4096`	Starting partition count
`spark.sql.adaptive.coalescePartitions.minPartitionNum`	`1`	`1-20`	Min after coalescing
`spark.sql.adaptive.advisoryPartitionSizeInBytes`	`64m`	`128m`	Target partition size
`spark.sql.adaptive.skewJoin.enabled`	`true`	`true`	Enable skew join
`spark.sql.adaptive.skewJoin.skewedPartitionFactor`	`5`	`2-5`	Skew detection sensitivity
`spark.sql.adaptive.skewJoin.skewedPartitionThresholdInBytes`	`256m`	`64-256m`	Min size for skew
`spark.sql.adaptive.autoBroadcastJoinThreshold`	`10m`	`100m`	Runtime broadcast conversion
`spark.sql.adaptive.localShuffleReader.enabled`	`true`	`true`	Avoid remote shuffle reads
`spark.sql.adaptive.forceOptimizeSkewedJoin`	`false`	`false`	Force skew handling
`spark.sql.shuffle.partitions`	`200`	`auto` (DBR)	Shuffle partition count

Configuration by Workload Size

python

# Small workloads (< 100 GB)
small_aqe_config = {
    "spark.sql.adaptive.enabled": "true",
    "spark.sql.adaptive.advisoryPartitionSizeInBytes": "64m",
    "spark.sql.adaptive.coalescePartitions.initialPartitionNum": "1024",
    "spark.sql.adaptive.skewJoin.skewedPartitionThresholdInBytes": "64m",
    "spark.sql.adaptive.autoBroadcastJoinThreshold": "200m",
}

# Large workloads (> 1 TB)
large_aqe_config = {
    "spark.sql.adaptive.enabled": "true",
    "spark.sql.adaptive.advisoryPartitionSizeInBytes": "256m",
    "spark.sql.adaptive.coalescePartitions.initialPartitionNum": "8192",
    "spark.sql.adaptive.skewJoin.skewedPartitionThresholdInBytes": "256m",
    "spark.sql.adaptive.skewJoin.skewedPartitionFactor": "2",
    "spark.sql.adaptive.autoBroadcastJoinThreshold": "100m",
}

When AQE Cannot Help

AQE is powerful but not a silver bullet. It cannot help with:

1. First-Stage Problems

AQE optimizes at shuffle boundaries. The first stage (reading data) has no prior

shuffle, so AQE has no statistics to work with:

python

# AQE cannot help here - this is a first-stage problem:
df = spark.read.parquet("/data/skewed/")  # Skewed file sizes
# Solution: Fix the source data (repartition, compact files)

2. Single-Stage Queries

If your query has no shuffles, AQE has no opportunity to intervene:

python

# No shuffle, no AQE optimization opportunity:
result = spark.read.parquet("/data/").filter(col("x") > 10).select("y")

3. UDF Performance

AQE cannot optimize Python UDFs. The serialization overhead is inherent:

python

# AQE cannot help with this:
@udf(returnType=StringType())
def slow_udf(value):
    return complex_python_logic(value)

df.withColumn("result", slow_udf(col("input")))
# Solution: Replace UDF with built-in functions or Pandas UDFs

4. Driver-Side Bottlenecks

AQE operates on executor-side execution. It cannot fix driver issues:

python

# AQE cannot help:
data = df.collect()  # Driver OOM is driver-side, not executor-side

5. Storage-Layer Issues

AQE cannot fix poorly organized storage:

python

# AQE cannot fix:
# - 100,000 small files (< 1 MB each)
# - Lack of partition pruning on physical storage
# - Wrong file format (CSV instead of Parquet)

Debugging AQE Decisions

Viewing AQE Plans

python

# The explain output shows AQE modifications with "AdaptiveSparkPlan":
df.explain(mode="formatted")

# Look for these markers:
# == Adaptive Spark Plan ==
# AdaptiveSparkPlan isFinalPlan=true
#   +- BroadcastHashJoin  <- AQE may have converted from SortMergeJoin
#      +- CustomShuffleReader coalesced  <- AQE coalesced partitions

Spark UI Indicators

In the Databricks Spark UI SQL tab:

1. Look for "CustomShuffleReader" nodes - these indicate AQE is active

2. Check partition counts before and after CustomShuffleReader

3. "coalesced" label means partitions were combined

4. "skew" label means skew join optimization was applied

Logging AQE Decisions

python

# Enable detailed AQE logging:
spark.conf.set("spark.sql.adaptive.logLevel", "debug")

# This produces log entries like:
# Coalescing partitions: 200 -> 23 (advisory size: 128 MB)
# Detected skewed partition: partition 42, size 2.1 GB (median: 50 MB)
# Converting SortMergeJoin to BroadcastHashJoin (right side: 45 MB)

Verifying AQE Is Working

python

def check_aqe_status(spark) -> dict:
    """Check if AQE is properly configured and active."""
    settings = {
        "aqe_enabled": spark.conf.get("spark.sql.adaptive.enabled"),
        "coalesce_enabled": spark.conf.get(
            "spark.sql.adaptive.coalescePartitions.enabled"
        ),
        "skew_join_enabled": spark.conf.get(
            "spark.sql.adaptive.skewJoin.enabled"
        ),
        "advisory_partition_size": spark.conf.get(
            "spark.sql.adaptive.advisoryPartitionSizeInBytes"
        ),
        "skew_factor": spark.conf.get(
            "spark.sql.adaptive.skewJoin.skewedPartitionFactor"
        ),
        "broadcast_threshold": spark.conf.get(
            "spark.sql.adaptive.autoBroadcastJoinThreshold"
        ),
    }
    return settings

# Usage:
status = check_aqe_status(spark)
for key, value in status.items():
    print(f"{key}: {value}")

Code Examples

Example 1: AQE Handling Skewed Join

python

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, rand, when, lit, sum as spark_sum

spark = SparkSession.builder.getOrCreate()

# Configure AQE for skew handling
spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "true")
spark.conf.set("spark.sql.adaptive.skewJoin.skewedPartitionFactor", "3")
spark.conf.set("spark.sql.adaptive.skewJoin.skewedPartitionThresholdInBytes", "64m")

# Generate skewed test data
orders = (
    spark.range(10_000_000)
    .withColumn(
        "customer_id",
        when(rand() < 0.3, lit(1))  # 30% of orders from customer 1 (skewed!)
        .otherwise((rand() * 100000).cast("int"))
    )
    .withColumn("amount", (rand() * 1000).cast("decimal(10,2)"))
)

customers = (
    spark.range(100_000)
    .withColumnRenamed("id", "customer_id")
    .withColumn("region", when(col("customer_id") % 4 == 0, "North")
                .when(col("customer_id") % 4 == 1, "South")
                .when(col("customer_id") % 4 == 2, "East")
                .otherwise("West"))
)

# Join - AQE will detect the skew on customer_id = 1
result = (
    orders.join(customers, "customer_id")
    .groupBy("region")
    .agg(spark_sum("amount").alias("total_revenue"))
)

# Check the plan - look for "skew" markers
result.explain(mode="formatted")

# Execute
result.show()

Example 2: AQE Coalescing After Aggressive Repartition

python

# Intentionally over-partition to let AQE coalesce
spark.conf.set("spark.sql.shuffle.partitions", "2000")
spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.adaptive.coalescePartitions.enabled", "true")
spark.conf.set("spark.sql.adaptive.advisoryPartitionSizeInBytes", "128m")

# This aggregation produces small output (~10 MB)
result = (
    spark.read.table("catalog.schema.large_table")
    .groupBy("category")
    .agg(
        spark_sum("value").alias("total"),
        count("*").alias("cnt")
    )
)

# Without AQE: 2000 tasks for 10 MB of output
# With AQE: coalesced to ~1-2 tasks
result.write.format("delta").mode("overwrite").save("/tmp/output")

# In the Spark UI, you'll see:
# Exchange (shuffle write): 2000 partitions
# CustomShuffleReader: coalesced to 2 partitions

Example 3: Comparing AQE On vs Off

python

import time
from pyspark.sql.functions import col, sum as spark_sum, count

def benchmark_with_aqe(spark, enabled: bool) -> dict:
    """Run a benchmark with AQE enabled or disabled.

    Args:
        spark: Active SparkSession
        enabled: Whether to enable AQE

    Returns:
        Dictionary with timing and plan information
    """
    spark.conf.set("spark.sql.adaptive.enabled", str(enabled).lower())

    # Clear cache
    spark.catalog.clearCache()

    start = time.time()

    result = (
        spark.read.table("catalog.schema.orders")
        .filter(col("order_date") >= "2025-01-01")
        .join(
            spark.read.table("catalog.schema.products"),
            "product_id"
        )
        .groupBy("category", "region")
        .agg(
            spark_sum("amount").alias("total"),
            count("*").alias("cnt")
        )
    )

    # Force execution
    result.write.format("noop").mode("overwrite").save()
    elapsed = time.time() - start

    return {
        "aqe_enabled": enabled,
        "elapsed_seconds": round(elapsed, 2),
        "plan": result._jdf.queryExecution().executedPlan().toString(),
    }


# Run both benchmarks
result_without = benchmark_with_aqe(spark, enabled=False)
result_with = benchmark_with_aqe(spark, enabled=True)

print(f"Without AQE: {result_without['elapsed_seconds']}s")
print(f"With AQE:    {result_with['elapsed_seconds']}s")
print(f"Speedup:     {result_without['elapsed_seconds'] / result_with['elapsed_seconds']:.1f}x")

Key Takeaways

1. Always enable AQE. There is almost no downside and significant upside.

2. Set shuffle.partitions to "auto" on Databricks or a high value (4096+) with AQE.

3. AQE handles skew automatically - but tune the sensitivity for severe skew.

4. AQE cannot fix first-stage problems. Fix input data organization separately.

5. Check the SQL tab for AQE markers: CustomShuffleReader, coalesced, skew.

6. AQE gets better with more shuffles - complex queries benefit the most.

7. Combine AQE with proper table statistics for the best results.

Next Chapter

Chapter 3: Shuffle Optimization - Deep dive into the

shuffle operation that AQE optimizes around, and how to minimize shuffles in the

first place.

Chapter 3

🔒 Available in full product

Chapter 3: Shuffle Optimization

Chapter 4

🔒 Available in full product

Chapter 4: Memory Tuning

Chapter 5

🔒 Available in full product

Chapter 5: Delta Lake Optimization

Chapter 6

🔒 Available in full product

Chapter 6: Photon Engine

Chapter 7

🔒 Available in full product

Chapter 7: Join Strategies

Chapter 8

🔒 Available in full product

Chapter 8: I/O Optimization

Chapter 9

🔒 Available in full product

Chapter 9: Streaming Performance

Chapter 10

🔒 Available in full product

Chapter 10: Troubleshooting Runbook

You’ve reached the end of the free preview

Get the full Spark Performance Masterclass and unlock everything.

All Chapters

Get the complete guide with every chapter unlocked, including code samples, diagrams, and best practices.

Full Tool Suite

Access all interactive tools with complete data, all workload profiles, and the full scenario library.

Source Files

Downloadable source code, configuration files, and working examples from every chapter.

Lifetime Updates

Free updates for life. Every new chapter, tool, and improvement included.

Buy Now — 9 →
📦 Free sample included — download another copy for the full product.

Spark Performance Masterclass v1.0.0 — Free Preview

Contents

Chapter 1: The Spark Execution Model

Table of Contents

The Big Picture

Driver and Executors

The Driver

Executors

The Relationship

From Code to Execution

Step 1: Unresolved Logical Plan

Step 2: Analysis (Catalog Resolution)

Step 3: Logical Optimization (Catalyst)

Step 4: Physical Planning

Step 5: Code Generation

Jobs, Stages, and Tasks

The Hierarchy

Jobs

Stages

Tasks

Narrow vs Wide Transformations

Narrow Transformations (No Shuffle)

Wide Transformations (Shuffle Required)

Performance Impact

The Shuffle: Spark's Most Expensive Operation

What Happens During a Shuffle

Shuffle Metrics to Monitor

Configuration

Partitions: The Unit of Parallelism

Input Partitions

Shuffle Partitions

Repartitioning Strategies

Before/After: Partition Sizing

Lazy Evaluation and the DAG

How Lazy Evaluation Works

Why This Matters for Performance

Breaking Lazy Evaluation (Performance Anti-Patterns)

Configuration Settings

Essential Settings for Execution Model Tuning

Code Examples

Example 1: Understanding Your Execution Plan

Example 2: Monitoring Execution Metrics Programmatically

Example 3: Minimizing Stage Count

Key Takeaways

Next Chapter

Chapter 2: Adaptive Query Execution (AQE) Deep Dive

Table of Contents

What Is AQE?

The Core Insight

AQE in Databricks vs Open Source

How AQE Works Under the Hood

The Re-Optimization Loop

What Statistics AQE Collects

Enabling AQE

Feature 1: Coalescing Post-Shuffle Partitions

The Problem

How AQE Solves It

Configuration

Before/After Metrics

Feature 2: Skew Join Optimization

The Problem

How AQE Solves It

Configuration

Before/After Metrics

Feature 3: Converting Sort-Merge to Broadcast Joins

The Problem

How AQE Solves It

Configuration

Before/After Metrics

Feature 4: Optimizing Skewed Aggregations

The Problem

How AQE Handles It (Databricks Enhanced)

Configuration

AQE Configuration Reference

Complete Configuration Table

Configuration by Workload Size

When AQE Cannot Help

1. First-Stage Problems

2. Single-Stage Queries

3. UDF Performance

4. Driver-Side Bottlenecks