如何使用pandas.groupby.apply替代agg实现用户分组合并输出单行结果

实现代码 #

import pandas as pd

# 构造测试数据
d = {'User':['A', 'A', 'B'],
     'time':[1,2,3],
     'state':['CA', 'CA', 'OR'],
     'type':['cd', 'dvd', 'cd']}
df = pd.DataFrame(data=d)

# 预获取全局所有type类型，保证最终结果所有type列都存在，缺失值自动补0
all_types = df['type'].unique().tolist()

def user_agg(user_df):
    # 你已实现的取值逻辑
    avg_time = user_df['time'].mean()
    state = user_df['state'].iloc[0]
    user_id = user_df['User'].iloc[0]
    type_counts = user_df['type'].value_counts().to_dict()
    
    # 整合固定字段与type计数字段
    res_dict = {
        'User': user_id,
        'avg_time': avg_time,
        'state': state
    }
    for t in all_types:
        res_dict[t] = type_counts.get(t, 0)
    
    # 返回单行结构DataFrame
    return pd.DataFrame([res_dict])

# groupby应用函数，group_keys=False避免生成多余索引层级
final_df = df.groupby('User', group_keys=False).apply(user_agg).reset_index(drop=True)

核心说明 #

将整合好的全量字段字典放在列表中传入pd.DataFrame构造函数，即可直接生成符合要求的单行返回结果，完全适配groupby.apply的调用逻辑，也满足后续并行化改造的要求。

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