Spark Stages and Tasks — Deep Dive Into the Unit of Parallelism

Job, stage, task — the three levels Spark actually schedules

Nothing runs until an action fires. Transformations build a plan; count(), collect(), save() and friends submit a job. One action usually means one job, but not always: a write with a sort, a show() that has to widen its limit, or CSV/JSON schema inference all launch extra jobs, which is why the Jobs tab often shows more rows than you have actions in your code.

The DAGScheduler takes the job's final RDD or physical plan and walks backward, cutting the graph into stages at every shuffle. Two kinds come out: a ShuffleMapStage, whose output is shuffle files consumed by a later stage, and exactly one ResultStage, whose output goes back to the driver or to storage. Stages are submitted in dependency order — a stage becomes runnable only when every stage producing its shuffle input has finished.

Each submitted stage becomes a TaskSet handed to the TaskSchedulerImpl, which wraps it in a TaskSetManager and hands individual tasks to executors as resource offers arrive. That is the whole ladder: an action makes a job, the DAG scheduler makes stages, and the task scheduler makes tasks. Stage numbers are global and monotonic across the application, so stage 47 is not "the 47th stage of this job" — a detail that confuses everyone reading the UI for the first time.

One more artifact you will see: skipped stages, shown greyed out. Those are stages whose shuffle output already exists from an earlier job, so Spark reuses the map output instead of recomputing it.

Where stage boundaries fall: narrow versus wide dependencies

A stage is a maximal run of transformations that need no data movement. The dividing line is the dependency type. A narrow dependency means each parent partition feeds at most one child partition — map, filter, mapPartitions, union, and a join whose sides are already co-partitioned. Narrow chains are pipelined: one task streams a partition through every operator in the chain without materializing anything in between, which is also what makes whole-stage codegen possible.

A wide dependency means a child partition draws from many parent partitions — groupByKey, reduceByKey, a hash join, repartition, a window without a matching existing partitioning. Rows must be redistributed by key, so the parent stage writes shuffle files to local disk and stops. That is the boundary. Everything expensive about Spark lives at these lines: network transfer, disk I/O, a hard synchronization barrier, and the place where a lost executor costs you recomputation.

Practically, you find your boundaries by reading df.explain() and counting Exchange nodes, or by counting stages in the UI's DAG visualization. n exchanges means n + 1 stages on that path. The dependency graph itself, and how it doubles as the fault-tolerance mechanism, is covered in the RDD lineage article; what happens inside the boundary — the sort, the spill, the fetch — is the shuffle article. Here we care only that the boundary is where one set of tasks ends and the next begins.

One task is one partition of one stage

The rule is exact and worth memorizing: a task computes exactly one partition of exactly one stage. A stage with 200 output partitions runs 200 tasks. There is no other source of parallelism inside a stage — you cannot make a stage faster by adding cores past its task count, and a stage with four partitions will use four cores on a thousand-core cluster.

Tasks come in the same two flavours as stages. A ShuffleMapTask runs the stage's pipelined operators over its partition and writes the result into shuffle files bucketed by destination partition. A ResultTask runs the operators and then applies the action's function, returning a value to the driver. Both are shipped to executors as a serialized task binary that the driver broadcasts once per stage, plus a small per-task description carrying the partition id and its preferred locations.

Each task occupies one slot — spark.task.cpus cores, default 1 — for its whole duration, and it is single-threaded from Spark's point of view. That gives you the two failure modes at the extremes. Too few partitions and each task must hold a large working set in one executor's memory share, so you spill or OOM while most of the cluster sits idle. Too many and per-task overhead dominates: serialization, launch, result handling, and one output file per task at the write. Spark's own guidance is that a task should run for at least a couple of hundred milliseconds; if your median task duration is 30 ms, you are paying more to schedule work than to do it.

Where the partition count comes from

Since task count equals partition count, tuning parallelism is really tuning three separate mechanisms that people routinely confuse.

Read side. For file sources, the planner packs file splits into partitions of at most spark.sql.files.maxPartitionBytes (128 MB default), charging spark.sql.files.openCostInBytes (4 MB default) per file so that thousands of tiny files do not become thousands of tiny tasks. spark.sql.files.minPartitionNum puts a floor under the result. Splittability matters: Parquet and ORC split at row-group and stripe boundaries, plain text splits anywhere, but a gzip file is not splittable at all — one file, one task, however large. In the RDD API the equivalent knob is the minPartitions argument to sc.textFile.

Shuffle side. This is where the two most-confused keys live. spark.sql.shuffle.partitions (default 200) sets the reduce-side partition count for DataFrame and SQL shuffles. spark.default.parallelism governs the RDD API — reduceByKey without an explicit count, parallelize — and defaults to the cluster's total core count. Setting the RDD key and expecting your SQL joins to change is a classic wasted afternoon.

