Definition
- groupByKey – shuffles every value for each key to the partition owning that key, then lets you aggregate the full collection.
- reduceByKey – combines values locally on each partition first using an associative reduce function, then shuffles only partial results.
- Shuffle volume – groupByKey shuffles the full raw data; reduceByKey shuffles one partial result per key per partition.
- Memory pressure – groupByKey loads all values for a key onto one partition, risking OOM on hot keys.
- DataFrame equivalent – df.groupBy().sum() already does map-side pre-aggregation automatically.
At the National Bank of India, suppose each branch has a pile of transaction amounts per account and head office wants the total per account. With groupByKey every branch ships every single slip to head office, which then adds them up – a mountain of paper travels. With reduceByKey each branch first totals its own slips and ships only one subtotal – almost no paper travels.
Understanding the Concepts
What is the difference between groupByKey and reduceByKey?
Both are key-value RDD transformations that group data by key, but they differ critically in where aggregation happens. This difference dominates performance.
from pyspark.sql import SparkSession
spark = SparkSession.builder.appName("GroupReduceDemo").getOrCreate()
sc = spark.sparkContext
# Transaction amounts per account at the National Bank of India
txns = sc.parallelize([
("SAV2001", 50000), ("SAV2001", 25000), ("SAV2001", 75000),
("CUR3001", 100000), ("CUR3001", 50000),
("SAV2001", 30000), ("CUR3001", 80000),
])
How does groupByKey work?
groupByKey shuffles every value to the partition owning the key, then gives you the full collection. The entire raw volume crosses the network.
# groupByKey: shuffles ALL values, then you aggregate
grouped = txns.groupByKey()
for key, values in grouped.collect():
total = sum(values)
print(f"{key}: values shipped = {list(values)}, total = {total}")
# Every single amount value crosses the network!
# SAV2001: 4 values shipped to one partition
# CUR3001: 3 values shipped to one partition
How does reduceByKey work?
reduceByKey takes an associative function and applies it locally on each partition first. Only the partial results are shuffled.
# reduceByKey: pre-aggregates locally, then shuffles partial results
totals = txns.reduceByKey(lambda a, b: a + b)
for key, total in totals.collect():
print(f"{key}: total = {total}")
# On partition 0: SAV2001 has 50000 + 25000 = 75000 (local)
# On partition 1: SAV2001 has 75000 + 30000 = 105000 (local)
# Shuffle: only 75000 and 105000 are sent, not all 4 raw values
# Final: 75000 + 105000 = 180000
How much data does each shuffle?
The shuffle volume difference can be enormous, especially when many values share the same key.
groupByKey shuffle volume:
Partition 0: SAV2001=50000, SAV2001=25000, CUR3001=100000
Partition 1: SAV2001=75000, SAV2001=30000, CUR3001=50000, CUR3001=80000
All 7 values cross the network -> 7 records shuffled
reduceByKey shuffle volume:
Partition 0: SAV2001=75000, CUR3001=100000 (pre-aggregated)
Partition 1: SAV2001=105000, CUR3001=130000 (pre-aggregated)
Only 4 partial results cross the network -> 4 records shuffled
At scale with millions of values per key:
groupByKey: shuffles millions of records
reduceByKey: shuffles one partial result per key per partition
Why does groupByKey cause memory pressure?
groupByKey loads all values for a key into memory on the destination partition. A hot key with millions of values can cause OOM.
# Simulate a hot key: one corporate account with 10M transactions
hot_txns = sc.parallelize(
[("CORP_POOL", i) for i in range(10000000)] +
[("SAV2001", 50000), ("CUR3001", 75000)]
)
# groupByKey: all 10M values for CORP_POOL go to ONE partition
# This can cause java.lang.OutOfMemoryError!
try:
hot_txns.groupByKey().mapValues(sum).collect()
except Exception as e:
print(f"OOM: {e}")
# reduceByKey: pre-aggregates on each partition, no OOM risk
result = hot_txns.reduceByKey(lambda a, b: a + b).collect()
print(result) # Works fine
What about the DataFrame API?
In the DataFrame API, groupBy with aggregate functions like sum and count already does map-side pre-aggregation automatically. You get reduceByKey behavior without thinking about it.
from pyspark.sql.functions import col, sum as spark_sum
txns_df = spark.createDataFrame([
("SAV2001", 50000), ("SAV2001", 25000), ("SAV2001", 75000),
("CUR3001", 100000), ("CUR3001", 50000),
], ["account_id", "amount"])
# This automatically does map-side pre-aggregation (partial_sum)
totals = txns_df.groupBy("account_id").agg(spark_sum("amount"))
totals.explain() # Shows HashAggregate with partial aggregation
== Physical Plan ==
*(2) HashAggregate(keys=[account_id], functions=[sum(amount)])
+- Exchange hashpartitioning(account_id, 200)
+- *(1) HashAggregate(keys=[account_id], functions=[partial_sum(amount)])
+- ... <-- partial_sum is the map-side pre-aggregation!
What is the common anti-pattern?
The most frequent anti-pattern is groupByKey followed by mapValues to sum, which shuffles everything when reduceByKey would shuffle almost nothing.
# WRONG: groupByKey + mapValues sum (shuffles all raw values)
bad = txns.groupByKey().mapValues(sum)
# CORRECT: reduceByKey (shuffles only partial results)
good = txns.reduceByKey(lambda a, b: a + b)
# Both produce the same result, but reduceByKey shuffles
# orders of magnitude less data
When is groupByKey appropriate?
groupByKey is appropriate only when you genuinely need the full collection of values per key and cannot express the operation as an incremental reduction.
# Legitimate use: need all transaction ids per account for audit
audit_ids = txns.map(lambda x: (x[0], str(x[1]))) \
.groupByKey() \
.mapValues(list)
# Cannot reduce transaction ids into a single value incrementally
# For this use case, manage hot-key risk:
# - Bound the per-key collection with .take(N)
# - Salt the hot key to spread the load
What should a senior engineer know?
- Always prefer reduceByKey, aggregateByKey, or combineByKey over groupByKey for associative aggregations.
- The DataFrame API handles this automatically with partial aggregation.
- At the National Bank of India, rewriting groupByKey as reduceByKey cut shuffle bytes several-fold and eliminated OOM on a corporate account with millions of transactions.
- The typed Dataset API's groupByKey can reintroduce full-shuffle behavior if followed by a non-combining operation, so watch for it.
- Spotting and rewriting the groupByKey-followed-by-sum pattern is a frequent, high-impact optimization.
Visual Overview
groupByKey: shuffles ALL raw values
┌──────────────┐ ┌──────────────┐
│ Partition 0 │ │ Partition 1 │
│ A=50K, A=25K │ │ A=75K, A=30K │
│ B=100K │ │ B=50K, B=80K │
└──────┬───────┘ └──────┬───────┘
│ ALL values │
└────────┬─────────┘
▼
┌──────────────────────────┐
│ Key A: [50K, 25K, 75K, 30K] │ Then sum = 180K
│ Key B: [100K, 50K, 80K] │ Then sum = 230K
└──────────────────────────┘
reduceByKey: pre-aggregates, shuffles partial results
┌──────────────┐ ┌──────────────┐
│ Partition 0 │ │ Partition 1 │
│ A=50K+25K=75K │ │ A=75K+30K=105K│ Local combine
│ B=100K │ │ B=50K+80K=130K│
└──────┬───────┘ └──────┬───────┘
│ PARTIAL results │
└────────┬─────────┘
▼
┌──────────────────────────┐
│ Key A: 75K + 105K = 180K │ Only partials shuffled!
│ Key B: 100K + 130K = 230K│
└──────────────────────────┘
Key Takeaway
groupByKey and reduceByKey both group key-value data, but they differ in where aggregation happens. groupByKey shuffles every raw value for a key across the network to the single partition owning that key, then lets you aggregate. reduceByKey combines values locally on each partition first, then shuffles only the partial results. Because the shuffle is the dominant cost, reduceByKey moves far less data and is almost always preferred for associative aggregations.
groupByKey also creates memory pressure. All values for a key pile into one partition, so a hot key with millions of values risks out-of-memory. reduceByKey avoids this because no partition ever holds a key's entire raw value set. In the DataFrame API, groupBy with sum or count already does map-side pre-aggregation automatically, so the distinction is mainly an RDD-level concern.
At the National Bank of India, computing per-account totals with reduceByKey instead of groupByKey cut the shuffle to a fraction and eliminated an OOM failure on a corporate account. The common anti-pattern to watch for is groupByKey followed by mapValues to sum, which shuffles everything when reduceByKey would shuffle almost nothing.
The only time to use groupByKey is when you genuinely need the full collection of values per key and cannot express the operation as an incremental reduction, and then you manage the hot-key risk explicitly.
Remember one sentence: reduceByKey pre-aggregates locally before the shuffle so it moves far less data than groupByKey, which ships all raw values across the network.
---