Count the jobs, stages, and tasks

Follow the work

Flow summarized from the recorded local run described below.

1. Read: 2 jobstake(1) reads the header line (1 task); inferSchema scans all 4 file partitions (4 tasks).
2. count(): 4 → 14 scan tasks write partial counts to a single-partition shuffle; 1 task adds them. One job with AQE off, two with AQE on.
3. Write: 1 job, 4 tasksNo shuffle: each of the 4 input partitions writes one Parquet part file.

spark.read is not free here. To name the columns, the CSV reader takes the first line with take(1), a one-task job. inferSchema then tokenizes the sampled input to choose types. At the default samplingRatio of 1.0 that is all of it: a second job with one task per file partition. Four small files become four partitions. The question 1 variant, inferSchema without header, measured the same two jobs. With a supplied schema, neither job ran.

count() is a global aggregate. Each scan task computes a partial count, an Exchange SinglePartition shuffles those four numbers, and one task adds them. With AQE off that is one job with two stages, 4 + 1 tasks. Two stages come from the aggregate’s shuffle, not from count() itself.

With AQE on, the shuffle stage is submitted as its own map-stage job (4 tasks). The result job then lists that stage as skipped and runs 1 task. The work is the same; the job count goes from 1 to 2.

The Parquet write has no shuffle, so it is one job with one stage of four tasks, and it wrote four part files. It rescans the CSV because nothing is cached. The schema, though, was fixed at read time and is not inferred again.

Check the plan and Spark UI

df.printSchema()  # inferred types: order_id, customer_id, and amt as integers
for group in ["csv-read", "csv-count", "csv-write"]:
    print_work(group)

1. csv-read should list two jobs: 1 task, then 4 tasks. Replace inferSchema with a schema string and the group has no jobs.
2. csv-count: with AQE off, one job with stages of 4 and 1 tasks. With AQE on, a 4-task job, then a job with one skipped stage and 1 task.
3. csv-write: one job, one stage, 4 tasks, the same under both settings. In the SQL tab, the count query shows its associated job IDs.

A short interview answer

Before count(), the reader already starts work. With header and inferSchema, Spark 3.5.7 runs a one-task job for the first line and a second job that scans all four input partitions to infer types. count() is a global aggregate: four scan tasks produce partial counts, a single-partition shuffle moves them, and one task adds them. With AQE off that is one job with two stages. With AQE on, the shuffle stage runs as its own job and the result job shows it as skipped. The Parquet write has no shuffle, so it is one stage of four tasks. On this input that was four jobs and 14 tasks with AQE off, five jobs with AQE on. Supplying a schema removes both read jobs.

Check the sources: csv docs · csv infer · dataset · shuffle exchange · dag scheduler

Follow the work

Flow summarized from the recorded local run described below.

1. Read: 1 job, 1 taskSpark reads one Parquet footer in a small job to get the schema (no mergeSchema).
2. filter show(): 1 taskshow() asks for 21 rows; the first file partition already had them.
3. SQL show(): shuffle, then roundsAQE off: 4 map tasks, then 1, 4, 20, 75, 8 of the 200 reduce partitions in five jobs. AQE on: 4 map tasks in one job, then 1 coalesced task.

Even without mergeSchema, the Parquet reader ran a one-task job to read a footer. createOrReplaceTempView started no job.

The filter is narrow and there is no exchange. show() calls take(21). The CollectLimit plan scans one partition first and scans more only if it is short of rows. The first file partition holds the largest file (event_id 1500–1999), with 250+ rows with age > 30, so one job with one task was enough.

The GROUP BY plan is partial count → Exchange hashpartitioning(user_id, 200) → final count. With AQE off, the first job ran 4 map tasks and 1 reduce partition. Only 38 of the 200 partitions hold any of the 40 users, and that partition held 2 rows. The next rounds asked for 4, 20, 75, and 8 partitions. Each round was a job that listed the map stage as skipped: 5 jobs, 112 tasks.

