Pandas：一次在两列上应用 agg

如何解决Pandas：一次在两列上应用 agg

我正在将一些 pySpark 代码迁移到 Pandas 中，并坚持在两列上实现 collect_set。
pySpark 代码如下所示：

df_collect = df.groupBy('col1').agg(collect_set('col2').alias('Col2Arr'),collect_set('col3').alias('Col3Arr'))

我可以通过在 lambda 上调用 agg 函数轻松实现其中一列，但不能同时在两列上执行此操作：

df_collect = df.groupby('col1')['col2'].agg({'Col2Arr': lambda x: set(x)})

我试过了：

df.groupby('col1').agg(Col2Arr = lambda x: set(x['col2']),Col3Arr = lambda x: set(x['col3']))

和

def count_set(x):
    d = {}
    d['Col2Arr'] = lambda a: set(a['col2'])
    d['Col3Arr'] = lambda a: set(a['col3'])
    return pd.Series(d,index=['Col2Arr','Col3Arr'])

df.groupby('col1').apply(count_set)

似乎没有任何效果。不确定我在这里缺少什么。

解决方法

根据您要查找的内容，正如 @anky 建议的那样，选择所需列的标准 groupby agg 可能会起作用：

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

df_collect：

   col1                         col2                         col3
0     1           {1,2,5,6,7,9}        {1,3,4,9}
1     2  {1,8,9}     {1,9}
2     3              {2,8}        {2,9}
3     4     {1,9}
4     5           {1,7}           {1,9}
5     6              {2,9}
6     7        {1,9}           {1,8}
7     8        {1,8}  {1,9}
8     9           {1,9}              {2,9}

或者，对于更类似于 PySpark 的外观，使用 Named Aggregation 来合并别名、列选择和单独的聚合选项：

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

df_collect：

   col1                      Col2Arr                      Col3Arr
0     1           {1,9}

---

使用的样本数据：

import numpy as np
import pandas as pd

np.random.seed(5)
df = pd.DataFrame(np.random.randint(1,10,(100,3)),columns=[1,3]).add_prefix('col')

df.head(10)：

   col1  col2  col3
0     4     7     7
1     1     9     5
2     8     1     1
3     8     2     6
4     8     1     2
5     5     7     3
6     2     3     8
7     1     6     1
8     1     5     5
9     4     3     5