Pipelining Narrow Operations

Concepts covered: Spark Execution Model

Narrow and wide

Sort every transformation by what 1 output partition needs. A narrow transformation builds each output partition from exactly 1 input partition: filter, select, withColumn and union work on the rows in front of them and nothing else. A wide transformation needs rows from many input partitions, because rows that share a key can sit in any of them: groupBy, distinct and orderBy are wide. Every shuffle and every stage cut you have seen so far sits at a wide transformation.

A chain of narrow steps runs as 1 pass

Spark runs consecutive narrow transformations as 1 pass, which is called pipelining. The task reads a row, applies the first operation, hands the result to the second, and carries on to the end of the chain before it reads the next row. No step runs over the whole partition before the next begins, and no intermediate DataFrame is stored anywhere.

The animation plays a 3-step chain on 6 order_items rows split into 2 partitions: keep rows with quantity >= 2, compute line_total, keep 2 columns.

(order_items
   .filter(F.col("quantity") >= 2)
   .withColumn("line_total", F.col("quantity") * F.col("unit_price"))
   .select("order_id", "line_total")
   .collect())

The board shows 1 stage of 2 tasks and no shuffle. Each task carries its rows through all 3 steps, and the driver collects 4 rows: order 1 with 36, order 2 with 50, order 3 with 60 and order 5 with 80. The 2 rows with quantity 1 stop at the filter and never reach the multiplication.

What pipelining buys, and what it hides

• A narrow step adds work per row, never a pass: 10 chained withColumn calls still read the input once.
• A task holds only the few rows in flight, so a 2GB partition can stream through a filter chain.
• The Spark UI times a stage as a whole. If 1 expression in the chain is slow, the stage is slow, with no step named.
• The chain ends at the first wide transformation, and every step after it belongs to the next stage.

Here is a longer job, with the stage each step lands in:

report = (transactions
    .filter(F.col("quantity") > 0)
    .withColumn("unit_price", F.col("total_amount") / F.col("quantity"))
    .groupBy("product_id")
    .agg(F.avg("unit_price").alias("avg_unit_price"))
    .filter(F.col("avg_unit_price") > 20)
    .orderBy(F.desc("avg_unit_price")))

report.collect()

Five transformations, 3 stages, 2 shuffles, plus the sampling job a sort always brings. Each stage is 1 pipelined pass, so when this job is slow, the question for the Spark UI is which of the 3 stages takes the time, then which steps that stage contains.

Run the narrow chain on the practice data. Then add .groupBy("order_id").agg(F.sum("line_total").alias("order_total")) to the end of result and say which steps now run in stage 1 and which in stage 2.

With the groupBy added, stage 1 holds the read, the filter, line_total, the select and the first half of the sum: each task adds up line_total per order_id for the rows it holds. The shuffle sends each order_id to a single task, and stage 2 adds those partial sums into order_total. The 3 narrow steps still cost 1 pass over the input, and the groupBy is the only new boundary in the job.

How many stages does this job run, and how does stage 1 treat each row?

result = (user_events
    .filter(F.col("event_type") == "purchase")
    .withColumn("year", F.substring("event_date", 1, 4))
    .groupBy("user_id", "year")
    .agg(F.count(F.lit(1)).alias("purchases"))
    .withColumn("is_repeat", F.col("purchases") > 1)
    .collect())

• Five stages, 1 per transformation, each reading what the previous stage stored
• Two stages: in stage 1 each row is filtered and given year in 1 pass before the shuffle
• Two stages: stage 1 filters the whole partition, stores it, then adds year in a second pass
• Three stages, because the withColumn after the aggregation needs a stage of its own

To place a step you have not met before, ask what 1 output partition needs. A filter on a column, a withColumn with arithmetic and a select need only the row in hand, so they join the current stage. A distinct must compare a row against copies that may sit in any partition, and an orderBy must know where every row ranks, so each one ends the stage. When a wide step cannot be avoided, putting the narrow filters before it means the shuffle moves fewer rows.

This job has 1 shuffle. Which steps run in stage 1, before it?

result = (order_items
    .filter(F.col("quantity") >= 2)
    .withColumn("line_total", F.col("quantity") * F.col("unit_price"))
    .distinct()
    .select("order_id", "line_total")
    .collect())

• The read, the filter, line_total and select, since Spark moves every narrow step ahead of the wide one
• The read and the filter only, since withColumn computes a new value and so starts a stage
• Every step, since distinct only removes rows and moves none of them between tasks
• The read, the filter, line_total and the first half of distinct; select runs in stage 2

About this lesson

This section is part of the Lazy Evaluation: Intermediate 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.