Replaying Spark’s rule on the measured rows per partition (2, 3, 7, and 20 rows after 1, 5, 25, and 100 partitions) gives exactly 1, 4, 20, 75, 8. Each later round asks for ceil(1.5 × rows still needed × partitions scanned ÷ rows found), capped at 4× the partitions scanned so far.

With AQE on, the shuffle stage runs as a map-stage job. AQE then coalesces the tiny output (160 records, 3.8 KiB) into one partition (AQEShuffleRead), and one result job finishes with 1 task.

Check the plan and Spark UI

spark.sql("""
    SELECT user_id, COUNT(*) AS cnt
    FROM events
    WHERE event_type = 'click'
    GROUP BY user_id
""").explain()
for group in ["parquet-read", "filter-show", "sql-show"]:
    print_work(group)

1. parquet-read: one job, one task. filter-show: one job, one task.
2. sql-show with AQE off: five jobs. The first has stages of 4 and 1 tasks; each later job has one skipped 4-task stage and runs 4, 20, 75, and 8 tasks.
3. sql-show with AQE on: two jobs, 4 tasks then 1 task after a skipped stage. The SQL tab’s final plan shows AQEShuffleRead over the ShuffleQueryStage.

A short interview answer

Reading Parquet without mergeSchema ran one small footer job. The filtered show() has no shuffle. show() asks for 21 rows, Spark scans one partition first, and that partition had enough, so one job with one task. The SQL aggregation adds a hash shuffle on user_id. With AQE off there are 200 shuffle partitions but only 40 users, so show() scanned them in rounds of 1, 4, 20, 75, and 8 partitions. That was five jobs, and the last four each listed the four-task map stage as skipped. With AQE on, the map stage ran as its own job, AQE coalesced the tiny output to one partition, and one more job finished it. The number of show() jobs depends on where rows land, so I would measure it.

Check the sources: parquet docs · schema merge · execute take · collect limit · coalesce shuffle · limit conf

Follow the work

Flow summarized from the recorded local run described below.

1. Read: 1 job, 1 taskheader only: take(1) for the first line; no inferSchema pass.
2. One shuffle: 4 map tasksExchange hashpartitioning(customer_id, 10), origin REPARTITION_BY_NUM.
3. Aggregate in place: 1, 4, 2 tasksThe partial and final count both run after the exchange; show() scanned 7 of 10 partitions in three rounds.

repartition(10, customer_id) produces hash partitioning on customer_id with 10 partitions. The aggregate needs rows with equal keys in the same partition, and that partitioning already provides it, so no second Exchange is added. The plan shows both HashAggregate operators above a single Exchange. There is no partial aggregation before the shuffle here, so all 2,000 rows are shuffled.

This exchange came from repartition with a number, so its origin is REPARTITION_BY_NUM. AQE’s coalescing rule skips that origin, and there was no AQEShuffleRead. The 10 partitions stayed 10 under both settings.

show() needs 21 rows from 50 keys spread over 10 partitions, about 5 per partition. The first partition held 3 rows, so the next round asked for 4 more partitions (17 rows after 5), then 2 more (28 rows after 7). With AQE off: three jobs, 4 + 1, then 4, then 2 tasks, with a skipped map stage in the last two. With AQE on, the map stage came first as its own job: four jobs with the same 11 tasks.

Check the plan and Spark UI

counts.explain()  # one Exchange: hashpartitioning(customer_id, 10)
for group in ["repartition-read", "repartition-show"]:
    print_work(group)

1. explain() shows one Exchange hashpartitioning(customer_id, 10), REPARTITION_BY_NUM, with both aggregates above it.
2. repartition-show with AQE off: three jobs running 4 + 1, 4, and 2 tasks; the second and third each list one skipped stage.
3. With AQE on: a 4-task map-stage job, then result jobs of 1, 4, and 2 tasks, each with one skipped stage. The final plan has no AQEShuffleRead.

A short interview answer

