使用PySpark对等价替代项分组并为每组分配唯一代表项

当前使用的等价项数据集构造逻辑如下：

# Creating the DataFrame
test =sqlContext.createDataFrame([(1,2),(2,1),
                               (1,3),(2,3),
                               (3,2),(3,1),
                               (4,5),(5,4)],
                               ['cod_item','alter_cod'])

已完成的聚合操作是按cod_item分组收集对应替代项，代码如下：

test.createOrReplaceTempView("teste")

teste = spark.sql("""select cod_item,
                  collect_list(alter_cod) as alternative_item 
                  from teste
                  group by cod_item""")

当前输出仅包含每个项对应的直接替代项列表，无法处理等价关系的传递性（比如1和2等价、2和3等价时，1、2、3同属一个等价组），需要给每个等价组分配唯一的统一代表编码，同组所有项都映射到该编码。

这类等价归并本质是无向图的连通分量计算问题：每个项编码是图的节点，两个项存在等价关系就在节点间连一条无向边，同一个连通分量内的所有节点属于同一个等价组，直接取组内最小的编码作为统一代表项即可。

方案1：GraphFrames连通分量算法（适配大规模数据集） #

GraphFrames是Spark生态的图计算库，内置的连通分量算法可以高效完成大规模数据的等价组归并：

from graphframes import GraphFrame
from pyspark.sql import functions as F
from pyspark.sql.window import Window

# 构造图的顶点表：所有出现过的项编码
vertices = test.selectExpr("cod_item as id").distinct()
# 构造图的边表：等价是双向关系，存双向边后去重
edges = test.selectExpr("cod_item as src", "alter_cod as dst").union(
    test.selectExpr("alter_cod as src", "cod_item as dst")
).distinct()

# 初始化图并计算连通分量
g = GraphFrame(vertices, edges)
connect_res = g.connectedComponents()

# 每个连通分量取最小id作为代表编码
final_result = connect_res.withColumn(
    "represent_cod",
    F.min("id").over(Window.partitionBy("component"))
).select(F.col("id").alias("cod_item"), "represent_cod")

执行后得到的结果如下：

方案2：递归CTE实现（无第三方依赖，适合中小规模数据集） #

如果不想引入额外依赖，可以用Spark SQL的递归CTE实现连通分量归并：

test.createOrReplaceTempView("item_pairs")

final_result = spark.sql("""
with recursive bidirect_edges as (
    -- 构造双向等价边
    select cod_item as n1, alter_cod as n2 from item_pairs
    union
    select alter_cod as n1, cod_item as n2 from item_pairs
),
node_groups as (
    -- 初始状态：每个节点自身为初始组代表
    select n1 as node, n1 as min_rep from bidirect_edges group by n1
    union all
    -- 递归遍历相连节点，更新组内最小代表
    select e.n2 as node, least(ng.min_rep, e.n2) as min_rep
    from node_groups ng
    join bidirect_edges e on ng.node = e.n1
    where least(ng.min_rep, e.n2) < ng.min_rep
)
-- 每个节点取最终的最小代表编码
select node as cod_item, min(min_rep) as represent_cod
from node_groups
group by node
""")

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