PySpark合并数组列中含至少一个共同值的所有行

问题描述 #

原始Spark DataFrame df 如下：

+------------+
|      values|
+------------+
|      [a, b]|
|[a, b, c, d]|
|   [a, e, f]|
|   [w, x, y]|
|      [x, z]|
+------------+

需要将所有至少有一个共同元素的数组合并，得到如下结果：

+-------------------+
|             values|
+-------------------+
| [a, b, c, d, e, f]|
|       [w, x, y, z]|
+-------------------+

之前尝试的代码仅能过滤掉被完全包含的子集数组（比如[a,b]被[a,b,c,d]包含），但无法处理有交集但非子集的情况（比如[a,e,f]和[a,b,c,d]），也无法合并[w,x,y]和[x,z]这类有共同元素的数组，因此输出不符合预期。

解决方案：基于连通分量的合并 #

这个问题本质是连通分量识别：将数组中的每个元素视为节点，同一数组内的元素互相建立连接，所有连通的节点即为需要合并的元素集合。我们可以用Spark的GraphFrames库实现这个逻辑：

步骤1：准备环境与导入依赖 #

确保已安装GraphFrames：

pip install graphframes

导入所需库：

from pyspark.sql import SparkSession
from pyspark.sql.functions import explode, collect_set, sort_array
from graphframes import GraphFrame

步骤2：处理原始数据，构建图结构 #

# 初始化SparkSession（若未初始化）
spark = SparkSession.builder.appName("MergeConnectedArrays").getOrCreate()

# 给每个原始行添加唯一ID，用于建立元素间的连接
df_with_id = df.withColumn("row_id", spark.sparkContext.range(df.count()).cast("int"))

# 提取所有元素作为图的顶点
vertices = df_with_id.select(explode("values").alias("id")).distinct()

# 构建边：同一行内的所有元素互相连接（实现数组内元素的连通）
edges = df_with_id.select("row_id", explode("values").alias("src")) \
                  .join(df_with_id.select("row_id", explode("values").alias("dst")), on="row_id") \
                  .select("src", "dst")

步骤3：计算连通分量并合并数组 #

# 创建GraphFrame图实例
g = GraphFrame(vertices, edges)

# 计算每个元素所属的连通分量
connected_components = g.connectedComponents()

# 按连通分量分组，收集所有元素、去重并排序
result = connected_components.groupBy("component") \
                            .agg(collect_set("id").alias("values")) \
                            .select(sort_array("values").alias("values"))

# 查看结果
result.show(truncate=False)

输出结果 #

+-------------------+
|values             |
+-------------------+
|[a, b, c, d, e, f]|
|[w, x, y, z]       |
+-------------------+

原理说明 #

• 同一数组内的所有元素会被互相连接，形成连通子图；若两个数组有共同元素，它们的子图会被合并为一个更大的连通分量。
• 通过连通分量计算，我们可以将所有相关元素归为同一组，最后收集并排序这些元素，得到合并后的数组。

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