There is one shuffle, not two. repartition(10, customer_id) hash-partitions by the grouping key, which already satisfies the aggregate’s requirement, so the plan has one Exchange and both aggregates run after it. The read has header but no inferSchema, so it is a single one-task job. The shuffle has four map tasks, one per file, and ten reduce partitions. AQE does not coalesce a repartition with an explicit number. show() needs 21 rows at about five per partition, so it scanned 1, then 4, then 2 partitions. With AQE off that is three jobs and 11 tasks. With AQE on, the map stage ran as its own job, so four jobs with the same tasks.

Check the sources: ensure requirements · coalesce shuffle · execute take · csv infer · shuffle exchange

Follow the work

Flow summarized from the recorded local run described below.

1. Reads: 2 jobsOne footer job per Parquet table, 1 task each.
2. Broadcast: 2 jobsCollect customers (1 task) for BroadcastExchange; then show() probes one orders partition (1 task).
3. Shuffle join: 4 + 1 + 1Hash-shuffle orders (4) and customers (1), then SortMergeJoin on 1 of 200 reduce partitions. One job with AQE off; three with AQE on.

Each spark.read.parquet ran a one-task footer job, two jobs before any action.

The customers files are far below the 10 MB autoBroadcastJoinThreshold, so the plan chose BroadcastHashJoin BuildRight. The build side is collected to the driver and broadcast in a separate job (it appears as $anonfun$withThreadLocalCaptured$1), 1 task for one file. show() then ran 1 task on the first orders partition, which had enough matching rows. Two jobs, no shuffle, under both settings.

With the threshold at -1 the plan became SortMergeJoin with two hash exchanges on customer_id (200 partitions). With AQE off, show()’s first job ran three stages: 4 orders map tasks, 1 customers map task, and 1 reduce task. Reduce partition 0 held customer 26’s 80 joined rows, enough for 21, so no second round ran.

With AQE on, each shuffle stage became its own job (4 tasks, then 1). AQE coalesced the reduce side to one partition, so the result job listed two skipped stages and ran 1 task. The same six tasks ran either way, as three jobs instead of one.

Check the plan and Spark UI

orders.join(customers, "customer_id").explain()  # threshold is still -1 here
for group in ["join-read", "join-broadcast", "join-shuffle"]:
    print_work(group)

1. explain() after setting the threshold to -1 shows SortMergeJoin with two Exchange hashpartitioning(customer_id, 200). Before that, the SQL tab shows BroadcastHashJoin with BroadcastExchange.
2. join-broadcast: two 1-task jobs under both settings, the first being the broadcast collection.
3. join-shuffle with AQE off: one job with stages of 4, 1, and 1 tasks. With AQE on: jobs of 4 and 1 tasks, then a job with two skipped stages and 1 task.

A short interview answer

Two reads come first: each Parquet table ran a one-task footer job. customers is tiny, so the default plan broadcasts it. The broadcast collection is its own one-task job, and show() then probes a single orders partition, so there is no shuffle. If I disable broadcast, the plan becomes a sort-merge join with two hash shuffles. With AQE off, one job has three stages: four map tasks for orders, one for customers, and one reduce task, because the first reduce partition already had 21 joined rows. With AQE on, each shuffle stage runs as its own job and the result job lists both as skipped. Same tasks, three jobs. Broadcasting avoids the shuffles, but it still adds a job.

Check the sources: broadcast exchange · aqe · shuffle exchange · execute take · schema merge

Follow the work

Flow summarized from the recorded local run described below.

1. Read + cache(): 1 jobThe Parquet footer job (1 task); cache() only marks the plan.
2. First count: fills cacheAQE off: 4 tasks scan Parquet and store cached partitions, then 1 task. AQE on: cache job (4), map job (4), result (1).
3. Filtered count: from cache4 tasks read cached partitions, then 1 task. AQE on splits it into two jobs.

cache() started no job. It marks the plan for caching, and the first action fills the cache. The only job before the first count was the Parquet footer read.

