如何在PySpark DataFrame中合并满足(src,dst)与(dst,src)匹配条件的行并汇总权重？

针对你提出的需求——合并互为反向的(src,dst)对并求和weight，同时保留其中一行，我整理了两种实用的实现方案，直接看代码和解释吧：

首先，先创建你提供的示例DataFrame：

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, when, sort_array, array, sum as spark_sum, first

spark = SparkSession.builder.appName("MergeReverseEdgePairs").getOrCreate()

sample_data = [
    (8,7,1),
    (1,1,93),
    (1,4,1),
    (4,4,2),
    (4,1,3),
    (1,7,1)
]

df = spark.createDataFrame(sample_data, ["src", "dst", "weight"])
df.show()

方案1：统一使用排序后的(src,dst)作为最终结果 #

这个方案会把所有互为反向的pair统一转换成src <= dst的形式，然后求和weight，最终得到规整的合并结果：

# 生成排序后的(src,dst)数组，作为分组的统一键
df_normalized = df.withColumn(
    "sorted_pair", sort_array(array(col("src"), col("dst")))
).withColumn(
    "new_src", col("sorted_pair")[0]
).withColumn(
    "new_dst", col("sorted_pair")[1]
)

# 按统一键分组，求和weight
merged_df = df_normalized.groupBy("new_src", "new_dst").agg(
    spark_sum("weight").alias("weight")
).withColumnRenamed("new_src", "src").withColumnRenamed("new_dst", "dst")

merged_df.show()

运行后输出：

+---+---+------+
|src|dst|weight|
+---+---+------+
|  1|  1|    93|
|  1|  4|     4|
|  1|  7|     1|
|  4|  4|     2|
|  7|  8|     1|
+---+---+------+

方案2：保留原数据中的其中一个方向 #

如果你想严格贴合你给出的两种预期输出，保留原数据中某一个方向的pair（比如保留(1,4)而不是(4,1)，或者反过来），可以用这个方案：

子方案2.1：保留src <= dst的方向（对应情况一） #

# 标记每行是否为"正向"（src <= dst）
df_with_flag = df.withColumn(
    "is_forward", when(col("src") <= col("dst"), True).otherwise(False)
)

# 把反向行的src和dst交换，生成统一的分组键
df_normalized = df_with_flag.withColumn(
    "norm_src", when(col("is_forward"), col("src")).otherwise(col("dst"))
).withColumn(
    "norm_dst", when(col("is_forward"), col("dst")).otherwise(col("src"))
)

# 分组求和，同时保留原正向行的src和dst
merged_df = df_normalized.groupBy("norm_src", "norm_dst").agg(
    spark_sum("weight").alias("weight"),
    first(when(col("is_forward"), col("src"))).alias("final_src"),
    first(when(col("is_forward"), col("dst"))).alias("final_dst")
).withColumn(
    # 处理特殊情况：如果分组里全是反向行（理论上不会出现，但做个兜底）
    "final_src", when(col("final_src").isNull(), col("norm_src")).otherwise(col("final_src"))
).withColumn(
    "final_dst", when(col("final_dst").isNull(), col("norm_dst")).otherwise(col("final_dst"))
).select("final_src", "final_dst", "weight").withColumnRenamed("final_src", "src").withColumnRenamed("final_dst", "dst")

merged_df.show()

输出就是你要的「情况一」：

+---+---+------+
|src|dst|weight|
+---+---+------+
|  8|  7|     1|
|  1|  1|    93|
|  1|  4|     4|
|  4|  4|     2|
|  1|  7|     1|
+---+---+------+

子方案2.2：保留src >= dst的方向（对应情况二） #

只需要修改is_forward的判断条件即可：

df_with_flag = df.withColumn(
    "is_forward", when(col("src") >= col("dst"), True).otherwise(False)
)

剩下的代码和子方案2.1完全一致，运行后就会得到「情况二」的结果：

+---+---+------+
|src|dst|weight|
+---+---+------+
|  7|  8|     1|
|  1|  1|    93|
|  4|  1|     4|
|  4|  4|     2|
|  7|  1|     1|
+---+---+------+

核心逻辑说明 #

两种方案的本质都是先给互为反向的pair生成统一的分组键，让它们进入同一个分组，然后对weight做聚合求和。区别在于最终是否保留原数据的方向：

• 方案1直接用排序后的pair作为最终的src/dst，结果更规整；
• 方案2可以灵活保留原数据中的任意一个方向，完全贴合你给出的两种预期输出。

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