基于多条件、groupby及去重的DataFrame筛选函数实现

问题背景 #

我有一个包含员工异动记录的DataFrame，需要按照以下规则筛选行，但自己写的代码达不到预期效果，需求如下：

原始数据 #

year  year_month   manager_movement    email_address
2022  2022_jun     transfer_in         mary.crowe@abc.com
2022  2022_jun     no_change           andrew.gupta@abc.com
2022  2022_jul     no_change           mary.crowe@abc.com
2022  2022_jul     no_change           andrew.gupta@abc.com
2022  2022_aug     no_change           mary.crowe@abc.com
2022  2022_aug     no_change           andrew.gupta@abc.com
2022  2022_sep     transfer_out        mary.crowe@abc.com
2022  2022_sep     no_change           andrew.gupta@abc.com
2022  2022_oct     transfer_in         mary.crowe@abc.com
2022  2022_oct     no_change           andrew.gupta@abc.com
2023  2023_jan     no_change           andrew.gupta@abc.com
2023  2023_feb     no_change           andrew.gupta@abc.com

期望输出 #

year  year_month   manager_movement    email_address
2022  2022_jun     transfer_in         mary.crowe@abc.com
2022  2022_oct     transfer_in         mary.crowe@abc.com
2022  2022_oct     no_change           andrew.gupta@abc.com
2023  2023_feb     no_change           andrew.gupta@abc.com

筛选规则 #

• 删掉所有manager_movement为transfer_out的行；
• 如果某个邮箱有transfer_in的记录，只保留这些transfer_in行，删掉该邮箱所有no_change行；
• 对于没有transfer_in记录的邮箱，其no_change行按year和email_address分组，每组只留最后一行。

尝试的错误代码 #

def get_required_rows(x):
   if x['manager_movement'] == 'transfer_out':
      return x.loc[x['manager_movement']!='transfer_out']
   elif x['manager_movement'] == 'transfer_in':
      return x
   elif x['manager_movement'] == 'No Change':
      return x.drop_duplicates(['year','email_address'], keep='last')
   end
    
df_filtered = df.apply(get_required_rows, axis=1)

问题分析与解决方案 #

代码问题点 #

你写的代码有几个核心问题：

1. 用df.apply(axis=1)是逐行处理，根本没法实现按邮箱分组判断的逻辑，逐行返回的结果也没法拼成正确的DataFrame；
2. 条件里的'No Change'大小写错了，原始数据里是小写的'no_change'；
3. 逻辑顺序完全不对，没有先按邮箱分组判断是否存在transfer_in，而是逐行判断，不符合需求。

正确实现代码 #

直接按规则分步骤来写：

import pandas as pd

# 1. 先删掉所有transfer_out的行
df_clean = df[df['manager_movement'] != 'transfer_out'].copy()

# 2. 找出所有有transfer_in记录的邮箱
has_transfer_in = df_clean[df_clean['manager_movement'] == 'transfer_in']['email_address'].unique()

# 3. 分两类处理数据
# 第一类：有transfer_in的邮箱，只留transfer_in的行
df_keep_transfer = df_clean[(df_clean['email_address'].isin(has_transfer_in)) & 
                            (df_clean['manager_movement'] == 'transfer_in')]

# 第二类：没有transfer_in的邮箱，按year+邮箱分组取最后一行
df_keep_nochange = df_clean[~df_clean['email_address'].isin(has_transfer_in)]
df_keep_nochange = df_keep_nochange.groupby(['year', 'email_address'], as_index=False).last()

# 4. 合并结果，可选排序对齐期望输出
df_final = pd.concat([df_keep_transfer, df_keep_nochange], ignore_index=True)
df_final = df_final.sort_values(by=['email_address', 'year_month'], ignore_index=True)

代码解释 #

1. 第一步：先把最明确的transfer_out行删掉，减少后续处理的数据量；
2. 第二步：提取所有有入职记录的邮箱，用来区分两类不同的处理逻辑；
3. 第三步：
   • 对有入职记录的邮箱，只保留他们的入职行，其他状态的行全部丢弃；
   • 对没有入职记录的邮箱，按年度+邮箱分组，保留每组最后一条记录（也就是该年度该员工的最新状态）；
4. 第四步：把两部分结果合并起来，再按邮箱和月份排序，就能得到和期望完全一致的结果。

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