Databricks Spark Developer Associate: The 4,000-Task Stage That Nobody Can Speed Up With More Executors

The 4,000-Task Stage That Nobody Can Speed Up With More Executors

A nightly job reads a 4,000-file Parquet directory (~1.8 TB), applies a handful of narrow withColumn transformations, and writes 4,000 output files. No shuffle occurs anywhere in the job. The cluster provides 40 executors x 4 cores, i.e. 160 task slots. In the Spark UI, the write stage shows 4,000 tasks with a median task duration of 900 ms, but the 75th-percentile scheduler delay is 2.4 s. The driver process sits at 95-100% CPU for the entire 70-minute stage while average executor CPU is 22%. spark.sql.shuffle.partitions is 200. Adding executors (40 -> 60) changed wall-clock time by less than 3%. A re-run on identical data reproduced the same driver CPU profile and the same runtime within two minutes.

1. Where is the bottleneck? Which observations rule out the two most common first guesses ("not enough parallelism" and "shuffle/IO pressure")?

Read the worked answer

Bottleneck: the driver, specifically per-task scheduling overhead.

Reasoning chain:

Slot math. 4,000 tasks over 160 slots = 25 waves. With a 900 ms task body, each wave is dominated by dispatch and bookkeeping, not work.
Scheduler delay of 2.4 s (75th pct) is time a task sits in the queue after being scheduled but before it starts running - that is driver-side, not executor-side. A task that runs 900 ms but waits 2.4 s means ~73% of its lifetime is driver latency.
Driver at ~100% CPU while executors average 22% is the signature of a driver that cannot hand out work fast enough: task serialization, DAG scheduling, status/heartbeat processing and stage cleanup all run there.
More executors did nothing (<3%) - that is decisive. If the constraint were slots, 50% more slots would have produced a visible gain. The constraint is upstream of the executors.

Why the usual guesses are wrong:

Not enough parallelism: 4,000 tasks is already far more than 160 slots; parallelism is excessive, not insufficient.
Shuffle/IO: there is no shuffle in this job, so spark.sql.shuffle.partitions=200 is irrelevant here (a trap - the value is a red herring).

The real problem is that the total task count was inherited from the input file granularity, and each task is far too cheap to justify its scheduling cost.

2. Propose concrete changes to task count and configuration. For each, state precisely what it fixes and what it costs - and explain why AQE will not rescue this job.

Read the worked answer

1. Decouple read partitions from write tasks.

Read side: one partition per file (plus split chunks) is why you get 4,000. If files are ~450 MB each, spark.sql.files.maxPartitionBytes / openCostInBytes tuning won't reduce task count much - the files are already large. The task count is a file artifact.
Write side: df.coalesce(400) before the write merges 10 input partitions into one task, giving ~2.5 waves over 160 slots. Each task now writes ~4.5 GB, which is fine for sequential Parquet writing but increases the blast radius of a single task failure and reduces retry granularity - that is the cost.
Do not use repartition(400) here: it inserts a full shuffle of 1.8 TB purely to fix a scheduling problem.

2. Why AQE does not help. spark.sql.adaptive.coalescePartitions only coarsens partitions at shuffle stage boundaries. With zero shuffles there is no query stage for AQE to re-optimize, so spark.sql.shuffle.partitions=200 and all AQE knobs are inert in this job.

3. Driver-side trims (secondary). Reduce status-update noise with a larger spark.ui.retainedTasks-style tuning only if needed; the primary fix is fewer, larger tasks. Also verify the driver has adequate heap and is not GC-thrashing (-XX:+UseG1GC, GC time in the UI), because a driver pinned at 100% CPU from GC behaves identically to a scheduler-bound driver.

Decision rule: target roughly 2-4 tasks per available slot per stage for cheap transformations; below that, per-task overhead dominates; far above it, you pay driver CPU.

3. What would you measure before and after the change to prove the diagnosis, rather than just observing that the job 'got faster'?

Read the worked answer

Measure the mechanism, not only the outcome:

Scheduler delay percentiles per stage (Spark UI -> Stage detail). If the diagnosis is right, p75 scheduler delay should collapse and task duration should rise (bigger tasks). Total task time may not drop - the win is in the gaps.
Driver CPU and GC time over the stage window (JVM metrics / event log). Expect driver CPU to fall well below saturation.
Executor CPU utilisation should rise from ~22% toward the practical ceiling for a write-bound stage.
Task deserialization time and result serialization time per task, which shrink proportionally with task count.
Wave count: 4,000/160 = 25 waves before vs ~2.5 after; if runtime does not improve roughly with wave count, another constraint (sink throughput, file commit cost) is now binding.
Output file sizes after coalescing, to confirm you did not simply trade driver overhead for pathological 4.5 GB files or, conversely, for tiny files at the sink.

A useful falsification test: keep the same data and same executor count but only change task count. If runtime does not move, the bottleneck was never scheduling - look at the sink commit protocol or driver GC.