如何基于任意日期动态拆分Python DataFrame为两个子数据框

我有一个包含date列及其他列的Python DataFrame，需要基于运行时确定的任意日期arbitrary_date将其拆分为两个子DataFrame：

1. 所有日期早于arbitrary_date的行
2. 序列中紧随arbitrary_date的下一个日期对应的所有行（该日期可能是arbitrary_date本身，也可能是未来一周内的其他日期）

限制条件 #

• arbitrary_date仅在运行时可知
• 无法通过索引拆分，不能保证特定索引归属
• 未知arbitrary_date与下一个日期的间隔天数
• 未知下一个日期对应的行数

示例 #

初始DataFrame：

date        data
20230101    e
20230108    E
20230115    S
20230117    L
20230117    f
20230118    u
20230118    B
20230121    r
20230127    n
20230128    o

示例1：arbitrary_date = 2023-01-17

拆分结果1（早于目标日期的行）：

20230101    e
20230108    E
20230115    S

拆分结果2（紧随的下一个日期的所有行）：

20230117    L
20230117    f

示例2：arbitrary_date = 2023-01-22

拆分结果1（早于目标日期的行）：

20230101    e
20230108    E
20230115    S
20230117    L
20230117    f
20230118    u
20230118    B
20230121    r

拆分结果2（紧随的下一个日期的所有行）：

20230127    n

1. 统一日期格式 #

首先确保date列和arbitrary_date都是datetime类型，避免字符串格式导致的比较错误：

import pandas as pd

# 将DataFrame的date列转为datetime（假设原始是字符串格式）
df['date'] = pd.to_datetime(df['date'], format='%Y%m%d')
# 将运行时传入的任意日期也转为datetime
arbitrary_date = pd.to_datetime(arbitrary_date)

2. 获取紧随目标日期的有效日期 #

从DataFrame中筛选出所有大于等于arbitrary_date的日期，去重排序后取第一个，就是需要的"紧随日期"：

# 提取符合条件的日期并去重排序
valid_dates = df[df['date'] >= arbitrary_date]['date'].unique()
valid_dates.sort()
target_date = valid_dates[0]

3. 拆分DataFrame #

根据条件拆分出两个子DataFrame：

# 子DataFrame1：所有早于arbitrary_date的行
df_before = df[df['date'] < arbitrary_date]

# 子DataFrame2：所有等于target_date的行
df_next = df[df['date'] == target_date]

完整代码示例 #

import pandas as pd

# 构造示例DataFrame
data = {
    'date': ['20230101', '20230108', '20230115', '20230117', '20230117',
             '20230118', '20230118', '20230121', '20230127', '20230128'],
    'data': ['e', 'E', 'S', 'L', 'f', 'u', 'B', 'r', 'n', 'o']
}
df = pd.DataFrame(data)

# 统一日期格式
df['date'] = pd.to_datetime(df['date'], format='%Y%m%d')
arbitrary_date = pd.to_datetime('2023-01-22')

# 获取紧随的目标日期
valid_dates = df[df['date'] >= arbitrary_date]['date'].unique()
valid_dates.sort()
target_date = valid_dates[0]

# 拆分操作
df_before = df[df['date'] < arbitrary_date]
df_next = df[df['date'] == target_date]

# 输出结果
print("早于目标日期的子DataFrame：")
print(df_before)
print("\n紧随目标日期的子DataFrame：")
print(df_next)

说明 #

• 代码完全不依赖索引，仅基于date列的值进行拆分，符合限制条件
• 自动适配两种场景：arbitrary_date存在于DataFrame中（直接取该日期）、arbitrary_date不存在（取之后最近的日期）
• 日期格式统一为datetime，确保比较逻辑的准确性

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