PySpark Transformations vs Actions

PySpark code looks like it runs line by line. It does not. Transformations only describe a new DataFrame; actions make Spark optimise and execute the whole chain. Almost every confusing PySpark behaviour — a “fast” line that is actually slow later, an error that appears far from the code that caused it, a job that recomputes everything twice — follows from this one rule.

Laziness in practice

from pyspark.sql import SparkSession, functions as F

spark = (SparkSession.builder.master("local[2]").appName("lazy")
         .config("spark.sql.shuffle.partitions", "4").getOrCreate())

events = spark.range(0, 1_000_000).withColumn("bucket", F.col("id") % 10)   # transformation
big = events.filter(F.col("id") > 500_000)                                   # transformation
counts = big.groupBy("bucket").count()                                        # transformation

print(type(counts).__name__)                 # still just a plan
print(counts.orderBy("bucket").first())      # action: now Spark runs a job

DataFrame
Row(bucket=0, count=49999)

Until first(), nothing was computed: the three assignments returned instantly because they only built a logical plan. When the action arrived, Spark could see the whole chain and optimise it as one unit — the filter runs before the grouping, only the columns that are needed are produced, and partial counts are computed before any data moves.

Common transformations and actions

Transformations (lazy, return a DataFrame)	Actions (run a job, return a value or write data)
select, withColumn, filter, drop	show, count, collect, first, take, head
join, groupBy().agg(), distinct, dropDuplicates	write...save(), saveAsTable, insertInto
orderBy, repartition, coalesce, union	toPandas, toLocalIterator, foreach, foreachPartition
cache, persist (they only mark the DataFrame)	Starting a streaming query with start()

A quick test: if the method returns a DataFrame, it is a transformation. If it returns rows, a number, None after writing, or a pandas object, it is an action.

Errors also wait for an action

Because nothing runs until an action, analysis errors (a missing column, a type mismatch Spark can see in the schema) are raised immediately, but data errors are raised only when a job processes the bad value:

bad = (events
       .withColumn("code", F.concat(F.lit("a"), F.col("id").cast("string")))
       .withColumn("code_num", F.col("code").cast("int")))   # every value is invalid
print("defined, no error yet")

print(bad.count())          # code_num is never needed, so Spark prunes it away

try:
    bad.select("code_num").first()
except Exception as e:
    print(type(e).__name__, "-", str(e).split("]")[0] + "]")

defined, no error yet
1000000
NumberFormatException - [CAST_INVALID_INPUT]

Two lessons here. First, the failing line in a stack trace is the action, not the cast that caused it; PySpark 4 helps by naming the DataFrame call site in the error message. Second, count() succeeded because the optimiser removed the column it did not need — so a “smoke test” with count() does not prove a pipeline is correct. Under Spark 3.x defaults (ANSI off) the cast would silently return NULL instead.

Narrow versus wide

• Narrow: each output partition depends on one input partition (filter, select, withColumn, coalesce, union). Spark pipelines consecutive narrow steps into one task, and no data moves between executors.
• Wide: an output partition needs data from many input partitions (groupBy, join without broadcast, distinct, orderBy, repartition, window functions with partitionBy). Spark must shuffle: write data to local disk partitioned by key, then have the next stage fetch it across the network. Each shuffle is a stage boundary.

You can see the shuffle in the plan as an Exchange:

counts.explain()

== Physical Plan ==
AdaptiveSparkPlan isFinalPlan=false
+- HashAggregate(keys=[bucket#1L], functions=[count(1)])
   +- Exchange hashpartitioning(bucket#1L, 4), ENSURE_REQUIREMENTS, [plan_id=72]
      +- HashAggregate(keys=[bucket#1L], functions=[partial_count(1)])
         +- Project [(id#0L % 10) AS bucket#1L]
            +- Filter (id#0L > 500000)
               +- Range (0, 1000000, step=1, splits=2)

Read it bottom up. Range → Filter → Project → partial HashAggregate is one narrow pipeline (stage 1). Exchange hashpartitioning(bucket, 4) is the shuffle into 4 partitions. The final HashAggregate is stage 2. Spark pre-aggregates on each partition before shuffling (partial_count), so only one row per bucket per input partition crosses the network instead of 500,000 raw rows.

AdaptiveSparkPlan isFinalPlan=false means Adaptive Query Execution (on by default since Spark 3.2) may still change this plan after it sees real shuffle sizes. Here it does:

print(events.rdd.getNumPartitions(), big.rdd.getNumPartitions(), counts.rdd.getNumPartitions())

2 2 1

The narrow steps keep the 2 input partitions. The shuffle was planned with 4 partitions, but AQE saw ten tiny groups and coalesced them into 1 at run time.

Why repeated actions can be slow

Each action re-executes the plan from the source unless results are cached. These two lines each run range → filter → groupBy again:

counts.count()
counts.show(3)

When an expensive DataFrame feeds several actions, cache it. cache() is itself lazy — the first action fills the cache, later actions read from it:

counts.cache()
print(counts.storageLevel)
counts.count()                      # materialises the cache
counts.orderBy("bucket").explain()  # now reads InMemoryTableScan

Disk Memory Deserialized 1x Replicated
== Physical Plan ==
AdaptiveSparkPlan isFinalPlan=false
+- Sort [bucket#1L ASC NULLS FIRST], true, 0
   +- Exchange rangepartitioning(bucket#1L ASC NULLS FIRST, 4), ENSURE_REQUIREMENTS, [plan_id=97]
      +- InMemoryTableScan [bucket#1L, count#2L]
            +- InMemoryRelation [bucket#1L, count#2L], StorageLevel(disk, memory, deserialized, 1 replicas)
               ...

counts.unpersist()

The DataFrame default level is memory-and-disk, so partitions that do not fit spill to local disk instead of being dropped. Do not cache everything: memory is shared with execution, cached data that is used once is pure overhead, and a stale cache hides changes in the source. The caching lesson covers storage levels and when caching hurts.

Long lineages

Each transformation adds to the plan. In a loop that keeps reassigning df = df.withColumn(...) hundreds of times, or an iterative algorithm that joins a DataFrame with itself repeatedly, the plan itself becomes large and slow to analyse. Prefer one select with many expressions (or withColumns with a dict) over hundreds of withColumn calls, and cut very long lineages by writing an intermediate result to a table, or with df.checkpoint() / df.localCheckpoint().

Common mistakes

1. Timing a transformation line and concluding it is fast (nothing ran).
2. Calling count() “just to check” between every step in production code, triggering a full job each time — and, as shown above, not even evaluating every column.
3. Expecting cache() to do anything before an action, or caching a DataFrame used only once.
4. collect() on large data.
5. Reading a stack trace from the action line and missing the transformation that caused it.

Interview angle

This is one of the most common Spark interview openers. A strong answer covers four things: transformations are lazy and return a new DataFrame; actions trigger a job; laziness lets Catalyst optimise the whole chain (predicate pushdown, column pruning, pre-aggregation); and wide transformations shuffle and create stage boundaries. See the interview answers for transformation vs action and what causes a shuffle.

Is cache() a transformation or an action?

Neither in the strict sense: it marks the DataFrame for caching and returns it, so it is lazy like a transformation. The data is stored only when the next action computes it.

Why did count() succeed on a DataFrame whose column fails when selected?

The optimiser prunes columns that the action does not need. count() needs no columns, so the failing cast was never evaluated. Validate with the columns you actually write.

Key takeaway

Transformations build a plan; actions run it. Know which of your transformations are wide, because those are where Spark shuffles — and remember that errors and costs surface at the action, not at the line that caused them.