Shuffling Data: When you call
repartition(256), Spark performs a
full shuffle of the data across the specified number of partitions.
The goal is to redistribute the data evenly across all partitions, which allows parallel processing across multiple nodes.
Default Repartitioning: If no column is specified, Spark randomly assigns rows to partitions.
The intention is to balance the number of rows across the 256 partitions for efficiency.
Even Distribution: Spark attempts to distribute the data as evenly as possible among the partitions.
While it doesn't guarantee perfectly equal splits, it ensures a reasonable balance.
Column-based Repartitioning
To repartition based on a specific column, you can use the following syntax:
In this case, Spark redistributes the data based on the values in the category column.
Rows with the same value in the category column will be placed in the same partition.
This is especially useful for operations like joins or aggregations that depend on column-based grouping.
Key Points
Without a Column: When using repartition(256) without specifying a column, Spark redistributes the data across 256 partitions without specific logic, balancing the rows across partitions.
With a Column: Using repartition(256, col("some_column")), Spark repartitions the data based on the column values, optimizing for operations like grouping or joins.
When to Use Repartition
Load Balancing: Repartitioning helps balance data across partitions when the data is skewed, ensuring that no partition has an oversized share of the data.
Parallel Processing: By redistributing the data across more partitions, repartitioning ensures efficient parallel processing, especially when working with a large cluster.
Alternatives
coalesce(n): Unlike repartition(),
coalesce(n) reduces the number of partitions without a full shuffle.
It's typically used after filtering operations or when you're writing out smaller data.
Example
# Repartition based on a specific column for optimized grouping
repartitioned_data = raw_data.repartition(256, col("customer_id"))
In this example, Spark repartitions the data so that rows with the same customer_id are placed in the same partition, which can be useful for later operations like grouping or joining on that column.
What a Shuffle Is
A shuffle is the step where Spark moves rows between executors so that rows sharing a key end up together. It is the boundary between two stages of a job. On the map side, each task computes a target partition for every row (for column-based repartitioning, a hash of the column value modulo the number of partitions) and writes sorted, partitioned shuffle files to local disk. On the reduce side, each new task fetches its piece from every map task over the network. Because it involves serialisation, disk writes and an all-to-all network exchange, a shuffle is usually the most expensive operation in a Spark job, and the Spark UI shows it as "Shuffle Write" and "Shuffle Read".
Joins, groupBy, distinct, window functions and repartition all trigger shuffles. Calling repartition(n, col) right before a groupBy(col) on the same column does not add work: Spark recognises the existing hash partitioning and skips a second exchange.
repartition vs coalesce
from pyspark.sql import SparkSession, functions as F
spark = SparkSession.builder.getOrCreate()
df = spark.range(0, 10_000_000).withColumn("category", (F.col("id") % 7).cast("string"))
print(df.rdd.getNumPartitions()) # depends on cores
by_cat = df.repartition(16, "category")
by_cat.groupBy("category").count().explain() # one Exchange, reused
# Rows per partition: shows skew (7 keys into 16 partitions leaves some empty)
(by_cat.withColumn("pid", F.spark_partition_id())
.groupBy("pid").count().orderBy("pid").show())
# Range partitioning: sorted, contiguous key ranges per partition
df.repartitionByRange(8, "id").write.mode("overwrite").parquet("/tmp/by_range")
# Shrink file count for a small output without a shuffle
df.filter("id < 1000").coalesce(1).write.mode("overwrite").parquet("/tmp/small")
The partition-id histogram makes an important point: repartitioning by a low-cardinality column cannot produce more non-empty partitions than there are distinct values. Seven categories spread over 16 partitions leave at least nine empty, and a dominant category creates one oversized partition that becomes the straggler task.
Tuning the Number of Partitions
The default number of post-shuffle partitions comes from spark.sql.shuffle.partitions (200 by default). With Adaptive Query Execution enabled (the default in Spark 3.2 and later and on Databricks), Spark coalesces small shuffle partitions at run time and can split skewed ones in joins, so hand-tuned repartition calls are needed less often. A common target is partitions of roughly 100 to 200 MB in memory, and at least two to three times as many partitions as total executor cores so all cores stay busy.
Operation
Shuffle?
Use for
repartition(n)
yes, round-robin
even sizes, increase parallelism
repartition(n, col)
yes, hash
co-locate keys before joins, per-key writes
repartitionByRange(n, col)
yes, sampled ranges
sorted output files, range queries
coalesce(n)
no
reduce file count after filtering
DataFrameWriter.partitionBy(col)
no extra
on-disk directory layout for pruning
Pitfalls
Skewed keys. A hot value (null, "unknown", a dominant customer) sends most rows to one task. Salt the key or rely on AQE skew join handling.
coalesce(1) on big data. It pulls the whole upstream stage into one task; use it only on small results.
Repartitioning too early. A shuffle before a filter moves rows that are then discarded; filter and project first.
Confusing in-memory and on-disk partitions.repartition does not create directories; partitionBy on the writer does.