PySpark如何对嵌套结构RDD执行Map操作实现共现元素聚合

实现思路 #

全程使用RDD分布式转换算子处理，全程不将全量数据拉取到Driver节点，适配数亿条记录的大规模场景，核心逻辑为：

• 遍历每条子列表记录，为列表内每个元素生成键值对：键为当前元素，值为同列表内其他所有元素
• 扁平化所有生成的键值对
• 按键分组，分组后每个键对应的value就是要求的ResultIterable类型可迭代对象

实现代码 #

基础版本（保留重复关联值，即两个元素共同出现在N个子列表中，关联列表里就会出现N次对应元素）：

desired = rdd.flatMap(
    lambda record: [(item, other) for item in record for other in record if other != item]
).groupByKey()

如果需要和示例输出一致，对关联元素去重（即两个元素无论共同出现在多少个子列表中，关联元素只保留1次），可以先对单条记录内元素去重再生成键值对，同时减少shuffle阶段传输的数据量，更适配大规模数据场景：

desired = rdd.flatMap(
    lambda record: 
        [(item, other) for item in set(record) for other in set(record) if other != item]
).groupByKey()

结果验证 #

执行测试代码：

for x in desired.collect()[0][1]:
    print(x)

会输出和示例完全一致的结果：

t5
t7
t3
t4

注意：数亿级数据场景下不建议直接调用collect()将全量结果拉取到Driver，否则会触发Driver端内存溢出，上述测试仅用于验证小批量样本结果正确性。groupByKey返回结果的value本身就是pyspark.resultiterable.ResultIterable类型，无需额外做类型转换。

内容的提问来源于stack exchange，提问作者abbassix