PRACTICE / LESSON 17

Predict what Spark will do

Three interview exercises. Read the code, predict the work, then check your reasoning against the plan and Spark UI.

25 min read + practicePySpark3.5.7 reference

THE IDEA TO KEEP

Predict the distribution each operator needs. State your assumptions before giving a task count, then verify the actual execution.

Set the rules before counting tasks

Start after the complete execution-flow lesson. These are original practice exercises for classic PySpark 3.5.7 batch execution. They are not questions reported from an interview and this page does not run Spark.

Use the setup below: four input partitions, four configured SQL shuffle partitions, AQE disabled, no cached/reused query results, and no failed or speculative attempts. We disable AQE only to make the first predictions easier to inspect; this is not a production tuning recommendation. Other versions and managed runtimes can produce different plans.

Each answer gives a source-reviewed expected shape, not an observed execution. Partition counts describe work to schedule; available CPU/resource slots determine how much can run concurrently. Stage IDs, job totals, and durations are runtime evidence.

Baseline references: Range partitions · AQE settings · task scheduling

One small dataset, three predictions

Run this setup once in a notebook with an existing spark session. Create a local session if needed. A new SQL session keeps these SQL settings separate, but shares the application’s SparkContext and cache. Do not call practice.stop() to clean up the exercise: it would stop the shared context.

PYTHON / PYSPARK

from pyspark.sql import functions as F
from uuid import uuid4

# Classic PySpark: use an existing spark session.
practice = spark.newSession()
practice.conf.set("spark.sql.adaptive.enabled", "false")
practice.conf.set("spark.sql.shuffle.partitions", "4")

orders = (
    practice.range(0, 12, 1, 4)
    .withColumn("bucket", F.col("id") % 3)
    .withColumn("amount", F.col("id") * 50)
)

# This default is for local mode on your own machine only.
# On a cluster, replace it with an accessible shared scratch URI.
output_path = f"file:///tmp/spark-fieldnotes-{uuid4().hex}/totals"
print("Spark version:", practice.version)
print("Spark UI:", practice.sparkContext.uiWebUrl)
print("Scratch destination:", output_path)

The write destination above is for local mode only. On a cluster, choose a fresh scratch URI on storage reachable by all participating executors before running exercise 2. The write uses errorifexists, so it will fail if you reuse an existing destination. Rerun the setup to create a new local destination.

IDs 0–11 form four input partitions: 0–2, 3–5, 6–8, and 9–11. Each row has bucket = id % 3 and amount = id × 50. The filter in every exercise retains IDs 3–11.

Setup references: SQL session isolation · Range API · write modes

EXERCISE 01 / AUTHORED PRACTICE

Filter → count

A filter keeps rows in place. Does that make the whole query shuffle-free?

Run the shared setup above first. Predict before running this snippet.

PYTHON / PYSPARK

filtered = orders.filter(F.col("amount") > 100)
rows = filtered.count()
print(rows)

Pause and predict

Which line requests distributed execution?
Does the filter change the four input partitions? Does the count need a different distribution?
How many tasks would you expect on each side of that boundary under the baseline?
What is the result, and how much data returns to the driver?

Say or write your prediction before revealing the answer.

Reveal the reasoning: Filter → count

Follow the work

Conceptual flow under the shared baseline; not a captured Spark UI trace.

4 input tasksGenerate IDs, evaluate the filter, compute partial counts.
Shuffle to one partitionExchange SinglePartition brings partial counts together.
1 final taskCombine partial counts; return one count to the driver.

The filter constructs a DataFrame expression; count() requests a result. For this Range input, generating and filtering rows is deferred until execution. Python assignments and constructing expressions still run on the driver.

IDs 3 through 11 survive. Filtering does not coalesce away an empty input partition: the task for IDs 0–2 still evaluates its input. The filter itself does not redistribute rows. The global count, however, combines partial counts using a SinglePartition exchange in the expected plan.

With four Range partitions, AQE off, no reused work, and no retries, predict four upstream tasks and one downstream task. This is a two-stage reference shape for this count, not a rule that every count always has two stages. Tasks can run in waves if fewer slots are available. Only a count is returned; the driver does not need all nine surviving rows.

Expected result

9. The input partitions contribute counts 0, 3, 3, and 3 before the final combination.

Check the plan and Spark UI

PYTHON / PYSPARK

# Preview an equivalent global-count query, not just the filter.
filtered.agg(F.count("*").alias("rows")).explain("formatted")
# The earlier filtered.count() is the action to inspect in the UI.

Look for partial and final count aggregates separated by Exchange SinglePartition. filtered.explain() alone omits the count operation.
In the SQL tab, open the count execution and follow its associated jobs to the Stages tab. Compare the four upstream tasks with the single final task; count attempts separately if retries occurred.
Check the filter output-row metric and exchange records where exposed. Expect nine retained rows overall, but partial counts rather than all input rows at the exchange.

A short interview answer

The filter is local to each input partition, but the final count is global. Under this four-partition baseline I expect four tasks to filter and partially count, then a single downstream task to combine those counts. The answer is nine. I would verify the count query, rather than only explaining the filtered DataFrame.

