如何在PySpark中实现双层级的集合聚合操作？

我明白你的需求啦——就是要做两层嵌套分组的聚合：先按第一列做“大分组”，再在每个大组里按第二列拆成“子分组”，最后把子分组里指定列的值用collect_set收集成集合对吧？我给你一步步演示怎么实现，先补全你的示例DataFrame，再上代码~

1. 先构造完整的示例数据 #

假设你的原始DataFrame是这样的（补全了你没写完的部分）：

from pyspark.sql import SparkSession

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

df = spark.createDataFrame([
    ('a', 'x', 11, 22),
    ('a', 'x', 11, 33),
    ('a', 'y', 44, 55),
    ('b', 'x', 66, 77),
    ('b', 'y', 88, 99),
    ('b', 'y', 88, 100)
], ['col1', 'col2', 'col3', 'col4'])

2. 分步实现嵌套聚合 #

我们分两步走，先处理内层子分组的聚合，再整理外层大分组的结果：

第一步：内层子分组聚合 #

先按col1+col2联合分组，对col3和col4分别调用collect_set，得到每个子分组的集合结果：

from pyspark.sql import functions as F

# 内层聚合：得到每个(col1, col2)对应的col3、col4集合
inner_agg_df = df.groupBy("col1", "col2") \
    .agg(
        F.collect_set("col3").alias("col3_set"),
        F.collect_set("col4").alias("col4_set")
    )

inner_agg_df.show(truncate=False)

这一步的输出是：

+----+----+---------+-----------+
|col1|col2|col3_set |col4_set   |
+----+----+---------+-----------+
|a   |x   |[11]     |[22, 33]   |
|a   |y   |[44]     |[55]       |
|b   |x   |[66]     |[77]       |
|b   |y   |[88]     |[99, 100]  |
+----+----+---------+-----------+

第二步：外层大分组整理 #

接下来按col1分组，把每个col2对应的聚合结果打包成键值对，最终生成嵌套结构（两种常见格式可选）：

格式1：单个Map存储所有子分组结果

把每个col2作为key，对应的col3_set和col4_set打包成结构体作为value：

final_df = inner_agg_df.groupBy("col1") \
    .agg(
        F.map_from_entries(
            F.collect_list(
                F.struct("col2", F.struct("col3_set", "col4_set"))
            )
        ).alias("subgroup_aggregations")
    )

final_df.show(truncate=False)

输出结果：

+----+-----------------------------------------------------------------------------+
|col1|subgroup_aggregations                                                         |
+----+-----------------------------------------------------------------------------+
|a   |{x -> {col3_set: [11], col4_set: [22, 33]}, y -> {col3_set: [44], col4_set: [55]}}|
|b   |{x -> {col3_set: [66], col4_set: [77]}, y -> {col3_set: [88], col4_set: [99, 100]}}|
+----+-----------------------------------------------------------------------------+

格式2：两个独立Map分别存储列的集合

如果需要把col3和col4的集合分开存储，可以这样写：

final_df = inner_agg_df.groupBy("col1") \
    .agg(
        F.map_from_entries(F.collect_list(F.struct("col2", "col3_set"))).alias("col2_to_col3_set"),
        F.map_from_entries(F.collect_list(F.struct("col2", "col4_set"))).alias("col2_to_col4_set")
    )

final_df.show(truncate=False)

输出结果：

+----+---------------------+-----------------------+
|col1|col2_to_col3_set     |col2_to_col4_set       |
+----+---------------------+-----------------------+
|a   |{x -> [11], y -> [44]}|{x -> [22, 33], y -> [55]}|
|b   |{x -> [66], y -> [88]}|{x -> [77], y -> [99, 100]}|
+----+---------------------+-----------------------+

核心思路 #

因为collect_set是分组级别的聚合函数，无法直接实现嵌套分组，所以我们需要先做「联合分组的内层聚合」，再通过collect_list+map_from_entries把内层结果按外层分组整理成嵌套结构，这样就能完美满足你的需求啦~

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