Spark jobs, stages and tasks: how to read the Spark UI

Module 2's script cleaned 8 orders in about a second, and nothing about the run was worth looking at. Real pipelines read millions of rows, and when one is slow you need to know where its time went. This module makes the shop bigger, runs the same kinds of query on it and reads what Spark did, first in the Spark UI and then, in part 4, from a script.

Make a bigger shop

This script writes 2 million orders from 50,000 customers, as 8 Parquet files, and a customers table that gives each customer a country. Save it as make_shop.py and run it once:

"""Make a bigger shop: 2 million orders from 50,000 customers, as Parquet files."""
from pyspark.sql import SparkSession
from pyspark.sql import functions as F

spark = SparkSession.builder.master("local[4]").appName("make-shop").getOrCreate()
LAKE = "lake"

orders = spark.range(2_000_000).select(
    F.col("id").alias("order_id"),
    (F.col("id") % 50_000).alias("customer_id"),
    ((F.col("id") % 9_000) / 100).cast("decimal(10,2)").alias("amount"),
)
orders.repartition(8).write.mode("overwrite").parquet(f"{LAKE}/orders")

customers = spark.range(50_000).select(
    F.col("id").alias("customer_id"),
    (F.col("id") % 30).alias("country_id"),
)
customers.coalesce(1).write.mode("overwrite").parquet(f"{LAKE}/customers")

master("local[4]") runs Spark on your computer with 4 worker threads, as if it had 4 cores, whatever yours has. That keeps your numbers the same as the ones on this page. On Databricks Free Edition, leave out the builder line, since spark is already there, and save tables instead of folders: .saveAsTable("orders") and .saveAsTable("customers").

Driver, executors and tasks

A Spark application has 1 driver and some executors. The driver runs your Python, turns each query into a plan and hands out the work. Executors do the work, each on its cores. The data is cut into partitions, and the unit of work is a task: 1 task does 1 step of the plan on 1 partition. A core runs 1 task at a time, so with 4 cores, 4 tasks run at once and the rest wait their turn.

On local[4] the driver and the executor are 1 process on your computer. On a cluster they are separate machines, but the plan, the jobs, the stages and the tasks are the same.

An action starts a job, a shuffle ends a stage

Save this as runs.py and run it. The last line keeps Spark running so you can look at it:

from pyspark.sql import SparkSession
from pyspark.sql import functions as F

spark = SparkSession.builder.master("local[4]").appName("how-spark-runs").getOrCreate()
orders = spark.read.parquet("lake/orders")
customers = spark.read.parquet("lake/customers")

print(orders.filter("amount > 50").count())

by_customer = orders.groupBy("customer_id").agg(F.sum("amount").alias("revenue"))
by_customer.explain()
by_customer.write.mode("overwrite").parquet("lake/revenue_by_customer")

input("The Spark UI is at http://localhost:4040. Press Enter to stop Spark.")

887778
== Physical Plan ==
AdaptiveSparkPlan isFinalPlan=false
+- HashAggregate(keys=[customer_id#1L], functions=[sum(amount#2)])
   +- Exchange hashpartitioning(customer_id#1L, 200), ENSURE_REQUIREMENTS, [plan_id=63]
      +- HashAggregate(keys=[customer_id#1L], functions=[partial_sum(amount#2)])
         +- FileScan parquet [customer_id#1L,amount#2] ...

Read the plan from the bottom. Each task reads its part of the files and sums each customer's orders in that part, partial_sum. A customer's orders are in every part, though, so to finish the totals Spark must bring each customer's partial sums together. That is the Exchange: every task writes its rows out by customer_id into 200 buckets, and the next tasks each read 1 bucket from all of them. Moving rows between tasks like this is a shuffle, and it is where a job is split into stages, since nothing after it can start until everything before it has written.

count, write, show and collect are actions, and each starts work. The 2 actions here run as 4 jobs, as the UI shows.

Read the Spark UI

Open http://localhost:4040 while the script waits. Start with the SQL / DataFrame tab: it lists each query an action ran, here a count and a parquet write, with the ids of the jobs it ran. Open the write to see its plan as a graph, with the rows that passed through each operator.

The Jobs tab lists 6 jobs. The first 2, with 1 task each, came from spark.read.parquet, which opens a file to learn the schema. The count ran jobs 2 and 3, and the write jobs 4 and 5. Their names, like $anonfun$withThreadLocalCaptured$1, say nothing useful, which is why the SQL tab is the better place to start.

The Stages tab holds the detail. For the write:

Stage 5 read the files in 4 tasks and wrote 3.3 MB of partial sums. Stage 7 read them back and wrote the result. Open a stage and its summary table shows the time and bytes of the quickest task, the median and the slowest: a slowest task far above the median is the first sign of skew, which part 3 takes apart.

Choose a query in the demo and turn adaptive query execution on and off. The numbers are from these runs.

Try itJobs, stages and tasks, as Spark ran them

Query

CountGroup byJoinJoin, no broadcast

Adaptive query execution

OnOff

orders.groupBy("customer_id").agg(F.sum("amount")).write.parquet(…)

Jobs2

Tasks run7

Shuffled3.3 MB

1. Job 1

   1. 4 tasks

      Read, sum each customer's part

      writes 3.3 MB

2. Job 2

   1. Skipped

      Read, sum

   2. →

      3 tasks

      Final sums, write the files

      reads 3.3 MB

Why a stage says skipped

The plan says AdaptiveSparkPlan isFinalPlan=false: Spark has not finished planning. Adaptive query execution runs each stage that ends in a shuffle on its own, as a job, measures what it wrote, then plans the rest. That is why stage 7 has 3 tasks, not the 200 of the plan: 3.3 MB did not need 200 partitions, so Spark merged them.

The second job needs stage 5's output, so it lists stage 5 again, as stage 6, and marks it skipped: its output was already on disk, so no task ran. Skipped stages are normal and cost nothing. With adaptive query execution off, the same write is 1 job of 2 stages, with 200 tasks in the second and 200 files at the end.

On Databricks serverless

Serverless compute has no Spark UI. Databricks points you to the query profile instead: under a cell that ran a query, choose See performance, pick the statement, then See query profile. It draws the query's operators as a graph with the time, rows and peak memory of each, and lists the most expensive ones. It answers the same question, which step took the time, without stages and tasks to read. Classic compute, which Free Edition does not have, keeps the Spark UI.

Part 2 explains where the 4 tasks of the first stage and the 200 of the shuffle come from, and when to change them.