Interviewer follow-up: If the filter removes every row, do the input tasks disappear?

For a predicate that Spark must evaluate against this input, tasks still inspect the partitions and the final count is zero. Empty output is not proof of zero work. A condition proved false during optimization, such as a literal false filter, can instead simplify the plan before scanning. Inspect the actual plan before assigning task counts.

Check the sources: count api · count source · aggregate planning · explain · spark ui

Record your own version in My notebook →

EXERCISE 02 / AUTHORED PRACTICE

GroupBy → write

Three keys produce three totals. Predict the shuffle, the writer tasks, and where those totals go.

Run the shared setup above first. Predict before running this snippet.

PYTHON / PYSPARK

totals = (
    orders.filter(F.col("amount") > 100)
    .groupBy("bucket")
    .agg(F.sum("amount").alias("total"))
)
writer = totals.write.mode("errorifexists")
writer.parquet(output_path)

Pause and predict

Does creating writer execute the query, or does parquet(output_path)?
Why is partial aggregation useful before the shuffle?
With four configured shuffle partitions and three keys, how many downstream tasks would you predict?
Does the data pass through the driver? Can you deduce the exact file count from the number of executors?

Say or write your prediction before revealing the answer.

Reveal the reasoning: GroupBy → write

Follow the work

Conceptual flow under the shared baseline; not a captured Spark UI trace.

4 input tasksFilter and partially sum amounts by bucket.
Shuffle by bucketHash distribution targets 4 downstream partitions in this baseline.
4 downstream tasksFinalize totals and write data to the scratch destination.

Creating the writer and selecting its mode configure a write. parquet(output_path) requests execution. There is no separate action needed to materialize totals first.

The expected shape has a partial sum before an exchange and a final sum after it. A partial sum can reduce the records transferred for repeated keys. Our deliberately small example has one of each surviving key per input partition, so it is not a demonstration of substantial data reduction.

Under the baseline, predict four upstream tasks and four tasks computing the shuffled output and writing it. Three distinct keys do not imply three partitions: keys may hash into the same partition, and some of the four partitions can be empty. A plain Parquet write does not inherently require another shuffle after this final aggregate.

Writer tasks write data to storage while the driver coordinates execution and commit. A writer task is not an executor: file counts also depend on empty partitions, file splitting, and writer behavior. Do not promise one file per key, task, or executor. The scratch path must be reachable by the executors; a driver-local directory is unsuitable for a distributed cluster.

Expected result

Written totals: bucket 0 → 900; bucket 1 → 1050; bucket 2 → 1200. Row order and exact part-file count are not promised. The mode rejects an existing destination.

Check the plan and Spark UI

PYTHON / PYSPARK

# Preview the data query; this does not show the entire write command.
totals.explain("formatted")

# Optional read-back AFTER inspecting the write's jobs and stages.
# This read and sort/show create separate work; do not count it as the write.
practice.read.parquet(output_path).orderBy("bucket").show()

Look for the bucket hash exchange between partial and final sum. Use the completed write execution in the UI to inspect the actual writer path, rather than assuming totals.explain() includes the write command.
Follow the write execution to its stages. Compare upstream tasks with the four downstream tasks; inspect output records and shuffle metrics if available.
Read the directory back only after recording the write execution. The verification read, sorting, and display add their own work. Compare the three totals rather than equating storage files with executor processes.

A short interview answer

The aggregation needs compatible distribution by bucket, so I expect partial sums, a shuffle, and final sums. With AQE disabled and four shuffle partitions, I predict four downstream writer tasks even though there are only three keys. The Parquet call triggers the work, and tasks write to storage while the driver coordinates it. File count needs separate inspection.

Interviewer follow-up: Would coalesce(1) before writing be a good way to produce one large production file?

It can force the final output through a single partition, limiting parallelism and creating a bottleneck. It can be reasonable for a tiny export, but choose file sizing and partitioning based on data volume and downstream readers for larger output. It does not remove the earlier aggregation shuffle. Compare both plans and the workload before adopting it.

Check the sources: parquet write · aggregate planning · partition controls · file writing · spark ui

Record your own version in My notebook →

EXERCISE 03 / AUTHORED PRACTICE

Join → aggregate

Broadcast a tiny lookup, then aggregate by its region. Is the entire query now free of shuffles?

Run the shared setup above first. Predict before running this snippet.

PYTHON / PYSPARK

regions = (
    practice.range(0, 3, 1, 1)
    .select(
        F.col("id").alias("bucket"),
        F.when(F.col("id") < 2, "Domestic")
         .otherwise("International").alias("region"),
    )
)
by_region = (
    orders.filter(F.col("amount") > 100)
    .join(F.broadcast(regions), "bucket", "inner")
    .groupBy("region")
    .agg(F.sum("amount").alias("total"))
)
rows_by_region = by_region.collect()  # Only two result rows here.
print(rows_by_region)

Pause and predict

Which side is being broadcast, and what does that change about the join?
Why might you see both BroadcastExchange and a hash-partitioning Exchange?
What task groups would you look for under the baseline?
What are the two totals? Would collect() still be suitable for millions of output rows?