With AQE off, the first count is one job with two stages. Its 4 scan tasks read 28,325 bytes (2,000 records) of Parquet, build the cached partitions, and compute partial counts, and 1 task adds them. The filtered count is also 4 + 1 tasks. Its plan uses InMemoryTableScan, and its scan stage read 8,496 bytes, 4 input records, from the cached blocks instead of the files.

With AQE on, Spark 3.5.7 treats the not-yet-built cache as a query stage (TableCacheQueryStage). It builds the cache first with its own job over every cached partition: a separate 4-task job that read the Parquet files and wrote no shuffle. Next came the shuffle map stage (4 tasks, reading the cache) and a result job with a skipped stage (1 task): three jobs for the first count, two for the filtered count.

The cache avoided rereading the files, but the tasks did not go away. Every count still scans every cached partition.

Check the plan and Spark UI

df.filter(df.amt > 1000).explain()  # InMemoryTableScan instead of a file scan
print(df.storageLevel)
for group in ["cache-read", "cache-count", "cache-filter-count"]:
    print_work(group)
df.unpersist()

1. explain() on the filter shows InMemoryTableScan over InMemoryRelation, not a Parquet FileScan.
2. In the stages list, the cache-building scan stage reports about 28 KB of input and 2,000 records. The filtered count’s scan stage reports 8,496 bytes and 4 records.
3. cache-count with AQE on: 3 jobs (4, 4, skipped + 1). With AQE off: 1 job (4 + 1). The Storage tab lists the cached relation after the first count.

A short interview answer

cache() does not start anything by itself, so the only job before the first count is the Parquet footer read. With AQE off, the first count is one job with two stages. Four tasks scan the files, fill the cache, and compute partial counts, and one task adds them. The filtered count is again four plus one tasks, but its plan reads InMemoryTableScan and its input metrics are cached blocks, not Parquet bytes. With AQE on, Spark 3.5.7 builds the cache as its own query stage first, so the first count ran three jobs: build the cache, the shuffle map stage, then the result. The filtered count ran two. Caching saves rereading the files, not tasks.

Check the sources: cache api · query stage · shuffle exchange · web ui · dag scheduler

Follow the work

Flow summarized from the recorded local run described below.

1. 3 map tasksparallelize’s 3 partitions are hash-partitioned by key into shuffle files.
2. Shuffle: 3 partitionsNo spark.default.parallelism set, so PySpark keeps the parent’s 3 partitions.
3. 3 reduce tasksGroup values per key, mapValues(list) in the same stage, return rows to the driver.

collect() is the only action, so there is one job. groupByKey introduces a shuffle, so the job has two stages. mapValues is narrow and runs in the reduce stage.

PySpark’s groupByKey defaults its partition count from _defaultReducePartitions(). That uses spark.default.parallelism only if the key is set in the configuration; otherwise it keeps the RDD’s own partition count. Here that is 3, even though sc.defaultParallelism is 4 under local[4]. getNumPartitions() on the grouped RDD returned 3.

AQE optimizes SQL/DataFrame plans only, so this job was identical with it on or off: 3 map tasks and 3 reduce tasks.

Check the plan and Spark UI

grouped = pairs.groupByKey().mapValues(list)
print(grouped.getNumPartitions())  # reduce-side partitions
print(grouped.toDebugString().decode())  # indentation marks the shuffle boundary
print_work("rdd-groupbykey")

1. getNumPartitions() prints 3; toDebugString() shows the ShuffledRDD from partitionBy with a (3) branch under it.
2. rdd-groupbykey: one job with two stages of 3 tasks each, the same with AQE on and off.

A short interview answer

One action, collect, so one job. groupByKey adds a shuffle, so there are two stages, and mapValues runs inside the second one. The first stage has three tasks, one per parallelize partition. In PySpark, groupByKey without a partition count uses spark.default.parallelism only if that key is set; otherwise it keeps the parent’s three partitions, so the reduce stage has three tasks. That makes six tasks. AQE does not apply to RDDs, and I measured the same counts with it on and off. The result groups a to [1, 3], b to [2, 4], and c to [5]. Key order depends on hashing, so I would not promise it.

