Partitions and shuffles in Spark

What we're doing

Understand the difference between narrow and wide transformations, watch a shuffle happen in real time, and inspect partition sizes on a real DataFrame.

Step 1: What a partition is

A DataFrame is split into chunks called partitions. Each partition is a subset of the rows, held on one worker in memory.

The number of partitions matters a lot for performance:

• Too few partitions — not enough parallelism. Big partitions on few workers = wasted cluster capacity.
• Too many partitions — overhead of managing tasks dominates the actual work.

Step 2: Narrow vs wide transformations

Spark transformations fall into two categories based on whether they move data between partitions.

Narrow transformations — data stays in place. Each partition is processed independently. No communication between workers is needed.

Wide transformations — data needs to move across partitions. Rows from one partition might need to end up on a different partition to complete the operation. This is called a shuffle.

Step 3: Shuffles

A shuffle is when Spark moves data between workers to complete a wide transformation. It's the most expensive operation in Spark. Here's what actually happens during a shuffle:

1. Write phase — each executor writes its data to local disk, split into files by destination partition
2. Fetch phase — executors on other workers read those files over the network
3. Sort phase — data gets sorted by key on the receiving side
4. Compute phase — the actual operation (groupBy, join, etc.) runs on the reshuffled data

Step 4: Start PySpark and set up

pyspark --master spark://localhost:7077

Wait for the Python prompt. Create a DataFrame with a million rows to work with:

df = spark.range(1, 1000001)

spark.range generates a DataFrame with a single column called id containing numbers 1 through 1,000,000.

Step 5: Inspect the partitions

How many partitions does this DataFrame have?

df.rdd.getNumPartitions()

You can also see how many rows are in each partition:

from pyspark.sql.functions import spark_partition_id
df.groupBy(spark_partition_id()).count().show()

Step 6: A narrow transformation without shuffle

Let's run a filter and see what happens:

filtered = df.filter(df.id > 500000)
filtered.count()

Open the Spark Master UI and click into the application. Look at the job that ran. You'll see it has one stage. --

Step 7: A wide transformation with the shuffle

Now let's do something that triggers a shuffle:

from pyspark.sql.functions import col

grouped = df.groupBy((col("id") % 10).alias("bucket")).count()
grouped.show()

We're grouping by the last digit of the id.

Open the Spark Master UI again. Look at the new job. You'll see it has two stages now:

• Stage 1 — the pre-shuffle stage, running the grouping locally on each partition
• Stage 2 — the post-shuffle stage, combining the results across the network

Step 8: Trigger a bigger shuffle with a sort

Try this:

sorted_df = df.orderBy(df.id.desc())
sorted_df.show()

Sorting a DataFrame globally requires seeing all the data at once, which means a full shuffle. Look at the Spark UI, the sort job takes noticeably longer than the filter did.

Exit the session with exit().