基于Pandas根据AM_PM字段匹配填充bag_weight列的实现问题

运行环境 #

Python 3.10.13 (main, Aug 25 2023, 13:20:03) [GCC 9.4.0]
pandas 版本 2.1.1
numpy 版本 1.26.2
pyarrow 版本 14.0.1

需求说明 #

为DataFrame df_entry_log 创建bag_weight列，填充规则如下：

• 当df_entry_log['AM_PM']为"AM"时，用df_entry_log['prev_wrk_day']匹配df_details的gate_date字段，取对应bag_weight值填充
• 当df_entry_log['AM_PM']为"PM"时，用df_entry_log['entry_date']匹配df_details的gate_date字段，取对应bag_weight值填充
  要求采用向量化方法实现，避免循环操作。

数据示例 #

df_entry_log 示例数据 #

guest     entry_date     AM_PM  prev_wrk_day   next_wrk_day
1     janet     2007-01-17     PM     2007-01-16     2007-01-18
2     janet     2007-04-25     AM     2007-04-24     2007-04-26
3     janet     2007-07-25     AM     2007-07-24     2007-07-26

df_details 示例数据 #

guest    gate_date  bag_weight
8     janet     2007-01-16     97.10    
9     janet     2007-01-17     94.95   
10    janet     2007-01-18     89.07  

向量化实现方案 #

以下两种方法均为向量化操作，效率远高于循环，适合大数据量场景：

方法1：字典映射 + np.where #

先将df_details转换为键值对映射，再通过np.where批量选择匹配键，最后映射取值：

import numpy as np
import pandas as pd

# 将df_details转为(guest, gate_date)为键、bag_weight为值的字典
weight_map = df_details.set_index(['guest', 'gate_date'])['bag_weight'].to_dict()

# 批量生成匹配键：根据AM_PM选择对应日期列
match_keys = np.where(
    df_entry_log['AM_PM'] == 'AM',
    list(zip(df_entry_log['guest'], df_entry_log['prev_wrk_day'])),
    list(zip(df_entry_log['guest'], df_entry_log['entry_date']))
)

# 映射填充bag_weight列，无匹配值则填充NaN
df_entry_log['bag_weight'] = np.array([weight_map.get(key, np.nan) for key in match_keys])

方法2：分条件合并填充 #

通过布尔掩码筛选AM/PM行，分别与df_details合并后填充对应值：

import pandas as pd

# 初始化bag_weight列为NaN
df_entry_log['bag_weight'] = pd.NA

# 处理AM场景：匹配guest+prev_wrk_day与guest+gate_date
am_mask = df_entry_log['AM_PM'] == 'AM'
df_entry_log.loc[am_mask, 'bag_weight'] = df_entry_log.loc[am_mask].merge(
    df_details,
    left_on=['guest', 'prev_wrk_day'],
    right_on=['guest', 'gate_date'],
    how='left'
)['bag_weight'].values

# 处理PM场景：匹配guest+entry_date与guest+gate_date
pm_mask = df_entry_log['AM_PM'] == 'PM'
df_entry_log.loc[pm_mask, 'bag_weight'] = df_entry_log.loc[pm_mask].merge(
    df_details,
    left_on=['guest', 'entry_date'],
    right_on=['guest', 'gate_date'],
    how='left'
)['bag_weight'].values

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