Check the sources: rdd source · dag scheduler · glossary · aqe

Follow the work

Flow summarized from the recorded local run described below.

1. 1 jobcollect() is the only action.
2. 1 stagemap and filter are narrow and pipeline together in each task.
3. 2 tasksPartition 1–5 returns nothing after the filter; partition 6–10 returns 12–20.

Both transformations are narrow, so there is no shuffle and the single job has a single stage.

The stage runs one task per partition: 2. The first partition, 1–5, doubles to 2–10 and keeps nothing, but its task still runs. An empty result does not remove work.

AQE does not apply to RDDs. Both settings measured 1 job, 1 stage, 2 tasks.

Check the plan and Spark UI

print_work("rdd-map-filter")

1. rdd-map-filter: one job, one stage, two tasks, two attempts, under both settings.

A short interview answer

map and filter are narrow and run together in the same task, and collect is the only action, so this is one job with one stage. The stage runs one task per partition, two tasks. The first partition keeps nothing after the filter, but its task still runs. The result is 12, 14, 16, 18, and 20. AQE does not change RDD jobs, and I measured the same counts either way.

Check the sources: glossary · dag scheduler · web ui

Follow the work

Flow summarized from the recorded local run described below.

1. Read: 1 job, 4 tasksSchema inference scans all four JSON files.
2. show(): roundsAQE off: 4 + 1, then 4, 20, 75, 8 reduce tasks in five jobs. AQE on: 4 map tasks, then 1 coalesced task.
3. write: shuffle againAQE off: 4 map + 200 write tasks in one job (38 files). AQE on: 4 map tasks, then 1 write task.

Without a schema, the JSON reader infers one by reading the records, one job with a task per file partition.

The aggregate is a partial avg, then Exchange hashpartitioning(user_id, 200), then the final avg. show() needs 21 rows. With AQE off, the 40 users land in 38 of 200 partitions, the same layout as exercise 2, so the rounds were identical: 1, 4, 20, 75, 8 reduce partitions over five jobs, the last four listing a skipped map stage.

The write is a separate action with its own physical plan and its own shuffle. Nothing is cached, so it rescanned the JSON files and reran the 4-task map stage. No stage was skipped from show(). With AQE off, the write stage ran all 200 partitions, 204 tasks in total, and wrote 38 part files because empty partitions write none. With AQE on, the map stage ran as its own job and the output was coalesced to one partition, so there was 1 write task.

Totals: AQE off, 7 jobs and 320 tasks. AQE on, 5 jobs and 14 tasks. Most of the AQE-off tasks are empty reduce partitions, which is why the 200 default matters for small data.

Check the plan and Spark UI

avg_duration.explain()
for group in ["json-read", "json-show", "json-write"]:
    print_work(group)

1. json-read: one job with 4 tasks. With a supplied schema, the read starts no job.
2. json-show with AQE off: five jobs with 4 + 1, 4, 20, 75, and 8 tasks, the last four each with one skipped stage.
3. json-write: AQE off, one job with 4 + 200 tasks and no skipped stage. AQE on, a 4-task job, then a job with one skipped stage and 1 task. Count the part files in the output directory.

A short interview answer

The JSON read infers a schema, which is a four-task job before any action. The aggregation adds a hash shuffle on user_id. With AQE off there are 200 shuffle partitions but only 40 users, so show() scanned in rounds of 1, 4, 20, 75, and 8 partitions: five jobs, and the later ones list the map stage as skipped. The write is a separate action with its own plan. Without a cache it rescans the JSON and reruns the shuffle, then runs 200 write tasks and writes 38 files, since only 38 partitions hold data. With AQE on, each action runs a map-stage job and a result job on one coalesced partition. Nothing is reused between show and write.

Check the sources: json docs · json infer · execute take · shuffle exchange · coalesce shuffle