Partitions and Shuffles in Python

Spark cuts a DataFrame's rows into partitions: pieces it can work on side by side. Each core of each machine works on one partition at a time, so a table in 400 partitions keeps a 100-core cluster busy. df.rdd.getNumPartitions() tells you how many pieces a DataFrame has (rdd is Spark's older, lower-level view of the same rows), and df.repartition(4) returns the same rows spread over 4 pieces.

Some steps work on each partition alone: a select or a where handles each row where it is. Others need rows with the same key together: to count a category, all its lines must reach the same partition. Moving rows between partitions is a shuffle. On a cluster it writes rows to disk and sends them over the network, so it is usually the slowest step of a job: when a job is slow, a data engineer looks for its shuffles first.

df.explain() prints the plan Spark will run, without running it, and a shuffle shows in it as a step named Exchange. The plan also holds column ids and the folder's path, which change from run to run, so the starting code's plan_of(df) catches the text in a string (contextlib.redirect_stdout into an io.StringIO) instead of printing it. You test that string with in:

paid = orders.where(F.col("status") == "paid")
print("where shuffles:", "Exchange" in plan_of(paid))
print("in 2:", orders.repartition(2).rdd.getNumPartitions())

Here Spark runs on one core and orders.csv is a small file. Your program must print:

partitions: 1
after repartition: 4
select shuffles: False
groupBy shuffles: True

The groupBy's plan has an Exchange even with a single partition: Spark writes the plan it would run on a cluster, which is why the same code works on one laptop core or on a hundred machines.

Your task: below plan_of, print how many partitions orders has and how many orders.repartition(4) has, keep product and quantity in picked, count the lines of each category in lines, then print whether each one's plan contains Exchange. A Check takes about 10 seconds here, because Spark starts a Java program first.

Assignment
Follow the Coding Tutorial to count the partitions, build a select and a groupBy, and find the shuffle in their plans.

Hint
Each print takes a label and a value, like print("partitions:", …). "Exchange" in plan_of(…) is already True or False: no if needed.