Lineage That Grows in Loops

Concepts covered: Spark Execution Model

Reassigning a DataFrame nests the old plan

A DataFrame is a plan, and every transformation returns a new DataFrame whose plan holds the old one as its child. Inside a loop, labels = new_labels rebinds a Python name to the newer plan. The older plan stays reachable as a subtree, so after 30 rounds the name labels points at a single DataFrame whose plan stacks 30 joins and 30 groupBys on top of the original nodes scan.

Here is the job from the intro. The edges table holds each connection in both directions plus a self-loop for every member, so a member's own label is always on offer. Each round, every member takes the smallest label it is offered. Labels can only fall, so once their sum stops falling, no label has moved and the loop stops.

labels = nodes.select("node_id", F.col("node_id").alias("label"))
previous_total = None

for round_no in range(1, 31):
    new_labels = (
        edges.join(labels, edges.src == labels.node_id)
        .groupBy(F.col("dst").alias("node_id"))
        .agg(F.min("label").alias("label"))
    )
    total = new_labels.agg(F.sum("label")).first()[0]
    if total == previous_total:
        break
    previous_total = total
    labels = new_labels

Read the loop the way Spark does. The single line that touches data is first(), the action at the end of the stopping check. Every other line adds operators to a plan held in driver memory. The table follows what labels refers to as each round begins.

Building that plan computed none of it. The description grows by 1 join and 1 groupBy every round whether an action runs, and each action decides how much of it gets executed. The next section prices that decision.

This loop runs 12 rounds with no action inside it. When the loop finishes, what has Spark done?

for i in range(12):
    df = df.join(dims[i], "user_id").filter(F.col("score") > 0)

• Run 12 jobs, one per round, and passed each round's rows to the next
• Run nothing; df is one DataFrame whose plan holds all 12 joins and filters
• Run nothing; df holds only round 12's join and filter, applied to the input
• Run the first round to learn the schema and deferred the other 11

A long plan is driver work

The driver pays for that plan in 2 places. Column references are resolved as each DataFrame is built, which is why a misspelled column fails on the line that names it. Every reassignment in the loop therefore checks its new join and groupBy against a plan one round deeper, so a deep loop slows down before any action runs. When an action does run, the driver rewrites the whole plan, chooses how each operator will run, and cuts the result into stages for the DAGScheduler, all on a single machine. At round 20, it does that for 20 joins to produce a single number, and the round after, for 21.

On the Jobs tab, planning shows up as a gap. A job's clock starts when the DAGScheduler receives it, after planning is done, so a check that spends 40 seconds in the driver leaves 40 idle seconds on the event timeline between the previous job's end and its own start. Gaps that widen round after round, on input that stays the same size, mean the plan is growing.

The animation runs the same growth on a loop small enough to watch. It starts from Monday's events, adds another day per round with union, and ends in a single collect().

week = mon
for day in [tue, wed, thu]:
    week = week.union(day)
week.collect()

The loop's 3 rounds launch nothing. When collect() runs, 1 job with 1 stage starts: each day's table brings its 2 partitions along, so the stage runs 8 tasks, and the driver gathers all 12 rows. Every round left its operators, and here its partitions too, in the job that finally ran.

Try it on the seed

This cell rewrites profit 5 times in a loop over orders, then totals it by status. Run it, change range(1, 6) to range(1, 41), and run it again. The totals change; the loop still finishes before any task starts, and all 40 rewrites wait for the collect() at the bottom.

Two lines decide what the cell computes. The reassignment df = df.selectExpr(...) hands each round the DataFrame the previous round built, so the plan stacks every rewrite; without it, each round would discard its rewrite and start over from the unchanged df. And collect() sits after the loop, so all the rounds run inside the single job it launches.

This loop leaves out the reassignment. What does collect() return?

df = orders.select("status", "profit")
for round_no in range(1, 6):
    df.selectExpr("status", "profit * 1.1 AS profit")
df.collect()

• Profit times 1.1 once, since only the last rewrite is kept
• Profit times 1.1 five times, since the plan stacks all 5 rewrites
• The original profit values, since df holds no rewrite at all
• No rows, since the first unassigned line raised an error

About this lesson

This section is part of the Lazy Evaluation: Advanced lesson on DataDriven, the home of data engineering. Lessons cover SQL, Python, Spark, data modeling and pipeline architecture, and each section pairs its explanation with worked examples and practice you run as you go.