循环创建事件前后6个月的Pandas DataFrame并按规则命名

场景与问题 #

现有一个包含约220万条记录的DataFrame dftxt，包含date（带时分的datetime类型）和text（字符串类型）两列。需要基于20个事件日期，为每个事件生成事件前6个月和事件后6个月的DataFrame，命名格式要求为preten1event1、postten1event1这类形式。

事件日期定义如下：

import datetime
from dateutil.relativedelta import relativedelta

ten1event1 = datetime.datetime(2011, 2, 15)
ten1event2 = datetime.datetime(2012, 9, 27)
ten1event3 = datetime.datetime(2013, 2, 10)
ten1event4 = datetime.datetime(2020, 3, 11)

ten2event1 = datetime.datetime(2013, 3, 25)
ten2event2 = datetime.datetime(2015, 10, 22)
ten2event3 = datetime.datetime(2018, 3, 7)
ten2event4 = datetime.datetime(2021, 9, 24)

ten3aevent1 = datetime.datetime(2010, 5, 22)
ten3aevent2 = datetime.datetime(2014, 1, 9)
ten3aevent3 = datetime.datetime(2020, 8, 11)
ten3aevent4 = datetime.datetime(2022, 1, 5)

ten3bevent1 = datetime.datetime(2013, 3, 25)
ten3bevent2 = datetime.datetime(2014, 7, 29)
ten3bevent3 = datetime.datetime(2017, 12, 1)
ten3bevent4 = datetime.datetime(2021, 10, 19)

ten4event1 = datetime.datetime(2011, 6, 14)
ten4event2 = datetime.datetime(2013, 7, 23)
ten4event3 = datetime.datetime(2014, 2, 11)
ten4event4 = datetime.datetime(2019, 2, 25)

原尝试代码如下，但运行后无法生成预期的DataFrame，报NameError: name 'preten1event1' is not defined：

events = [ten1event1,ten1event2,ten1event3,ten1event4,ten2event1,ten2event2,ten2event3,ten2event4,ten3aevent1,
          ten3aevent2,ten3aevent3,ten3aevent4,ten3bevent1,ten3bevent2,ten3bevent3,ten3bevent4,ten4event1,
          ten4event2,ten4event3,ten4event4]

for x in events:
    dfpre = dftxt.loc[(dftxt['date'] >= x - relativedelta(months=6)) & (dftxt['date'] <= x)]
    dfpre = dfpre.reset_index()
    dfpre.drop(columns=['index'], inplace=True)
    f"pre{x}" == dfpre
    
    dfpost = dftxt.loc[(dftxt['date'] >= x) & (dftxt['date'] <= x + relativedelta(months=6))]
    dfpost = dfpost.reset_index()
    dfpost.drop(columns=['index'], inplace=True)
    f"post{x}" == dfpost

错误原因 #

原代码中f"pre{x}" == dfpre只是执行字符串与DataFrame的比较操作，并没有将DataFrame赋值给对应名称的变量。同时，直接通过字符串无法直接创建变量，必须借助命名空间或容器来实现。

解决方案 #

方法1：动态创建全局变量（符合命名格式要求） #

利用globals()字典（当前全局命名空间的映射），通过变量名字符串作为键，将DataFrame赋值到全局命名空间中。注意：需要把事件与其名称配对，才能生成对应的变量名。

修改后的代码：

# 重新定义事件列表，每个元素是(事件名称, 事件日期)的元组
events = [
    ('ten1event1', ten1event1),
    ('ten1event2', ten1event2),
    ('ten1event3', ten1event3),
    ('ten1event4', ten1event4),
    ('ten2event1', ten2event1),
    ('ten2event2', ten2event2),
    ('ten2event3', ten2event3),
    ('ten2event4', ten2event4),
    ('ten3aevent1', ten3aevent1),
    ('ten3aevent2', ten3aevent2),
    ('ten3aevent3', ten3aevent3),
    ('ten3aevent4', ten3aevent4),
    ('ten3bevent1', ten3bevent1),
    ('ten3bevent2', ten3bevent2),
    ('ten3bevent3', ten3bevent3),
    ('ten3bevent4', ten3bevent4),
    ('ten4event1', ten4event1),
    ('ten4event2', ten4event2),
    ('ten4event3', ten4event3),
    ('ten4event4', ten4event4),
]

for name, event_date in events:
    # 生成事件前6个月的DataFrame
    pre_df = dftxt.loc[
        (dftxt['date'] >= event_date - relativedelta(months=6)) & 
        (dftxt['date'] <= event_date)
    ].reset_index(drop=True)  # 简化索引重置操作
    globals()[f'pre{name}'] = pre_df
    
    # 生成事件后6个月的DataFrame
    post_df = dftxt.loc[
        (dftxt['date'] >= event_date) & 
        (dftxt['date'] <= event_date + relativedelta(months=6))
    ].reset_index(drop=True)
    globals()[f'post{name}'] = post_df

执行后即可直接使用preten1event1、postten1event1等变量。

方法2：用字典统一存储（更推荐） #

不直接创建大量全局变量，而是用字典作为容器存储所有生成的DataFrame，避免污染全局命名空间，同时方便后续批量处理。

代码示例：

event_dfs = {}  # 存储所有事件前后的DataFrame

for name, event_date in events:
    # 事件前6个月数据
    pre_df = dftxt.loc[
        (dftxt['date'] >= event_date - relativedelta(months=6)) & 
        (dftxt['date'] <= event_date)
    ].reset_index(drop=True)
    event_dfs[f'pre{name}'] = pre_df
    
    # 事件后6个月数据
    post_df = dftxt.loc[
        (dftxt['date'] >= event_date) & 
        (dftxt['date'] <= event_date + relativedelta(months=6))
    ].reset_index(drop=True)
    event_dfs[f'post{name}'] = post_df

# 访问示例：获取ten1event1的前6个月数据
# event_dfs['preten1event1']

这种方式更易维护，后续遍历、修改所有事件数据时更便捷。

内容的提问来源于stack exchange，提问作者Jesse-Burton Nicholson