Spark GraphFrame aggregateMessages树形节点聚合及多属性扩展问题

一、需求与预期结果 #

树形结构：

(A) --> (B) --> (D)
     \
      \--> (C)

节点总价值计算规则：T(i) = V(i) + 所有子节点的T(i)，其中V(i)为节点自身值。

给定节点自身值：

V(A) = 1
V(B) = 1
V(C) = 1
V(D) = 1

预期计算结果：

T(D) = 1
T(B) = T(D) + V(B) = 2
T(C) = 1
T(A) = V(A) + T(B) + T(C) = 4

二、原始代码与错误结果 #

错误输出 #

+---+----+
| id|cost|
+---+----+
|1-A|   2|
|1-B|   1|
+---+----+

原始代码 #

df = spark.createDataFrame(
        [
            (1,"A", None, 1, 1),
            (1,"B", "A", 2, 1),
            (1,"C", "A", 2, 1),
            (1,"D", "B", 3, 1),
            (2,"A", None, 1, 1),
            (2,"B", "A", 2, 1),
        ], 
        ["frame_id", "node_id", "parent_node_id", "depth", "cost"]
    )
node_df = df.selectExpr("concat_ws('-', frame_id, node_id) as id", 
                            "cost as cost", 
                            "depth as depth", 
                            "frame_id as frame_id", 
                            "node_id as node_id", 
                            "parent_node_id as parent_node_id",
                            )
edge_df = df.selectExpr("concat_ws('-', frame_id, node_id) as src", "concat_ws('-', frame_id, parent_node_id) as dst")
    
g = GraphFrame(node_df, edge_df)
g.aggregateMessages(
        sum(AM.msg).alias("cost"), 
        sendToDst=AM.src["cost"],
    ).orderBy("id").show()

三、错误原因分析 #

1. 消息传递逻辑错误：sendToDst=AM.src["cost"]仅发送了子节点的自身价值V(i)，没有包含子节点的总价值T(i)。根据规则，父节点的T值需要累加子节点的完整T值，而非仅子节点的自身值。
2. 单轮聚合的局限性：aggregateMessages默认是单轮消息传递，只能收集直接子节点的自身值，无法实现树形结构所需的递归累加计算。你后续调整的代码能得到正确结果只是巧合，仅适用于当前深度较浅的树，对于更深层级的树会失效。

四、多属性聚合的实现方案 #

如果节点存在多个需要按相同规则聚合的属性（如value2、value3），推荐使用Pregel API实现递归式计算，它更适合树形结构的迭代累加；若仅需单轮直接子节点聚合，也可通过封装结构体在aggregateMessages中处理。

方案1：Pregel API实现递归多属性计算（推荐） #

Pregel支持迭代式消息传递，能完美实现树形结构的递归累加：

from pyspark.sql.functions import lit, struct, when, col

# 给节点添加额外属性value2、value3
node_df_with_extra = node_df.withColumn("value2", lit(2)).withColumn("value3", lit(3))
g = GraphFrame(node_df_with_extra, edge_df)

# 使用Pregel从叶子节点向上迭代计算总价值
result = g.pregel(
    # 初始消息：叶子节点无子女，初始累加值为0
    initialMsg=struct(lit(0).alias("cost"), lit(0).alias("value2"), lit(0).alias("value3")),
    # 最大迭代次数：设置为树的最大深度即可
    maxIter=3,
    # 子节点向父节点发送自己的总价值（自身值+已累加的子节点值）
    sendMsgToDst=lambda e: when(e.src["depth"] > e.dst["depth"], 
                                struct(
                                    (e.src["cost"] + e.src["agg_cost"]).alias("cost"),
                                    (e.src["value2"] + e.src["agg_value2"]).alias("value2"),
                                    (e.src["value3"] + e.src["agg_value3"]).alias("value3")
                                )),
    # 父节点合并所有子节点的消息，累加总价值
    mergeMsg=lambda m1, m2: struct(
        (m1["cost"] + m2["cost"]).alias("cost"),
        (m1["value2"] + m2["value2"]).alias("value2"),
        (m1["value3"] + m2["value3"]).alias("value3")
    ),
    # 更新节点自身的累加值：自身值 + 子节点累加的总价值
    agg=lambda v, msg: struct(
        (v["cost"] + msg["cost"]).alias("agg_cost"),
        (v["value2"] + msg["value2"]).alias("agg_value2"),
        (v["value3"] + msg["value3"]).alias("agg_value3")
    )
)

# 展示最终结果：agg_cost为T(cost)，agg_value2为T(value2)，以此类推
result.select("id", "cost", "agg_cost", "value2", "agg_value2", "value3", "agg_value3").show()

方案2：aggregateMessages实现单轮多属性聚合（仅直接子节点） #

若仅需计算直接子节点的聚合值（不递归），可将多属性封装为结构体传递：

from pyspark.sql.functions import struct, sum, col

# 先给节点添加额外属性
node_df_with_extra = node_df.withColumn("value2", lit(2)).withColumn("value3", lit(3))
g = GraphFrame(node_df_with_extra, edge_df)

# 聚合直接子节点的所有属性
agg_result = g.aggregateMessages(
    sum(AM.msg["cost"]).alias("total_cost"),
    sum(AM.msg["value2"]).alias("total_value2"),
    sum(AM.msg["value3"]).alias("total_value3"),
    # 发送包含所有属性的结构体消息
    sendToDst=struct(AM.src["cost"], AM.src["value2"], AM.src["value3"]).alias("msg")
).orderBy("id")

# 合并节点自身值与聚合结果，得到最终T值
final_result = agg_result.join(g.nodes, on="id") \
    .withColumn("T_cost", col("cost") + col("total_cost")) \
    .withColumn("T_value2", col("value2") + col("total_value2")) \
    .withColumn("T_value3", col("value3") + col("total_value3")) \
    .select("id", "T_cost", "T_value2", "T_value3")

final_result.show()

内容的提问来源于stack exchange，提问作者月牙天冲