DE Learning

Spark Core · Module 3 · Examples

Step through it, then count the jobs

Watch a plan grow line by line until an action fires, then practise the skill interviewers test most: looking at code and saying how many jobs and stages it creates.

Lazy evaluation, step by step

Run the pipeline one line at a time. The plan grows while "rows read" stays at zero, until the action.

A quick method for counting

  1. Jobs: count the actions (count, show, collect, write, toPandas…). Add one for schema inference on CSV/JSON, and expect extras for sorts and broadcasts.
  2. Stages: for each job, count the shuffles (groupBy, sort-merge join, distinct, orderBy, repartition), then add one per independent input branch plus the final stage. Or count Exchange nodes in explain().
  3. Tasks: per stage, the number of partitions: input splits for scans, spark.sql.shuffle.partitions after a shuffle.
  4. Reads: each action re-reads its source unless something upstream is cached.

Count the jobs and stages

Work each one out before opening the answer.

1 · A simple filter and write
spark.read.parquet(p).filter("amount > 0").write.parquet(q)
1 job (plus possibly a tiny file-listing job), 1 stage: no shuffle, so read → filter → write is pipelined. Tasks = number of input splits.
2 · Group and show
df.groupBy("segment").count().show()
1 action, 2 stages: scan + partial count, then the final count after the shuffle. With AQE the Jobs tab may show it as two jobs because query stages run separately; the SQL tab shows it as one query.
3 · Join, aggregate, write
(orders.filter("country = 'AE'")
       .join(customers, "cust_id")
       .groupBy("segment").agg(F.sum("amount"))
       .write.saveAsTable("gold.sales"))
1 job, 4 stages: scan orders and scan customers (both shuffle by cust_id), the join plus partial sum (shuffles by segment), then the final sum and write. If customers is broadcast: 3 stages, because the orders side isn't shuffled for the join.
4 · Count, then write
df = spark.read.parquet(p).filter("country = 'AE'")
print(df.count())
df.write.parquet(q)
2 jobs, and the Parquet files are read twice. With df.cache() before the count it's still 2 jobs, but the second reads from the cache. Often the better fix is to drop the debug count().
5 · A loop of writes
for d in ["2026-10-01", "2026-10-02", "2026-10-03"]:
    df.filter(F.col("dt") == d).write.mode("overwrite").parquet(f"{q}/dt={d}")
3 jobs, run one after another, each with its own scan (pruned to that date if the source is partitioned by dt). One write with .write.partitionBy("dt") does it in 1 job.
6 · CSV with schema inference
(spark.read.option("header", True).option("inferSchema", True).csv(p)
      .filter("amount > 0").count())
2 jobs: one scans the data to infer column types, then the count. Passing an explicit schema removes the first job, which on large files is a full extra pass.
7 · Harder: distinct, sort, collect
df.select("segment").distinct().orderBy("segment").collect()
Two shuffles (distinct, then a range-partitioned sort), so 3 stages of real work. The sort needs to know the key ranges, so Spark also runs a small sampling job first. Expect at least 2 jobs in the UI for one collect.