Say or write your prediction before revealing the answer.

Reveal the reasoning: Join → aggregate

Follow the work

Conceptual flow under the shared baseline; not a captured Spark UI trace.

1 lookup taskProduce the tiny lookup for broadcast to participating executors.
4 fact-side tasksFilter orders, join against the broadcast lookup, partially sum by region.
Shuffle by region → 4 tasksCombine region totals and return the two output rows.

For this supported inner equi-join with an explicit broadcast hint, expect BroadcastHashJoin with the tiny regions side as the build side. That side must be materialized and distributed. The join can then probe it locally while processing the orders partitions; broadcasting is still data movement.

The later groupBy(region) has its own distribution requirement. Expect a hash shuffle for region totals even though the join avoids shuffling the larger orders side for the join itself. BroadcastExchange is not a hash shuffle. Explain what each exchange does instead of counting every Exchange label as the same operation.

With this input and baseline, look for one task producing the one-partition lookup, four tasks processing orders/join/partial sums, and four downstream tasks combining shuffled totals. Broadcast materialization can have its own job, so do not infer one job from one collect(). Treat this as a predicted task-group shape, not fixed stage IDs or a captured run.

Buckets 0 and 1 map to Domestic; bucket 2 maps to International. The lookup has exactly one row per bucket. If it contained duplicate bucket keys, the join could multiply rows and totals before aggregation.

Expected result

Domestic → 1950; International → 1200. The two collected rows have no guaranteed ordering. Collect is reasonable for this tiny result, not an unbounded production output.

Check the plan and Spark UI

PYTHON / PYSPARK

by_region.explain("formatted")
# Inspect the already-completed collect execution in the UI.
# Explaining an unexecuted query predicts operators; it does not measure them.

Distinguish the lookup BroadcastExchange and BroadcastHashJoin from the later hash-partitioning Exchange on region.
In the SQL execution, follow the associated jobs and compare their stages. Inspect the lookup materialization separately from the fact-side pipeline and final aggregate.
Check join output rows against the nine filtered input rows, and final output against two regions. Unexpected row growth can indicate duplicate dimension keys.

A short interview answer

Broadcasting the tiny lookup lets the fact-side tasks perform this join without reshuffling orders for the join. The subsequent region aggregation still requires a shuffle. I would distinguish the broadcast exchange from the hash exchange in the plan, check the associated jobs, and validate that the lookup keys do not multiply rows.

Interviewer follow-up: If the lookup grows, should you keep the broadcast hint?

Reassess it. An explicit hint can prioritize broadcasting beyond the automatic size threshold, but collecting and distributing a large build side can exceed driver or executor memory or time out. Inspect sizes, statistics, join eligibility, and memory before choosing broadcast versus a shuffle join. A hint is not a universal speedup.

Check the sources: broadcast api · join planning · broadcast exchange · aggregate planning · collect api

Record your own version in My notebook →

Change one assumption: turn AQE on

Run this in a PySpark notebook with an existing spark session. Create a local session if needed. Examples are illustrative; this website does not execute Spark.

PYTHON / PYSPARK

# After the shared setup, compare AQE on the grouped query.
# Rebuild the query so it uses the changed SQL configuration.
practice.conf.set("spark.sql.adaptive.enabled", "true")
try:
    adaptive_totals = (
        orders.filter(F.col("amount") > 100)
        .groupBy("bucket")
        .agg(F.sum("amount").alias("total"))
    )
    adaptive_totals.explain("formatted")  # Initial plan
    adaptive_rows = adaptive_totals.collect()  # Only three rows
    adaptive_totals.explain("formatted")  # Inspect final adaptive plan
    print(adaptive_rows)
finally:
    practice.conf.set("spark.sql.adaptive.enabled", "false")

Expected result

The same three bucket totals: 0 → 900, 1 → 1050, 2 → 1200, in unspecified order. Compare initial and final plans and observed downstream task counts. AQE may coalesce shuffle partitions; do not prescribe an exact adaptive task count.

Result reasoned through on the small input; not runtime-tested against a Spark cluster.

Watch the assumption

Do not compare unrelated runs: the optional read-back, show(), collect(), or another action can create additional work. Inspect the SQL execution associated with the action you are explaining. WholeStageCodegen labels and AQE query stages are not scheduler stage IDs.

EXPLAIN IT BACK

How would you explain a PySpark snippet you have never seen before?

Say your answer out loud before opening the notes.

Reveal the reasoning

First I state the runtime, data shape, and settings. I identify the action and trace the plan from input to output, explaining each operator’s required distribution. I distinguish local processing, broadcasting, and shuffling. Then I predict task groups from the relevant partitions, keeping concurrency separate. Finally I check the executed plan and associated jobs and stages, validate the result, and explain any differences from my prediction.

Sources & verification

Core claims checked against Apache Spark 3.5.7 · classic PySpark batch documentation and source references on 2026-09-14. Explanations, scenarios, and diagrams are original teaching material. Other releases and managed runtimes can differ.