🔥 Spark & Databricks · Chapter 1

Spark: jobs, stages, and the shuffle that costs you everything

Spark is fast until it shuffles. Almost every Spark performance problem is one question: how many times does the data cross the network?

Study this lesson free in the app →

The basics

Apache Spark is an engine that runs one piece of data-processing code on many machines at once, each machine working on its own slice of the data.

The problem it solves: a day of ride events for a city like Bengaluru can be 400 GB. One laptop would take hours just to read it. Spark splits it into partitions (say 3,200 chunks of ~128 MB) and hands each to a worker core, so 100 cores finish in minutes.

The mental model: think of an exam hall with 100 invigilators. Checking that every answer sheet has a name on it needs no coordination — each invigilator checks their own pile. But "total marks per school" forces everyone to walk across the hall and regroup sheets by school. That regrouping is the shuffle, and it is where Spark jobs spend their time.

Explained

Your job takes 4 minutes. You add one groupBy and it takes 40.

You did not add much computation. You added a shuffle — and a shuffle means every row is written to disk, sent across the network, and read again.

Spark splits data into partitions and runs the same code on each, on different machines.

Any operation that only needs the rows in front of it — filter, select, withColumn — is narrow. No data moves. These are almost free.

Any operation that needs rows from other partitions — groupBy, join, distinct, orderBy — is wide. Rows must be redistributed so that everything with the same key lands together. That redistribution is the shuffle: write to disk, transfer over the network, read back.

Narrow is cheap. Wide is the bill.

How the words fit together

Transformation — lazy. Building a plan, running nothing. df.filter(...) does no work.

Action — count(), write(), collect(), show(). This is what triggers execution.

Job — one action.

Stage — a run of narrow operations that can be pipelined together. A new stage begins at every shuffle. So counting stages in the UI tells you how many shuffles you have.

Task — one stage on one partition. The unit that is actually scheduled.

Laziness is why an error in line 3 surfaces at line 40: nothing ran until the action.

Reading the plan before you run it

df = spark.read.parquet('s3://lake/rides/')       # narrow
     .filter(col('ride_date') == '2026-03-14')     # narrow — pushed to the file scan
     .withColumn('fare_inr', col('fare') * 1.0)    # narrow

out = (df.groupBy('city').agg(sum('fare_inr'))    # WIDE — shuffle 1
         .join(cities, 'city')                    # WIDE — shuffle 2
         .orderBy(desc('sum(fare_inr)')))         # WIDE — shuffle 3

out.explain()      # read this BEFORE .write()

Three shuffles. explain() shows them as Exchange — and every Exchange is a full round trip of your data.

The join is often avoidable: if cities is small, broadcast it. Spark sends the small side to every executor and the shuffle disappears.

Broadcast the small side. Under ~10MB (tunable), broadcast(cities) turns a shuffle join into a map-side join. Single biggest win available in most jobs.

Filter before the shuffle, always. Shuffling rows you are about to discard is pure waste, and it is astonishing how often the filter sits after the join.

Partition count. Default spark.sql.shuffle.partitions is 200. On a small job that is 200 tiny tasks with scheduling overhead exceeding the work. On a huge one it is 200 enormous tasks that spill to disk. Adaptive Query Execution fixes much of this automatically in Spark 3 — turn it on.

Cache only what is reused. cache() on a DataFrame used once is pure cost — it materialises and holds memory for nothing.

A job that ran in 6 minutes now takes 90.

The data grew 20%. Nothing else changed.

In the Spark UI, one stage has 199 tasks finishing in seconds and one task running for 80 minutes. That is not a capacity problem — adding executors will not help, because 199 of them are idle.

It is skew: one key holds most of the rows. Usually null, or a default like 'unknown', or one enormous customer.

The fix is not a bigger cluster. It is to filter the nulls out before the join, or salt the hot key so it spreads across partitions.

collect() on anything large. Pulls the whole dataset into the driver and kills it. Use take(20) or write to storage.

A filter placed after the join rather than before it.

cache() everywhere. Memory pressure, spilling, and slower than not caching.

Assuming more executors fixes it. If the problem is skew, 199 of them will sit idle.

Ignoring explain(). The plan is right there, before you spend forty minutes finding out.

Narrow is free. Wide shuffles. A new stage means a new shuffle.

Before optimising anything: run explain(), count the Exchange nodes, and ask which one you can delete.

Filter early. Broadcast the small side. Never collect().

Words to know

In two minutes

• Spark splits data into partitions and runs one task per partition, in parallel across executors.
• Transformations are lazy: they build a plan. Only an action (count, write, collect) executes it.
• Narrow operations (filter, select, withColumn) stay inside a partition and are nearly free.
• Wide operations (groupBy, join, distinct, orderBy) shuffle: write to disk, send over network, read back.
• Every shuffle is a stage boundary, so counting stages or Exchange nodes in explain() counts your cost.
• The big wins: filter before the shuffle, broadcast the small side of a join, size shuffle partitions (or let AQE do it).
• One slow task among many fast ones is skew — fix the hot key, not the cluster size.

Test yourself

There are 8 questions, 3 puzzles and 6 recall cards for this lesson in the app — the recall cards come back on later days, just before you would forget.