AggregateByKey

Overview

Learn how to use aggregateByKey to aggregate Pair RDD values using different functions for local combining and cross-partition merging.

What You Will Learn

Detailed Concept Explanation

aggregateByKey is a powerful wide transformation on Pair RDDs. It is used when the output value of your aggregation has a different type than the input values (for example, taking a list of integers and calculating the average—which requires tracking both sum and count as a tuple).

Parameters

It takes three parameters:

1. zeroValue: The starting value for the aggregation (e.g. (0, 0)).
2. seqOp: The function used to aggregate values locally on each partition (map-side combining).
3. combOp: The function used to merge the aggregates from different partitions over the network.

Code Examples

Input Dataset Preview

Below is the sales scores dataset:

Python (PySpark) Implementation

from pyspark.sql import SparkSession

spark = SparkSession.builder.appName("AggregateByKey").getOrCreate()
sc = spark.sparkContext
pair_rdd = sc.parallelize([("A", 10), ("B", 15), ("A", 20)])

# Calculate sum and count for each key
zero_val = (0, 0)
seq_op = lambda acc, value: (acc[0] + value, acc[1] + 1)
comb_op = lambda acc1, acc2: (acc1[0] + acc2[0], acc1[1] + acc2[1])

aggregated = pair_rdd.aggregateByKey(zero_val, seq_op, comb_op)
print(aggregated.collect())

Expected Output

[('A', (30, 2)), ('B', (15, 1))]

Execution Plan Diagram (Python & Scala)

Scala Implementation

import org.apache.spark.sql.SparkSession

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

val pairRDD = sc.parallelize(Seq(("A", 10), ("B", 15), ("A", 20)))

val zeroVal = (0, 0)
val seqOp = (acc: (Int, Int), value: Int) => (acc._1 + value, acc._2 + 1)
val combOp = (acc1: (Int, Int), acc2: (Int, Int)) => (acc1._1 + acc2._1, acc1._2 + acc2._2)

val aggregated = pairRDD.aggregateByKey(zeroVal)(seqOp, combOp)
println(aggregated.collect().mkString(", "))

Expected Output

(A,(30,2)), (B,(15,1))

Common Mistakes

• Zero Value Mutation: Designing zero values that mutate inside the seqOp. Always return a new tuple/object to prevent partition memory corruption.

Interview Perspective

Interactive Challenges

Related Topics

• ReduceByKey
• CombineByKey
• Pair RDDs