Explicit. repartition(n) forces a shuffle and gives you exactly n partitions. coalesce(n) is narrow — it merges partitions without a shuffle, which is cheap but has a nasty property covered below.

Slots, waves, and the cost of a ragged last wave

Available parallelism is executors × executor.cores ÷ spark.task.cpus. Call that the slot count. A stage with more tasks than slots runs in waves: 600 tasks on 200 slots is three waves. Wave arithmetic explains a lot of otherwise mysterious timing.

The failure case is the ragged tail. 210 tasks on 200 slots is not "5% more work" than 200 tasks — it is two waves, and during the second one 190 slots idle while 10 tasks finish. Wall clock nearly doubles for a 5% increase in work. The same effect at the opposite end is why a stage with 8 partitions on a 200-slot cluster is not eight times slower than optimal, it is twenty-five times slower.

The standard guidance — two to four tasks per core — exists precisely to smear this out. With many small tasks per slot, an uneven task duration distribution averages away, and a slow executor simply picks up fewer tasks. With exactly one task per slot, every duration outlier lands directly on the critical path.

The coalesce trap belongs here. coalesce(1) before a write does not add a shuffle boundary, so the reduced parallelism propagates upstream through the whole narrow chain: your 400-task scan-and-filter stage becomes a 1-task stage that reads everything on one core. If you want one output file but full read parallelism, use repartition(1) and accept the shuffle, or leave the partitioning alone and control file size at the writer.

Locality levels and delay scheduling

