问题描述
我正在将一些 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