如何封装pandas中基于groupby apply的通用分组掩码计算函数

原自定义函数的问题点 #

• 列访问语法错误：用变量传递列名时不能使用df.列名的点访问写法，必须用df[变量名]的索引方式，否则会直接查找固定名为col的列触发报错
• 索引错位：原代码计算m1后没有做索引对齐，直接在apply的子分组里cumsum会因为loc筛选丢了非目标行的索引，导致编号和原行匹配错误
• 语法疏漏：函数调用时传入var参数的字符串少了闭合引号
• 分组逻辑冲突：把函数放在groupby.apply里调用时，传入的子DataFrame已经是分组后的切片，原代码里又重复做groupby会导致计算逻辑异常

可复用通用封装代码 #

函数无任何硬编码列名，所有逻辑参数全部可配置，直接传入原始全量DataFrame即可，不需要额外嵌套groupby调用：

import pandas as pd

def gen_target_group(df, group_by, filter_col, filter_val, calc_col, new_col='new_group'):
    # 按分组列取上一行的计算列值，索引和原表完全对齐
    prev_val = df.groupby(group_by)[calc_col].shift()
    # 筛选目标行，判断当前行与上一行同目标类型记录的差值是否大于上一行的基准值
    target_row_mask = df[filter_col] == filter_val
    break_mask = df.loc[target_row_mask, calc_col].groupby(df[group_by]).diff().gt(prev_val.loc[target_row_mask])
    # 非目标行掩码
    non_target_mask = df[filter_col].ne(filter_val)
    # 按分组累加断点生成连续编号，非目标行置为空值
    df[new_col] = break_mask.groupby(df[group_by]).cumsum().add(1).mask(non_target_mask)
    return df

调用方式 #

# 对应原示例的参数传入即可，结果和手动写的原生逻辑完全一致
df = gen_target_group(
    df,
    group_by='person',
    filter_col='product',
    filter_val='fruit',
    calc_col='time_bought'
)

注意事项 #

如果time_bought列是带单位的字符串格式（如示例里的10min/5min），需要先转换成可比较的时间类型再调用函数，否则大小比较会报错，转换代码参考：

# 转成pandas时间差类型
df['time_bought'] = pd.to_timedelta(df['time_bought'])
# 若需要转成数值类型的分钟数，可使用下面的写法
# df['time_bought'] = pd.to_timedelta(df['time_bought']).dt.total_seconds() // 60

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