当前位置：首页 > article >正文

Spark RDD各种join算子从源码层分析实现方式

article 2024/11/20 10:01:09

在 Spark RDD 中，join、leftOuterJoin、rightOuterJoin、fullOuterJoin 等多个 Join 操作符都使用了 cogroup 进行底层实现。cogroup 是 Spark 中的一种底层分组操作，可以将两个或多个 RDD 中同一键的数据分组到一起，为各种 Join 操作提供了基础。下面我们从源码实现角度来分析这些 Join 操作符的实现原理，并列出相关的核心代码。

1. `join`

join 是最常用的连接操作，它会返回两个 RDD 中键相同的元素对。join 操作底层依赖 cogroup 实现：

def join[W](other: RDD[(K, W)], numPartitions: Int): RDD[(K, (V, W))] = {
  this.cogroup(other, numPartitions).flatMapValues {
    case (vs, ws) => for (v <- vs.iterator; w <- ws.iterator) yield (v, w)
  }
}

实现过程：

cogroup 将两个 RDD 中相同的键的数据组合成 (K, (Iterable[V], Iterable[W])) 形式。
然后通过 flatMapValues 遍历所有相同键的值对 (v, w)，形成最终的 (K, (V, W)) 格式。

示例：

rdd1 = sc.parallelize([("a", 1), ("b", 2)])
rdd2 = sc.parallelize([("a", 3), ("b", 4)])
joined = rdd1.join(rdd2)
print(joined.collect())
# 输出: [('a', (1, 3)), ('b', (2, 4))]

2. `leftOuterJoin`

leftOuterJoin 会返回左侧 RDD 的所有键，并对右侧 RDD 中匹配的键进行连接，未匹配到的则返回 None。

def leftOuterJoin[W](other: RDD[(K, W)], numPartitions: Int): RDD[(K, (V, Option[W]))] = {
  this.cogroup(other, numPartitions).flatMapValues {
    case (vs, ws) =>
      if (ws.isEmpty) vs.iterator.map(v => (v, None))
      else for (v <- vs.iterator; w <- ws.iterator) yield (v, Some(w))
  }
}

实现过程：

使用 cogroup 分组得到 (K, (Iterable[V], Iterable[W])) 格式。
对于左侧 RDD 的值 vs，如果右侧 RDD 的 ws 为空，则返回 None；否则返回 (v, Some(w))。

示例：

rdd1 = sc.parallelize([("a", 1), ("b", 2)])
rdd2 = sc.parallelize([("a", 3)])
left_joined = rdd1.leftOuterJoin(rdd2)
print(left_joined.collect())
# 输出: [('a', (1, Some(3))), ('b', (2, None))]

3. `rightOuterJoin`

rightOuterJoin 返回右侧 RDD 的所有键，未匹配的左侧 RDD 值则为 None。

def rightOuterJoin[W](other: RDD[(K, W)], numPartitions: Int): RDD[(K, (Option[V], W))] = {
  this.cogroup(other, numPartitions).flatMapValues {
    case (vs, ws) =>
      if (vs.isEmpty) ws.iterator.map(w => (None, w))
      else for (v <- vs.iterator; w <- ws.iterator) yield (Some(v), w)
  }
}

实现过程：

cogroup 将两个 RDD 分组。
对右侧 RDD 中的值 ws，如果左侧 RDD 中 vs 为空，则返回 None；否则返回 (Some(v), w)。

示例：

rdd1 = sc.parallelize([("a", 1)])
rdd2 = sc.parallelize([("a", 3), ("b", 4)])
right_joined = rdd1.rightOuterJoin(rdd2)
print(right_joined.collect())
# 输出: [('a', (Some(1), 3)), ('b', (None, 4))]

4. `fullOuterJoin`

fullOuterJoin 会返回左右两个 RDD 的所有键。未匹配的键对应的值为 None。

def fullOuterJoin[W](other: RDD[(K, W)], numPartitions: Int): RDD[(K, (Option[V], Option[W]))] = {
  this.cogroup(other, numPartitions).flatMapValues {
    case (vs, ws) =>
      if (vs.isEmpty) ws.iterator.map(w => (None, Some(w)))
      else if (ws.isEmpty) vs.iterator.map(v => (Some(v), None))
      else for (v <- vs.iterator; w <- ws.iterator) yield (Some(v), Some(w))
  }
}

实现过程：

使用 cogroup 进行键的分组。
对于每个键的 vs 和 ws，分别判断它们是否为空，如果为空则返回 None。

示例：

rdd1 = sc.parallelize([("a", 1), ("b", 2)])
rdd2 = sc.parallelize([("a", 3), ("c", 4)])
full_joined = rdd1.fullOuterJoin(rdd2)
print(full_joined.collect())
# 输出: [('a', (Some(1), Some(3))), ('b', (Some(2), None)), ('c', (None, Some(4)))]