Spark aggregateByKey()

aggregateByKey() is the most flexible key-based aggregation in the Spark RDD API. It lets you define two separate combining functions: one for combining values within a partition and one for combining the partition subtotals across partitions. This enables aggregations that reduceByKey can’t express — like computing averages, collecting into sets, or tracking multiple statistics simultaneously.

Syntax

rdd.aggregateByKey(
    zeroValue,          # Starting accumulator for each key in each partition
    seqFunc,            # (accumulator, value) → accumulator  [within partition]
    combFunc,           # (accumulator, accumulator) → accumulator  [across partitions]
    numPartitions=None  # Optional: output partition count
)

Example 1: Computing Average (Cannot Use reduceByKey)

from pyspark.sql import SparkSession

spark = SparkSession.builder.appName("aggregateByKey").getOrCreate()
sc = spark.sparkContext

scores = sc.parallelize([
    ("Alice", 85), ("Bob", 72), ("Alice", 91),
    ("Bob", 88),   ("Alice", 79), ("Bob", 65),
])

# Accumulator: (sum, count)
zero = (0, 0)

def seq_func(acc, value):
    return (acc[0] + value, acc[1] + 1)

def comb_func(acc1, acc2):
    return (acc1[0] + acc2[0], acc1[1] + acc2[1])

sum_count = scores.aggregateByKey(zero, seq_func, comb_func)
averages  = sum_count.mapValues(lambda sc: sc[0] / sc[1])
averages.collect()
# [("Alice", 85.0), ("Bob", 75.0)]

Example 2: Collect Unique Values per Key

purchases = sc.parallelize([
    ("C001", "Laptop"), ("C002", "Mouse"),  ("C001", "Monitor"),
    ("C001", "Laptop"), ("C002", "Keyboard"), ("C002", "Mouse"),
])

unique_products = purchases.aggregateByKey(
    set(),                                        # zero: empty set per key per partition
    lambda acc, v: acc | {v},                     # seqFunc: add value to set
    lambda a, b: a | b                            # combFunc: merge sets
)
unique_products.collect()
# [("C001", {"Laptop", "Monitor"}), ("C002", {"Mouse", "Keyboard"})]

Example 3: Min and Max Simultaneously

temperatures = sc.parallelize([
    ("NYC", 22), ("LA", 28), ("NYC", 18), ("LA", 31), ("NYC", 25), ("LA", 24),
])

stats = temperatures.aggregateByKey(
    (float("inf"), float("-inf")),                         # zero: (min, max)
    lambda acc, v: (min(acc[0], v), max(acc[1], v)),       # seqFunc
    lambda a, b: (min(a[0], b[0]), max(a[1], b[1]))        # combFunc
)
stats.collect()
# [("NYC", (18, 25)), ("LA", (24, 31))]

Example 4: Count with Custom Logic (Conditional Count)

clicks = sc.parallelize([
    ("ad_001", {"clicked": True,  "converted": False}),
    ("ad_001", {"clicked": True,  "converted": True}),
    ("ad_002", {"clicked": True,  "converted": False}),
    ("ad_001", {"clicked": False, "converted": False}),
    ("ad_002", {"clicked": True,  "converted": True}),
])

metrics = clicks.aggregateByKey(
    (0, 0),   # (impressions, conversions)
    lambda acc, v: (acc[0] + 1, acc[1] + (1 if v["converted"] else 0)),
    lambda a, b: (a[0] + b[0], a[1] + b[1])
)
cvr = metrics.mapValues(lambda m: (m[0], m[1], m[1] / m[0] if m[0] > 0 else 0))
cvr.collect()
# [("ad_001", (3, 1, 0.333)), ("ad_002", (2, 1, 0.5))]

aggregateByKey vs Other Aggregations

Method	Use When
`reduceByKey`	Simple aggregation (sum, max, min) — same type in and out
`groupByKey`	Need full list of values — use sparingly for large data
`aggregateByKey`	Different input/output types, compound accumulators (avg, set, stats)
`combineByKey`	Most flexible — custom `createCombiner`, `mergeValue`, `mergeCombiner`
DataFrame `groupBy().agg()`	Structured data — almost always preferred for DataFrames