Definition
- repartition(n) – distributes records round-robin into
nevenly sized partitions, ignoring column values. - repartition(cols) – applies hash-modulo on the given columns so all rows with the same key land in the same partition, using the default
spark.sql.shuffle.partitionscount (200). - repartition(n, cols) – hashes on the columns into exactly
npartitions, combining co-location with a controlled count. - All three are wide transformations that trigger a full shuffle with the same network cost.
Imagine the National Bank of India has a huge pile of transaction slips and wants to hand them to clerks. repartition(n) is like dealing the slips one-by-one to n clerks like cards, so every clerk gets the same number but Rajesh's slips may be scattered. repartition("accountId") says 'every slip for the same account must go to the same clerk,' using the bank's standard number of counters. repartition(n, "accountId") says both things: keep each account together AND use exactly n counters.
Understanding the Concepts
What are the three repartition variants?
All three variants are wide transformations that trigger a shuffle. Data is written to shuffle files on the map side, transferred across the network, and read on the reduce side. The cost is the same for every variant; only the partitioning strategy differs.
from pyspark.sql import SparkSession
import pyspark.sql.functions as F
spark = SparkSession.builder.appName("NBI-Repartition-Variants").getOrCreate()
txn = spark.createDataFrame([
("SAV2001", "Mumbai", 50000.00),
("SAV2001", "Mumbai", 75000.00),
("SAV2002", "Delhi", 12000.00),
("CUR3001", "Chennai", 300000.00),
("SAV2002", "Delhi", 8000.00),
("SAV2001", "Mumbai", 25000.00)
], ["account_id", "branch", "amount"])
How does repartition(n) work?
The first variant uses a round-robin partitioner. Spark assigns each record to the next partition in turn, wrapping around. The result is n evenly-sized partitions regardless of the data's key distribution. Because it ignores values entirely, rows with the same account can land in different partitions.
Use this when you only care about even balance and parallelism, such as spreading a skewed dataset evenly before a CPU-heavy map operation.
# repartition(n): round-robin, even distribution
even = txn.repartition(4)
print("Partitions:", even.rdd.getNumPartitions())
# Check: rows are spread evenly, same account may be split
for i, part in enumerate(even.rdd.glom().collect()):
accounts = [r[0] for r in part]
print(f" Partition {i}: {len(part)} rows, accounts: {accounts}")
Partitions: 4
Partition 0: 2 rows, accounts: ['SAV2001', 'CUR3001']
Partition 1: 1 rows, accounts: ['SAV2001']
Partition 2: 2 rows, accounts: ['SAV2002', 'SAV2001']
Partition 3: 1 rows, accounts: ['SAV2002']
How does repartition(cols) work?
The second variant uses a hash partitioner. Spark computes hash(column_value) % numPartitions and sends all rows with the same hash to the same partition. The number of partitions defaults to spark.sql.shuffle.partitions, which is 200 unless changed.
Use this when you want all rows for the same key in one partition so a downstream groupBy or join on that key avoids an extra shuffle.
# repartition(cols): hash-modulo co-location
by_account = txn.repartition("account_id")
print("Partitions:", by_account.rdd.getNumPartitions()) # 200 (default)
# All rows for same account are in same partition
for i, part in enumerate(by_account.rdd.glom().collect()):
if len(part) > 0:
accounts = [r[0] for r in part]
print(f" Partition {i}: {len(part)} rows, accounts: {accounts}")
Partitions: 200
Partition 42: 3 rows, accounts: ['SAV2001', 'SAV2001', 'SAV2001']
Partition 87: 2 rows, accounts: ['SAV2002', 'SAV2002']
Partition 155: 1 rows, accounts: ['CUR3001']
How does repartition(n, cols) work?
The third variant is the hash partitioner into a specific count you choose instead of the default 200. It gives you key co-location and control of the partition count.
Use this when you want co-location and a sensible file count, for instance to match the parallelism of a downstream join or to avoid creating 200 tiny files on write.
# repartition(n, cols): hash co-location into explicit count
by_account_4 = txn.repartition(4, "account_id")
print("Partitions:", by_account_4.rdd.getNumPartitions()) # exactly 4
# Same-key rows co-located, partition count controlled
for i, part in enumerate(by_account_4.rdd.glom().collect()):
accounts = [r[0] for r in part]
print(f" Partition {i}: {len(part)} rows, accounts: {accounts}")
Partitions: 4
Partition 0: 2 rows, accounts: ['SAV2002', 'SAV2002']
Partition 1: 0 rows, accounts: []
Partition 2: 3 rows, accounts: ['SAV2001', 'SAV2001', 'SAV2001']
Partition 3: 1 rows, accounts: ['CUR3001']
How do the three variants compare?
| Feature | repartition(n) | repartition(cols) | repartition(n, cols) |
|---|---|---|---|
| Strategy | Round-robin | Hash on columns | Hash on columns |
| Partition count | Explicit n | Default 200 | Explicit n |
| Balance | Even sizes | Uneven (depends on key) | Uneven (depends on key) |
| Key co-location | No | Yes | Yes |
| Use case | Pure balance | Co-locate by key | Co-locate + control count |
# Comparison: write output file counts
txn.repartition(4).write.mode("overwrite").parquet("/tmp/by_count") # 4 files
txn.repartition("account_id").write.mode("overwrite").parquet("/tmp/by_key") # 200 files!
txn.repartition(4, "account_id").write.mode("overwrite").parquet("/tmp/by_both") # 4 files
What are common mistakes?
The most common mistake is forgetting that repartition(cols) uses 200 partitions by default. This silently creates 200 output files when writing.
# WRONG: accidentally creating 200 files
# Forgot that repartition(cols) defaults to 200 partitions
txn.repartition("account_id").write.parquet("/tmp/oops") # 200 files!
# CORRECT: specify the count explicitly
txn.repartition(16, "account_id").write.parquet("/tmp/good") # 16 files
# WRONG: using repartition(n) when you need co-location
# Round-robin scatters same-key rows across partitions
txn.repartition(4).groupBy("account_id").agg(F.sum("amount")) # triggers another shuffle
# CORRECT: repartition by key to avoid downstream shuffle
txn.repartition(4, "account_id").groupBy("account_id").agg(F.sum("amount")) # no extra shuffle
When should you use each variant?
The choice depends on whether you want raw even balance, co-location by key, or both.
- repartition(n) – use when you only care about even parallelism. Example: spreading a skewed dataset before a CPU-heavy UDF.
- repartition(cols) – use when you want co-location by key and the default partition count is acceptable. Example: pre-grouping before a downstream groupBy.
- repartition(n, cols) – use when you want co-location and need to control the file or partition count. Example: writing co-located Parquet files for a downstream join.
# Use case 1: Balance before heavy UDF (repartition(n))
scored = txn.repartition(32).withColumn(
"risk", F.when(F.col("amount") > 100000, "HIGH").otherwise("LOW")
)
# Use case 2: Co-locate for downstream groupBy (repartition(cols))
pre_grouped = txn.repartition("account_id")
totals = pre_grouped.groupBy("account_id").agg(F.sum("amount"))
# Use case 3: Co-locate + control file count (repartition(n, cols))
txn.repartition(16, "account_id").write.parquet("/bank/by-account/")
How does hash co-location interact with downstream operations?
When you repartition by key, the downstream groupBy or join on that same key can skip the shuffle because Spark recognizes the data is already partitioned correctly. This is a significant optimization for banking pipelines that repeatedly aggregate by account.
# Pre-partition by account, then multiple aggregations without extra shuffles
by_acct = txn.repartition(16, "account_id")
# These groupBys on account_id can reuse the existing partitioning
total = by_acct.groupBy("account_id").agg(F.sum("amount").alias("total"))
max_txn = by_acct.groupBy("account_id").agg(F.max("amount").alias("max_amount"))
count = by_acct.groupBy("account_id").agg(F.count("*").alias("txn_count"))
total.show()
+----------+--------+
|account_id| total|
+----------+--------+
| SAV2001|150000.0|
| SAV2002| 20000.0|
| CUR3001|300000.0|
+----------+--------+
What should a senior engineer know?
A senior engineer states the rule concisely: a plain number gives balance, columns give co-location at the default count, and the combined form gives both into a count you control. All three cost a shuffle. The subtle detail is that hash partitioning does not guarantee even sizes: if one key has far more rows than others, its partition is large regardless. For truly even partitioning by key, you need salting or AQE.
# Senior tip: check partition balance after hash repartition
def partition_stats(df, label):
sizes = df.rdd.mapPartitions(lambda it: [sum(1 for _ in it)]).collect()
non_empty = [s for s in sizes if s > 0]
print(f"{label}: total_parts={len(sizes)}, non_empty={len(non_empty)}, "
f"min={min(non_empty)}, max={max(non_empty)}")
partition_stats(txn.repartition(4), "round-robin(4)")
partition_stats(txn.repartition(4, "account_id"), "hash(4, account_id)")
round-robin(4): total_parts=4, non_empty=4, min=1, max=2
hash(4, account_id): total_parts=4, non_empty=3, min=1, max=3
Visual Overview
Three repartition variants:
1. repartition(n) – Round-Robin
┌───────────────┐ ┌─────┐
│ All records │──RR──→│ P0 │ even
│ (any order) │──RR──→│ P1 │ even
│ │──RR──→│ P2 │ even
└───────────────┘ └─────┘
Keys scattered, sizes balanced
2. repartition(cols) – Hash on key, default 200
┌───────────────┐ ┌──────┐
│ hash(key) % 200│──→ │ P0 │ SAV2001 rows
│ │──→ │ P42 │ SAV2002 rows
│ │──→ │ P155 │ CUR3001 rows
└───────────────┘ │ ... │ (197 empty)
└──────┘
Keys co-located, 200 partitions (many empty)
3. repartition(n, cols) – Hash on key, explicit n
┌───────────────┐ ┌─────┐
│ hash(key) % 4 │──→ │ P0 │ SAV2002 rows
│ │──→ │ P1 │ (empty)
│ │──→ │ P2 │ SAV2001 rows
│ │──→ │ P3 │ CUR3001 rows
└───────────────┘ └─────┘
Keys co-located, count controlled
Key Takeaway
Repartition always performs a full shuffle, and it comes in three forms that differ by how records are placed. The first form, repartition with just a number, uses a round-robin partitioner that deals records evenly into exactly that many partitions. The result is perfectly balanced sizes, but it completely ignores the values, so rows that share a key get scattered. I reach for this when my only concern is even parallelism, for example after a heavy filter that left lopsided partitions before a CPU-intensive map.
The second form, repartition with one or more columns, uses a hash partitioner. Spark computes the hash of the key modulo the partition count and sends every row with the same key to the same partition. The important detail that trips people up is that this form uses the default shuffle partition count, which is two hundred, so it can silently produce two hundred partitions and two hundred output files. I use this when I want co-location by key so a downstream groupBy or join avoids an extra shuffle.
The third form, repartition with both a number and columns, is the hash partitioner into an explicit count. It gives me key co-location and control of the partition count at the same time, which is what I want when I need the output to align with a consumer or to avoid too many small files.
In my project at the National Bank of India I used repartition with a plain number to rebalance a skewed daily slice before fraud scoring, repartition on the account column to pre-group transactions so the per-account totals avoided an extra shuffle, and repartition with a count plus the account column to produce exactly sixteen co-located partitions that matched a downstream join.
Remember one sentence: a plain number gives balance, columns give co-location at the default count, and the combined form gives both into a count you control, and all three cost a shuffle.
---