如何在Spark中对DataFrame的两列双向组合进行GroupBy分组

要实现将(src, dst)和(dst, src)视为同一分组并合并计数的需求，核心思路是统一反向组合的分组标识——让所有反向组合拥有相同的分组键，之后再进行分组聚合即可。下面是两种简洁可行的实现方案：

方案一：使用least和greatest函数（推荐） #

Spark 2.0+提供了least和greatest函数，可以直接获取两列中的较小值和较大值，用这两个值作为分组键，就能自动把反向组合归为同一组。

步骤1：导入必要的函数 #

import org.apache.spark.sql.functions.{least, greatest, count}

步骤2：处理DataFrame并分组计数 #

val result = dff
  // 创建统一的分组键：较小值作为group_src，较大值作为group_dst
  .withColumn("group_src", least($"src", $"dst"))
  .withColumn("group_dst", greatest($"src", $"dst"))
  // 基于统一分组键聚合计数
  .groupBy("group_src", "group_dst")
  .agg(count("*").alias("count"))
  // 重命名列以匹配你期望的输出格式
  .withColumnRenamed("group_src", "src")
  .withColumnRenamed("group_dst", "dst")

执行结果 #

运行上述代码后，会得到合并计数后的结果：

+---+---+-----+
|src|dst|count|
+---+---+-----+
| B| C| 1|
| C| D| 2|
| A| C| 4|
| A| B| 5|
+---+---+-----+

方案二：使用when条件判断 #

如果你的Spark版本较低（不支持least/greatest），可以用when函数手动判断并生成统一分组键：

代码实现 #

import org.apache.spark.sql.functions.{when, count}

val result = dff
  .withColumn("group_src", when($"src" < $"dst", $"src").otherwise($"dst"))
  .withColumn("group_dst", when($"src" < $"dst", $"dst").otherwise($"src"))
  .groupBy("group_src", "group_dst")
  .agg(count("*").alias("count"))
  .withColumnRenamed("group_src", "src")
  .withColumnRenamed("group_dst", "dst")

这个方案的逻辑和方案一完全一致，只是用条件判断替代了内置函数，最终效果相同。

说明 #

两种方案都会将反向组合统一为字典序升序的分组（比如(B,A)会被转为(A,B)），如果你的业务需要保留原数据中出现次数更多的组合作为显示的分组名，可以在聚合后额外添加一步逻辑：先统计每个原组合的次数，再为每个统一分组匹配出现次数最多的原组合。不过对于大多数场景，统一为字典序的分组已经能满足需求，且实现最简单。

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