Pandas如何实现类似PySpark collect_set的两列同时聚合功能

两处写法的错误点 #

• 第一种写法问题：agg在未指定对应列时，传入每个聚合函数的参数x是单个列的分组Series，不存在x['col2']这类跨列索引的写法，因此执行报错。
• 第二种写法问题：自定义函数里你把lambda函数本身赋值给了字典值，没有实际执行计算，返回结果存储的是函数对象而非去重集合；同时用apply做这类聚合的效率远低于原生agg，属于不必要的实现方式。

正确实现方式 #

写法1：命名聚合（推荐，Pandas 0.25及以上版本支持，最接近PySpark的写法） #

df_collect = df.groupby('col1', as_index=False).agg(
    Col2Arr=('col2', set),
    Col3Arr=('col3', set)
)

这里的命名聚合语法直接指定了新列名、要聚合的原列、对应的聚合函数（set本身就可以直接当聚合函数用，不需要额外包lambda），和你原PySpark的逻辑完全对齐。

写法2：兼容旧版本Pandas的写法 #

如果你的Pandas版本低于0.25，用字典指定列与聚合函数的映射，再重命名列即可：

df_collect = df.groupby('col1').agg({
    'col2': set,
    'col3': set
}).rename(columns={
    'col2': 'Col2Arr',
    'col3': 'Col3Arr'
}).reset_index()

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