Spark Python/SQL如何对两列DataFrame关联唯一组合分组生成指定输出

你的需求本质是计算无向图的连通分量，将互相连通的节点归为同一组，再取每组的最小节点作为公共标识拼接为目标格式，以下提供两种实现方式：

方案1：PySpark + GraphFrames（推荐，适合大规模数据） #

该方案使用Spark官方的图计算库GraphFrames实现，分布式优化完善，处理大数据性能更高。

前置依赖 #

提交任务时需携带GraphFrames依赖，本地测试可通过pip安装graphframes库。

实现代码 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import least, greatest, col
from graphframes import GraphFrame

# 初始化SparkSession
spark = SparkSession.builder.appName("group_connected_nodes").getOrCreate()
# 必须设置检查点目录，避免连通分量计算栈溢出
spark.sparkContext.setCheckpointDir("./spark_checkpoint")

# 示例数据，替换为你的实际DataFrame即可
data = [("a","b"), ("b","c"), ("b","d"), ("d","e"), ("f","g")]
df = spark.createDataFrame(data, schema=["col1", "col2"])

# 1. 构造图的顶点（所有去重节点）和边
vertices = df.select(col("col1").alias("id")).union(df.select(col("col2").alias("id"))).distinct()
edges = df.select(least("col1", "col2").alias("src"), greatest("col1", "col2").alias("dst")).distinct()

# 2. 计算连通分量，每个分量会得到唯一的component_id
g = GraphFrame(vertices, edges)
component_result = g.connectedComponents()

# 3. 取每个分量的最小节点作为分组公共标识
component_root = component_result.groupBy("component").agg(col("id").min().alias("root"))
node_root_map = component_result.join(component_root, on="component").select("id", "root")

# 4. 转换为目标输出格式，过滤掉根节点自身映射的行
final_result = node_root_map.filter(col("id") != col("root")) \
    .select(col("root").alias("col1"), col("id").alias("col2")) \
    .orderBy("col1", "col2")

# 输出结果
final_result.show()

输出结果和你期望的格式完全一致：

+----+----+
|col1|col2|
+----+----+
|   a|   b|
|   a|   c|
|   a|   d|
|   a|   e|
|   f|   g|
+----+----+

方案2：纯Spark SQL实现（无需额外依赖） #

使用递归CTE实现连通分量计算，适合中小规模数据，无需引入第三方库。

实现代码 #

首先将你的源DataFrame注册为临时视图：

df.createOrReplaceTempView("node_pairs")

然后执行如下SQL：

-- 调整最大递归深度，根据你的连通链最长长度调整
SET spark.sql.recursiveCTE.maxIterations = 1000;

WITH RECURSIVE 
-- 统一边的方向，避免重复计算
edges AS (
    SELECT LEAST(col1, col2) AS u, GREATEST(col1, col2) AS v FROM node_pairs
),
-- 收集所有去重节点
all_nodes AS (
    SELECT col1 AS node FROM node_pairs UNION SELECT col2 AS node FROM node_pairs
),
-- 递归计算每个节点对应的最小根节点
connected AS (
    -- 初始状态：每个节点的根是自身
    SELECT node AS root, node AS member FROM all_nodes
    UNION ALL
    -- 迭代扩展连通节点，始终取最小节点作为根
    SELECT 
        LEAST(c.root, e.u, e.v) AS root,
        CASE WHEN c.member = e.u THEN e.v ELSE e.u END AS member
    FROM connected c
    JOIN edges e ON c.member = e.u OR c.member = e.v
    WHERE LEAST(c.root, e.u, e.v) < c.root -- 避免死循环，只有根更新才继续迭代
),
-- 去重得到每个节点对应的唯一最小根
node_root AS (
    SELECT member, MIN(root) AS root FROM connected GROUP BY member
)
-- 输出目标格式，过滤掉根节点自身
SELECT root AS col1, member AS col2 
FROM node_root 
WHERE member != root
ORDER BY col1, col2;

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