The TaskSetManager does not hand tasks out purely first-come. Each task carries preferred locations — the hosts holding its input block or its cached partition — and Spark tries to honour them through five locality levels, best to worst: PROCESS_LOCAL (the data is in this executor's memory), NODE_LOCAL (on this host, in HDFS or another executor), NO_PREF (no preference at all), RACK_LOCAL, and ANY.

When no slot at the desired level is free, Spark uses delay scheduling: it waits briefly for a better-placed slot rather than immediately launching remotely. spark.locality.wait (3s default) sets that patience, with per-level overrides spark.locality.wait.process, .node and .rack. The level actually achieved is printed per task in the Stages tab, and the aggregate shows up as a "Locality Level Summary" line on the stage page.

Two operational notes. On HDFS with cached RDDs, the wait usually pays for itself. On object storage there is no meaningful node locality — tasks are NO_PREF and the wait never engages — so tuning it there is cargo cult. But a stage reading a cached DataFrame on a busy cluster can spend seconds per task waiting for the executor that holds the block; if scheduler delay is high and locality levels are all PROCESS_LOCAL, dropping spark.locality.wait to 0s trades a little remote reading for a lot less waiting.

Which job's tasks win a free slot in the first place is a separate question, governed by spark.scheduler.mode — FIFO by default, or FAIR with named pools set via the spark.scheduler.pool local property.

Task retry versus stage retry

These are different mechanisms with different limits, and conflating them makes failure logs unreadable.

Task retry is the TaskSetManager's job. An ordinary task failure — an exception in your code, an executor OOM, a lost container — causes that single task to be resubmitted, preferably somewhere else. The counter is per task and the limit is spark.task.maxFailures, default 4. The fourth failure of the same task aborts the stage and fails the job, which is why the driver log says "Task 137 in stage 9.0 failed 4 times" rather than naming the stage.

Stage retry is the DAGScheduler's job and is triggered almost exclusively by FetchFailedException: a reduce task cannot read map output because the executor or host that wrote it is gone. Retrying the reduce task is pointless — the data does not exist. So Spark unregisters that executor's map outputs, resubmits the missing tasks of the parent ShuffleMapStage in a new stage attempt, and then re-runs the blocked reducers. Fetch failures do not count toward spark.task.maxFailures; they count toward spark.stage.maxConsecutiveAttempts, default 4.

In the UI a retried stage shows as a second attempt with its own task table, and in-flight tasks from the old attempt become "zombie" tasks that finish but are ignored. The structural fix for repeated fetch failures is to stop losing shuffle files with the executor: enable the external shuffle service, which is also what makes aggressive dynamic allocation safe.

Speculative execution and when it backfires

Speculation is Spark's answer to the slow machine, not to the big partition. With spark.speculation=true the scheduler samples running tasks every spark.speculation.interval (100 ms) and, once spark.speculation.quantile (0.75) of a stage's tasks have finished, launches a duplicate copy of any task running longer than spark.speculation.multiplier (1.5) times the median of the completed ones. spark.speculation.minTaskRuntime suppresses this for tasks too short to be worth duplicating, and spark.speculation.task.duration.threshold lets a task that blows an absolute time budget be speculated before the quantile is reached. Whichever copy finishes first wins; the other is killed and shows in the UI as a killed task with reason "another attempt succeeded".

It is off by default for good reasons. Against data skew it is actively harmful: the duplicate reads the same 2 GB partition and is exactly as slow, so you have burned a slot for nothing — that is AQE's skew splitter's job, not speculation's. Against non-idempotent work it is dangerous: the output commit coordinator prevents two attempts from both committing files, but nothing protects a foreachPartition that posts to an external API. And it is outright incompatible with barrier execution mode, where all tasks must run concurrently.

Turn it on when your fleet is heterogeneous or preemptible and you have observed single-node stragglers with even input sizes. That is the case it was built for.

What AQE changes about your task counts

Adaptive query execution rewrites the second half of everything above, because it changes task counts after you set the config. The enabling fact is that a shuffle is a full materialization barrier: when the map side finishes, the exact byte count of every reduce partition is known, not estimated.

Coalescing is the one that matters here. With spark.sql.adaptive.coalescePartitions.enabled, contiguous small reduce partitions are merged toward spark.sql.adaptive.advisoryPartitionSizeInBytes (64 MB), so your configured 200 becomes however many tasks the data justifies. Note the asymmetry: coalescing acts on the reduce side only — map-side task count is fixed by the input layout and AQE cannot touch it. Note also spark.sql.adaptive.coalescePartitions.parallelismFirst, which defaults to true and quietly overrides your advisory size in favour of keeping cores busy; batch jobs that care about output file size usually want it set to false.

Skew splitting goes the other way, subdividing any partition larger than spark.sql.adaptive.skewJoin.skewedPartitionFactor (5) times the median and above spark.sql.adaptive.skewJoin.skewedPartitionThresholdInBytes (256 MB), duplicating the matching side to preserve join semantics — so one stage can end up with more tasks than partitions you asked for.

The practical consequence: with AQE on, spark.sql.shuffle.partitions stops being a decision and becomes a ceiling. Set it generously and let the runtime shrink it. The mechanics of re-planning are in the AQE article.

Reading the stage timeline and the summary metrics

Everything above is diagnosable from the stage detail page, and most people never scroll far enough down it.

The event timeline draws one lane per executor and one bar per task, coloured by phase: scheduler delay, task deserialization, shuffle read, executor computing, shuffle write, result serialization. Read its shape first. Clean vertical bands mean full waves and healthy parallelism. A dense block that thins into a few long bars is the ragged tail or a straggler. Long bars in one lane only points at a bad node. A wide green scheduler-delay band across everything means the driver is the bottleneck — usually too many tiny tasks, or a driver busy with broadcast and result collection.

Below it, Summary Metrics gives min / 25th / median / 75th / max for duration, GC time, shuffle read size and records, and spill. This table is the single most useful thing in the Spark UI, because it turns "the stage is slow" into a distribution. The ratio of max to median is the skew measurement you actually want, and it is available without instrumenting anything.

Two more panels earn attention: Aggregated Metrics by Executor, which isolates a single misbehaving host, and the DAG visualization, where greyed boxes are skipped stages and an AQEShuffleRead node annotated coalesced or skewed tells you the runtime already intervened. If the plan header still says isFinalPlan=false, the query is mid-flight and the plan you are reading is not the one that will run.

A straggler diagnosis playbook

When one task holds a stage hostage, work the distribution rather than guessing. The order below resolves the large majority of cases.

1. Compare max and median shuffle read. If max is many times median, it is data skew — one key dominates. Fix at the source: enable AQE skew join, lower skewedPartitionThresholdInBytes if the hot partition sits under it, salt the key, or convert to a broadcast join if one side is small enough. Speculation will not help here.

2. If input and record counts are even but one task is slow, it is the machine, not the data. Check Aggregated Metrics by Executor for a single bad host — a degraded disk, a noisy neighbour, a preempted node retrying. This is the speculation case; spark.excludeOnFailure.enabled handles the repeat-offender version.

3. Check GC time. GC above roughly a tenth of task duration means memory pressure — the partition is too big for its executor's share. More partitions or more memory per executor; see unified memory management.

4. Check spill. Non-zero disk spill with even sizes means the working set genuinely exceeds execution memory. Raise the partition count so each task handles less.

5. Check scheduler delay and task count. High delay with thousands of sub-second tasks is over-partitioning; coalesce or let AQE do it. The cluster-level context for all of this — driver, executors, slots — is in the execution architecture article.