Learn / Lesson 9 · 6 min

Where Your Code Runs

The One Row Per Group lesson's groupBy("size") is one line. On a real cluster it is also the most expensive line on the page, and filter, which looks the same size, costs almost nothing. The difference is not in the code. It is in where the rows are: your code describes work that gets split across machines, and some operations need those machines to talk to each other. This lesson is about telling the two kinds apart, with a tool you can run right here.

Watch it happen

Ask the engine what it plans to do before it does it:

from pyspark.sql import functions as F

summary = orders.groupBy("size").agg(F.sum("price").alias("revenue"))
summary.explain()
result = summary

The plan prints to the Console tab. Read it bottom up, the order the work happens in. Trimmed of its column bookkeeping, this is what Sail printed:

AggregateExec: mode=FinalPartitioned, gby=[size], aggr=[sum(price)]
  RepartitionExec: partitioning=Hash([size], 2), input_partitions=1
    AggregateExec: mode=Partial, gby=[size], aggr=[sum(price)]
      DataSourceExec: partitions=1

Bottom up: read the data; total up each size within each chunk of rows; move rows so that every row with the same size lands in the same place; finish the totals. The middle line is the shuffle, and this lesson is about that line.

The shuffle is the line in the middle of the plan: rows move so that every row with the same key is in one place, and that is the line that costs on a cluster.

One word

Data partition. A DataFrame is not one lump. It is split into chunks of rows, and each chunk is processed as a unit, by one task on one executor at a time. That chunk is a data partition. The plan above started from partitions=1, because eleven coffee orders fit in one; a real table has thousands, and the previous lesson's executors each hold some of them.

One warning for the window functions lesson: Window.partitionBy uses the same word for a different thing, a group of rows a window function treats together. We will always say data partition for the chunks and window partition for the groups.

Two kinds of operation

filter, select and withColumn each look at one row at a time. Every executor can do that to its own partitions without asking anyone. The plan shows it:

FilterExec: size = large
  DataSourceExec: partitions=1

No RepartitionExec. Rows stay where they are.

groupBy, distinct and join are different. To total up the large orders, every large order has to be in the same place. If they are spread across ten executors, nine have to send their large rows to the tenth. That movement, across the network, with the data written down and read back on the way, is a shuffle, and RepartitionExec: partitioning=Hash([size], 2) is its shape in the plan: hash each row's key, send the row where the hash says. You can see the effect directly:

result = (
    orders.repartition(3, "size")
    .withColumn("partition", F.spark_partition_id())
    .select("order_id", "size", "partition")
)

Every large and small row comes back marked partition 0; both medium rows come back marked 1. Same key, same place. That is what a shuffle buys, and moving the rows is what it costs.

After the shuffle, rows with the same key share a data partition; the two medium rows are the ones that moved, and that movement is the shuffle's cost.

distinct plans exactly like groupBy, because it is one: a group per value with nothing to compute. join plans as HashJoinExec: here Sail collects the smaller table whole and looks each order up in it, which a cluster calls a broadcast join.

What you will never feel here

Sail here runs in its local mode: one process on one machine. (Sail also has a cluster mode, with a driver and workers of its own; this site does not use it.) The plan shows you the shape of a shuffle, honestly, but partitions=1 means there is nothing to move and nobody to move it to. A groupBy on this site costs about what a filter costs. On a cluster, with a billion rows across a hundred executors, the same RepartitionExec line is the slowest thing that happens all day. This lesson can show you where the cost is. It cannot make you wait for it.

Three habits that follow

Shrink the data before it reaches a shuffle. Filter early, so fewer rows reach a groupBy or a join, and don't orderBy unless you need sorted output. A sort of everything is a sort of everything.

Let the small side be broadcast. When one table in a join is small, a cluster ships it whole to every executor and skips the shuffle; the joins lesson says when that happens on its own and when to ask for it.

Know the word "skew". A skewed job is one where a single key is far more common than the others, so one data partition holds most of the rows and one executor does nearly all the work while the rest wait. The symptom is a job that is almost finished for a very long time.

Nothing to grade here. The next lesson, joins, is the most shuffle-heavy thing you will do, and the window functions lessons end by calling .explain() on a window so you can see where the rows have to be. explain() can print more than the physical plan: Reading a Plan, in Part 2, reads all four.