pyspark的join.py代码记录

pyspark RDD中join算子实现代码分析

代码版本为 spark 2.2.0

1.join.py

这个代码单独作为一个文件在pyspark项目代码中,只有一个功能即实现join相关的几个方法

# 这个方法是所有join动作的基础方法,实现了join的基本需求
# 通过mapValues的方法把所有的rdd value都变成一个(rddNumber,rdd.value)的元组
# 通过转换之后所有的value都变成了相同的格式，然后再用union联合起来
# 之后再通过groupByKey的方式把所有key分组到相同的分区,然后再把value联合起来
def _do_python_join(rdd, other, numPartitions, dispatch):
    vs = rdd.mapValues(lambda v: (1, v))   # 转换成相同的格式
    ws = other.mapValues(lambda v: (2, v)) # 转换成相同的格式
    # union联合起来,在groupByKey分组到一起
    return vs.union(ws).groupByKey(numPartitions).flatMapValues(lambda x: dispatch(x.__iter__()))

# 如果有join的需求可以使用类似这种方式实现,性能应该会好一点
# 这是个人添加的示例代码,并不在源代码文件中
def _do_self_join(rdd, other, numPartitions, dispatch):
    V=namedtuple("V",["v1","v2"])
    vs = rdd.mapValues(lambda v: V(v,None))
    ws = other.mapValues(lambda v: V(None,v))
    agg = lambda v1,v2: V(v1.v1,v2.v2) if v1.v1 else V(v2.v1,v1.v2)
    return vs.union(ws).reduceByKey(agg)

# 这个方法被RDD.join算子调用
def python_join(rdd, other, numPartitions):
    # 这里构造了一个把value放在同一个元组中的函数
    def dispatch(seq):
        vbuf, wbuf = [], []
        for (n, v) in seq:
            if n == 1:
                vbuf.append(v)
            elif n == 2:
                wbuf.append(v)
        return ((v, w) for v in vbuf for w in wbuf)
    # 调用了上面的函数
    return _do_python_join(rdd, other, numPartitions, dispatch)


# 右join方法,跟join方法实现的方式相同,只是聚合value的方法略有区别
def python_right_outer_join(rdd, other, numPartitions):
    def dispatch(seq):
        vbuf, wbuf = [], []
        for (n, v) in seq:
            if n == 1:
                vbuf.append(v)
            elif n == 2:
                wbuf.append(v)
        if not vbuf:
            vbuf.append(None)
        return ((v, w) for v in vbuf for w in wbuf)
    return _do_python_join(rdd, other, numPartitions, dispatch)


def python_left_outer_join(rdd, other, numPartitions):
    def dispatch(seq):
        vbuf, wbuf = [], []
        for (n, v) in seq:
            if n == 1:
                vbuf.append(v)
            elif n == 2:
                wbuf.append(v)
        if not wbuf:
            wbuf.append(None)
        return ((v, w) for v in vbuf for w in wbuf)
    return _do_python_join(rdd, other, numPartitions, dispatch)


def python_full_outer_join(rdd, other, numPartitions):
    def dispatch(seq):
        vbuf, wbuf = [], []
        for (n, v) in seq:
            if n == 1:
                vbuf.append(v)
            elif n == 2:
                wbuf.append(v)
        if not vbuf:
            vbuf.append(None)
        if not wbuf:
            wbuf.append(None)
        return ((v, w) for v in vbuf for w in wbuf)
    return _do_python_join(rdd, other, numPartitions, dispatch)


def python_cogroup(rdds, numPartitions):
    def make_mapper(i):
        return lambda v: (i, v)
    vrdds = [rdd.mapValues(make_mapper(i)) for i, rdd in enumerate(rdds)]
    union_vrdds = reduce(lambda acc, other: acc.union(other), vrdds)
    rdd_len = len(vrdds)

    def dispatch(seq):
        bufs = [[] for _ in range(rdd_len)]
        for n, v in seq:
            bufs[n].append(v)
        return tuple(ResultIterable(vs) for vs in bufs)

    return union_vrdds.groupByKey(numPartitions).mapValues(dispatch)