Checkpointing to Cut Lineage

Saving the rows, dropping the recipe

df.checkpoint() runs the plan of df, writes the resulting rows as files under the checkpoint directory, and returns a new DataFrame whose plan is a single scan of those files. The returned DataFrame keeps no link to the joins that produced it. Any plan built on it starts from the saved rows, so a later check can reach back no further than the most recent save.

spark.sparkContext.setCheckpointDir("/shared/checkpoints/label_propagation")

CHECKPOINT_EVERY = 5
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"))
    )
    if round_no % CHECKPOINT_EVERY == 0:
        new_labels = new_labels.checkpoint()
    total = new_labels.agg(F.sum("label")).first()[0]
    if total == previous_total:
        break
    previous_total = total
    labels = new_labels

Two lines carry the change. setCheckpointDir names a directory on storage that the driver and every executor can reach and that outlives any executor; call checkpoint() without it and Spark raises an error saying the checkpoint directory has not been set. The reassignment new_labels = new_labels.checkpoint() is what cuts the lineage. The call leaves the DataFrame it was invoked on unchanged, so dropping its return value pays for the save and carries on with the uncut plan.

Eager and lazy saves

checkpoint() is eager by default: the call itself launches the save, one of those jobs that run before your own action. With checkpoint(eager=False), the call returns at once and the files are written when the first action runs on the returned DataFrame. Loops usually keep the default, so the save lands on a known line and the check right after it reads files.

A reliable save makes 2 passes

A checkpoint runs the plan twice. The first pass is a job that counts the rows; the second is a job that writes them, and since the first pass kept nothing, the second recomputes every row. Spark's documentation for the underlying checkpoint call recommends keeping the data in executor storage before saving so the write reads stored rows instead; the caching lessons later in the track cover how. In the loop above, the save at round 5 replays 5 rounds once to count and once to write.

Choosing the interval

With a save every n rounds, no job ever replays more than n rounds. Within each block of n, the checks before the save replay 1, 2, up to n-1 rounds, and the save replays n rounds twice. Over 30 rounds:

• 4 hours
• 345

Rounds of work fall as the interval shrinks, but each save writes the whole labels table to shared storage and adds 2 jobs for the driver to schedule. Saving every round cuts 465 rounds of work to 60, at the price of 30 full writes of the table and 60 extra jobs. The right interval is where another save costs about as much as the replay it removes, and you find it by timing a save job and a round on the Jobs tab. Spark's own ALS recommender exposes the same choice as checkpointInterval, which defaults to 10 and takes effect only once a checkpoint directory is set.

Here step holds a round's join and groupBy, and the loop is meant to save every 5 rounds. How many rounds does the check at round 30 replay, and why?

for round_no in range(1, 31):
    new_labels = step(labels)
    if round_no % 5 == 0:
        new_labels.checkpoint()
    total = new_labels.agg(F.sum("label")).first()[0]
    labels = new_labels

• All 30, because the saved DataFrame is never assigned back to new_labels
• Only 5, because the save at round 25 cut every round before it
• None, because the save at round 30 runs right before the check reads it
• All 30, because a save cuts lineage only after its files are cleaned up

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.