The Spark Execution Model: Jobs, Stages, and Caching

Spark’s API hides the execution model well enough that you can write jobs without knowing it, right up until performance matters. This lesson opens it up.

Four layers

Your DataFrame code goes through four representations before running:

1. Unresolved logical plan — parsed, column names not yet verified
2. Analysed logical plan — names and types resolved against the catalog
3. Optimised logical plan — Catalyst has applied its rewrite rules
4. Physical plan — concrete operators, join strategies, exchanges

explain(True) shows all four:

from pyspark.sql import SparkSession, functions as F

spark = SparkSession.builder.appName("execution").master("local[*]").getOrCreate()

df = spark.range(1_000_000).select(
    F.col("id"),
    (F.col("id") % 100).alias("bucket"),
    (F.rand() * 1000).alias("amount"),
)

q = df.filter(F.col("amount") > 500).filter(F.col("bucket") < 10).select("id", "amount")
q.explain(True)

== Parsed Logical Plan ==
'Project [unresolvedalias('id), unresolvedalias('amount)]
+- Filter (bucket#2L < cast(10 as bigint))
   +- Filter (amount#3 > cast(500 as double))
      +- Project [id#0L, (id#0L % 100) AS bucket#2L, (rand(...) * 1000.0) AS amount#3]
         +- Range (0, 1000000, step=1, splits=8)

== Analyzed Logical Plan ==
id: bigint, amount: double
Project [id#0L, amount#3]
+- Filter (bucket#2L < 10)
   ...

== Optimized Logical Plan ==
Project [id#0L, amount#3]
+- Filter (((id#0L % 100) < 10) AND ((rand(...) * 1000.0) > 500.0))
   +- Range (0, 1000000, step=1, splits=8)

== Physical Plan ==
*(1) Project [id#0L, amount#3]
+- *(1) Filter (((id#0L % 100) < 10) AND ((rand(...) * 1000.0) > 500.0))
   +- *(1) Range (0, 1000000, step=1, splits=8)

The optimised plan is where the work shows. Two Filter nodes became one. The Project that created bucket disappeared entirely — Catalyst inlined id % 100 into the filter, so the intermediate column is never materialised. You wrote four operations; Spark will run two.

Narrow and wide

The *(1) prefix marks whole-stage code generation — Spark compiled those operators into a single Java method. Everything with the same number runs in one stage, in one pass over the data, with no intermediate objects.

Add a groupBy and the picture changes:

df.groupBy("bucket").agg(F.sum("amount")).explain()

== Physical Plan ==
AdaptiveSparkPlan isFinalPlan=false
+- HashAggregate(keys=[bucket#2L], functions=[sum(amount#3)])
   +- Exchange hashpartitioning(bucket#2L, 200), ENSURE_REQUIREMENTS
      +- HashAggregate(keys=[bucket#2L], functions=[partial_sum(amount#3)])
         +- Project [(id#0L % 100) AS bucket#2L, (rand(...) * 1000.0) AS amount#3]
            +- Range (0, 1000000, step=1, splits=8)

Exchange is a stage boundary. Everything below it is stage 0; everything above is stage

1. Stage 0 writes its output to local disk as shuffle files, stage 1 fetches them over the network.

That is the whole rule:

Counting Exchange nodes in a plan tells you how many times your data crosses the network, which is usually the number that decides runtime.

Jobs, stages, tasks

result = df.groupBy("bucket").agg(F.sum("amount")).collect()

One action triggers one job. The job has two stages because of the one shuffle. Each stage runs one task per partition — 8 tasks in stage 0 (the Range had 8 splits), then 200 in stage 1 before AQE coalesces them.

Check the partition counts directly:

print("input partitions: ", df.rdd.getNumPartitions())
print("after groupBy:    ", df.groupBy("bucket").agg(F.sum("amount")).rdd.getNumPartitions())
print("after coalesce:   ", df.coalesce(2).rdd.getNumPartitions())
print("after repartition:", df.repartition(16).rdd.getNumPartitions())

input partitions:  8
after groupBy:     200
after coalesce:    2
after repartition: 16

coalesce merges partitions without a shuffle and can only reduce. repartition shuffles and can do either, and gives you even sizes. Use coalesce before a write to cut file count; use repartition when partitions are unevenly sized.

Lineage and recomputation

Spark stores a recipe, not results. Two actions mean two full computations:

import time

expensive = df.filter(F.col("amount") > 100).groupBy("bucket").agg(F.avg("amount").alias("avg"))

start = time.perf_counter(); expensive.count();          t1 = time.perf_counter() - start
start = time.perf_counter(); expensive.collect();        t2 = time.perf_counter() - start
print(f"first action:  {t1:.2f}s")
print(f"second action: {t2:.2f}s")

first action:  1.84s
second action: 1.79s

No speed-up. The second action re-ran the scan, the filter, and the shuffle from scratch.

Cache it:

expensive.cache()

start = time.perf_counter(); expensive.count();   t1 = time.perf_counter() - start
start = time.perf_counter(); expensive.collect(); t2 = time.perf_counter() - start
print(f"first action (populates cache): {t1:.2f}s")
print(f"second action (from cache):     {t2:.2f}s")

first action (populates cache): 1.97s
second action (from cache):     0.04s

Forty times faster. The first action paid a little extra to write the cache; every action after reads it.

Caching changes the plan, not just the timing:

expensive.explain()

== Physical Plan ==
InMemoryTableScan [bucket#2L, avg#31]
   +- InMemoryRelation [bucket#2L, avg#31], StorageLevel(disk, memory, deserialized, 1 replicas)
         +- *(3) HashAggregate(keys=[bucket#2L], functions=[avg(amount#3)])
            +- ...

InMemoryTableScan replaced the whole subtree. The computation below is only run once, to populate the relation.

Cache is lazy too

other = df.filter(F.col("bucket") == 5)
other.cache()
print(spark.catalog.isCached("..."), "— nothing computed yet")

cache() marks a DataFrame for caching but computes nothing. Until an action runs, the cache is empty. This trips people up when they cache and immediately check the Storage tab and find it empty. Force it:

other.cache().count()

Free it when done — cached data holds executor memory that other work needs:

expensive.unpersist()

When caching hurts

cheap = df.select("id", "amount")

start = time.perf_counter(); cheap.count(); cheap.count()
print(f"uncached, two counts: {time.perf_counter() - start:.2f}s")

cheap.cache()
start = time.perf_counter(); cheap.count(); cheap.count()
print(f"cached,   two counts: {time.perf_counter() - start:.2f}s")
cheap.unpersist()

uncached, two counts: 0.34s
cached,   two counts: 1.12s

Slower. Serialising a million rows into the cache cost more than recomputing a trivial projection twice. Cache when the work below is expensive — after a shuffle, a wide join, or a UDF — not reflexively.

Non-determinism without caching

r = df.select("id", F.rand().alias("r")).filter(F.col("r") > 0.5)
print("count 1:", r.count())
print("count 2:", r.count())

count 1: 500204
count 2: 499876

Different answers from the same DataFrame. rand() is re-evaluated on each action because the lineage is recomputed. Anything downstream of a non-deterministic expression must be cached — or given a seed — before you use it twice.

r = df.select("id", F.rand(seed=42).alias("r")).filter(F.col("r") > 0.5).cache()
print("count 1:", r.count())
print("count 2:", r.count())

count 1: 500118
count 2: 500118

Practice

coalesce(1):    2.31s
repartition(1): 0.88s

print(expensive.storageLevel)

Disk Memory Deserialized 1x Replicated

0, 1, 2, 8589934592, 8589934593, 17179869184

Next: turning all of this into a tuning checklist.