Pandas：按条件为timestamp列的NaT值指定datetime

需求场景 #

需要实现：当A列的值满足指定阈值条件时，为timestamp列中值为NaT的行赋值指定的datetime值，同时必须保持timestamp列的datetime64[ns]类型（原列仅包含该类型值或NaT）。

示例数据 #

import pandas as pd

dates = [pd.to_datetime('2022-10-14 10:13:52', format="%Y-%m-%d %H:%M:%S"), 
         pd.to_datetime('2022-10-14 17:43:52', format="%Y-%m-%d %H:%M:%S"),
         pd.to_datetime('2022-10-14 09:00:10', format="%Y-%m-%d %H:%M:%S")]
data = {'A': [-0.5, -0.5, 0.7, 1, 0.65, 0.5], 'timestamp': pd.Series(dates, index=[1, 3, 5])}
df = pd.DataFrame(data=data, index=[0, 1, 2, 3, 4, 5])

初始方法的问题 #

使用numpy.where的尝试会破坏列类型，导致非预期结果：

import numpy as np

threshold = 0.65
null_date = pd.to_datetime('2022-09-01 09:00:00', format="%Y-%m-%d %H:%M:%S")
df.timestamp = np.where(df.A >= threshold, null_date, df.timestamp)

出现的问题：

• timestamp列类型从datetime64[ns]变为object
• 不满足条件的行中，原NaT被替换为None
• 正常的datetime值被转换为时间戳数值

可行解决方案 #

使用DataFrame.apply结合lambda函数，既实现条件赋值，又能严格保持列的datetime类型：

threshold = 0.65
null_date = pd.to_datetime('2022-09-01 09:00:00', format="%Y-%m-%d %H:%M:%S")
df.timestamp = df[['A', 'timestamp']].apply(lambda x: null_date if x['A'] >= threshold else x['timestamp'], axis=1)

执行后，timestamp列将维持datetime64[ns]类型：

• 满足A >= threshold的行，timestamp被赋值为指定的null_date
• 不满足条件的行，保留原有的datetime值或NaT

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