Pandas中为groupby函数添加可选变量的实现方法

你的问题出在两个核心点：**kwargs的调用方式错误，以及未处理var_a为None时的分组列表构建。以下是两种正确实现方式：

方法一：使用默认参数（更简洁） #

这种方式比**kwargs更适配单个可选参数的场景：

import pandas as pd

d = {'col1': ['1', '2', '2', '4', '5'], 'vehicle': ['car', 'car', 'truck', 'truck', 'bike'], 'eng': [2, 2, 6, 6, 0], 'id' : ['1','2','3','4','5']}
df_attempt = pd.DataFrame(data=d)

def new_fun(df, var_a=None):
    # 动态生成分组字段列表
    group_cols = ['vehicle']
    if var_a is not None:
        group_cols.append(var_a)
    return df.groupby(group_cols, as_index=False)['id'].count()

# 测试：不传入var_a，仅按vehicle分组
print(new_fun(df_attempt))
# 测试：传入var_a，按vehicle和var_a共同分组
print(new_fun(df_attempt, 'col1'))

方法二：正确使用**kwargs #

如果一定要用**kwargs，需注意调用时必须传关键字参数，同时动态构建分组列表：

import pandas as pd

d = {'col1': ['1', '2', '2', '4', '5'], 'vehicle': ['car', 'car', 'truck', 'truck', 'bike'], 'eng': [2, 2, 6, 6, 0], 'id' : ['1','2','3','4','5']}
df_attempt = pd.DataFrame(data=d)

def new_fun(df, **kwargs):
    var_a = kwargs.get('var_a')
    group_cols = ['vehicle']
    if var_a is not None:
        group_cols.append(var_a)
    return df.groupby(group_cols, as_index=False)['id'].count()

# 测试：不传入var_a
print(new_fun(df_attempt))
# 测试：传入var_a（必须用关键字参数）
print(new_fun(df_attempt, var_a='col1'))

错误原因说明 #

1. **kwargs仅接收关键字参数，你之前调用new_fun(df_attempt, 'col1')用了位置参数，无法被kwargs捕获，必须写成var_a='col1'。
2. 当var_a为None时，直接使用['vehicle', var_a]会生成包含None的列表，pandas无法将None识别为列名，因此需要动态构建分组字段列表，仅在var_a有效时添加。

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