如何使用自定义列条件函数拆分pandas DataFrame

你的自定义函数earlier_or_later仅支持单个时间戳字符串作为输入，但直接传入整列Series对象时，函数无法直接处理批量输入，因此无法返回可用于筛选的布尔序列。

方法1：使用Series.apply()逐元素调用自定义函数 #

这是最直观的修改方式，apply会将自定义函数作用于Series的每一个元素，最终返回和原Series长度一致的布尔序列，可直接用于筛选：

# 直接调用apply逐行处理timestamp列
mask = df['timestamp'].apply(earlier_or_later)
# 拆分两个子集
df1 = df[mask] # 满足条件的子集
df2 = df[~mask] # 不满足条件的子集

如果你的自定义函数需要传入额外参数（比如动态指定阈值），可以用lambda包装：

threshold = "2024-01-01 00:00:00"
mask = df['timestamp'].apply(lambda x: earlier_or_later(x, threshold))

该方法适合数据量较小、函数逻辑复杂无法向量化的场景

方法2：改写为向量化操作（优先推荐） #

pandas的内置时间操作支持向量化运算，执行效率远高于逐行遍历，适合数据量较大的场景。你可以直接将timestamp列转为datetime类型后和阈值比较，不需要自定义函数：

# 先将时间戳列转为pandas datetime类型
df['timestamp'] = pd.to_datetime(df['timestamp'])
# 定义阈值，格式和列数据对齐即可
threshold = pd.to_datetime("2024-01-01 00:00:00")
# 直接向量化比较得到布尔序列
mask = df['timestamp'] < threshold
# 拆分
df1 = df[mask]
df2 = df[~mask]

该方法效率比逐行调用高1~2个数量级，是时间筛选场景的最优选择

方法3：使用numpy.vectorize包装函数 #

如果你的自定义函数逻辑非常复杂、无法改写为向量化操作，可以用numpy的vectorize包装后直接传入Series：

import numpy as np
# 包装自定义函数
vec_func = np.vectorize(earlier_or_later)
# 直接传入整列得到布尔序列
mask = vec_func(df['timestamp'])
# 拆分
df1 = df[mask]
df2 = df[~mask]

注意：该方法本质还是逐元素遍历，效率和apply接近，仅作为写法简化的替代方案

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