Spark RDD中常用聚合算子源码层面的对比分析

news2024/11/20 1:27:47

在 Spark RDD 中，groupByKey、reduceByKey、foldByKey 和 aggregateByKey 是常用的聚合算子，适用于按键进行数据分组和聚合。它们的实现方式各不相同，涉及底层调用的函数也有区别。以下是对这些算子在源码层面的分析，以及每个算子适用的场景和代码示例。

1. `groupByKey`

功能：将相同键的值分组，形成一个 (key, Iterable<values>) 的 RDD。

源码分析：
groupByKey 底层使用了 combineByKeyWithClassTag 方法进行数据分组。

def groupByKey(): RDD[(K, Iterable[V])] = {
    combineByKeyWithClassTag(
      (v: V) => mutable.ArrayBuffer(v),
      (c: mutable.ArrayBuffer[V], v: V) => { c += v; c },
      (c1: mutable.ArrayBuffer[V], c2: mutable.ArrayBuffer[V]) => { c1 ++= c2; c1 }
    ).asInstanceOf[RDD[(K, Iterable[V])]]
}

适用场景：适合需要按键分组、无聚合的场景，但由于需要把所有键的值都传输到驱动端，数据量大时可能导致内存问题。

示例：

rdd = sc.parallelize([("a", 1), ("b", 2), ("a", 3)])
result = rdd.groupByKey().mapValues(list)
print(result.collect())

输出：[('a', [1, 3]), ('b', [2])]

2. `reduceByKey`

功能：基于给定的二元函数（如加法）对每个键的值进行聚合。
源码分析：
reduceByKey 底层也是基于 combineByKeyWithClassTag 方法进行处理，但与 groupByKey 不同的是，它在每个分区内执行局部聚合，再进行全局聚合，减少了数据传输。
```
def reduceByKey(func: (V, V) => V): RDD[(K, V)] = {
    combineByKeyWithClassTag[V]((v: V) => v, func, func)
}
```
- 适用场景：适用于需要对数据进行聚合计算的场景，能够有效减少 shuffle 数据量。

示例：

rdd = sc.parallelize([("a", 1), ("b", 2), ("a", 3)])
result = rdd.reduceByKey(lambda x, y: x + y)
print(result.collect())

输出：[('a', 4), ('b', 2)]

3. `foldByKey`

功能：与 reduceByKey 类似，但提供了初始值，分区内和分区间合并时都使用这个初始值。
源码分析：
foldByKey 的实现中调用了 aggregateByKey 方法，初始值会在每个分区中传递，确保聚合逻辑一致。
```
def foldByKey(zeroValue: V)(func: (V, V) => V): RDD[(K, V)] = {
    aggregateByKey(zeroValue)(func, func)
}
```
- 适用场景：当聚合操作需要一个初始值时使用，如从初始值开始累积计算。

示例：

rdd = sc.parallelize([("a", 1), ("b", 2), ("a", 3)])
result = rdd.foldByKey(0, lambda x, y: x + y)
print(result.collect())

输出：[('a', 4), ('b', 2)]

4. `aggregateByKey`

功能：支持更复杂的聚合操作，提供了分区内和分区间不同的聚合函数。
源码分析：
aggregateByKey 是最通用的聚合算子，调用了 combineByKeyWithClassTag 方法来控制分区内和分区间的计算方式。
```
def aggregateByKey[U: ClassTag](zeroValue: U)(
    seqOp: (U, V) => U,
    combOp: (U, U) => U): RDD[(K, U)] = {
    // Implementation detail here
}
```
- 适用场景：适合复杂的聚合逻辑需求，例如在分区内和分区间使用不同的函数。

示例：

rdd = sc.parallelize([("a", 1), ("b", 2), ("a", 3)])
result = rdd.aggregateByKey(0,
                            lambda x, y: x + y,   # 分区内加和
                            lambda x, y: x + y)   # 分区间加和
print(result.collect())