将多重SQL CASE语句转换为Pandas的np.where语法

原SQL逻辑拆解 #

原SQL是两层嵌套的CASE判断，核心逻辑如下：

1. 外层判断：当SCH_END_LOCN_ID != ATS_STA_ID，且PLC_ACTUAL_DEPART_TIME、BEACON_ACTUAL_DEPART_TIME均为空时，进入内层分支；否则取BEACON_ACTUAL_DEPART_TIME、PLC_ACTUAL_DEPART_TIME、ITRAC_ACTUAL_DEPART_TIME中第一个非空值（对应SQL的COALESCE）。
2. 内层分支：按优先级依次判断，返回第一个满足条件的计算值：
   • 若ITRAC_ACTUAL_DEPART_TIME非空，直接返回该值
   • 若PLC_ACTUAL_DEPART_TIME_CLEAR非空，返回该值减去15秒（15/(24*60*60)是将15秒转换为天数单位，适配Pandas datetime运算）
   • 若PLC_ACTUAL_ARRIVE_TIME_DWELL非空，返回该值加上median_dwell
   • 若PLC_ACTUAL_ARRIVE_TIME非空，返回该值加上median_track_occ
   • 若ITRAC_ACTUAL_ARRIVE_TIME非空，返回该值加上30秒
   • 所有条件都不满足时返回NaN

Pandas代码实现（嵌套np.where写法） #

通过多层嵌套np.where严格对应SQL的分支优先级：

import numpy as np
import pandas as pd

# 定义外层判断条件
outer_condition = (df['SCH_END_LOCN_ID'] != df['ATS_STA_ID']) & \
                  (df['PLC_ACTUAL_DEPART_TIME'].isna()) & \
                  (df['BEACON_ACTUAL_DEPART_TIME'].isna())

# 内层分支：按优先级嵌套np.where
inner_result = np.where(df['ITRAC_ACTUAL_DEPART_TIME'].notna(),
                        df['ITRAC_ACTUAL_DEPART_TIME'],
                        np.where(df['PLC_ACTUAL_DEPART_TIME_CLEAR'].notna(),
                                 df['PLC_ACTUAL_DEPART_TIME_CLEAR'] - pd.Timedelta(seconds=15),
                                 np.where(df['PLC_ACTUAL_ARRIVE_TIME_DWELL'].notna(),
                                          df['PLC_ACTUAL_ARRIVE_TIME_DWELL'] + df['median_dwell'],
                                          np.where(df['PLC_ACTUAL_ARRIVE_TIME'].notna(),
                                                   df['PLC_ACTUAL_ARRIVE_TIME'] + df['median_track_occ'],
                                                   np.where(df['ITRAC_ACTUAL_ARRIVE_TIME'].notna(),
                                                            df['ITRAC_ACTUAL_ARRIVE_TIME'] + pd.Timedelta(seconds=30),
                                                            pd.NaT)))))

# 组合外层判断与结果
final_result = np.where(outer_condition,
                        inner_result,
                        df[['BEACON_ACTUAL_DEPART_TIME', 'PLC_ACTUAL_DEPART_TIME', 'ITRAC_ACTUAL_DEPART_TIME']].bfill(axis=1).iloc[:, 0])

# 将结果赋值给DataFrame新列
df['CALCULATED_DEPART_TIME'] = final_result

代码说明 #

• 用isna()/notna()替代直接和np.datetime64('NaT')比较，更符合Pandas datetime空值判断规范
• 用pd.Timedelta(seconds=15)替代15/(24*60*60)，可读性更强，避免浮点运算误差
• 外层else部分用bfill(axis=1).iloc[:,0]实现COALESCE逻辑：按列顺序取第一个非空值

更简洁的np.select实现（推荐） #

多分支场景下np.select可读性更高，逻辑更直观：

# 定义内层条件与对应结果列表
inner_conditions = [
    df['ITRAC_ACTUAL_DEPART_TIME'].notna(),
    df['PLC_ACTUAL_DEPART_TIME_CLEAR'].notna(),
    df['PLC_ACTUAL_ARRIVE_TIME_DWELL'].notna(),
    df['PLC_ACTUAL_ARRIVE_TIME'].notna(),
    df['ITRAC_ACTUAL_ARRIVE_TIME'].notna()
]

inner_values = [
    df['ITRAC_ACTUAL_DEPART_TIME'],
    df['PLC_ACTUAL_DEPART_TIME_CLEAR'] - pd.Timedelta(seconds=15),
    df['PLC_ACTUAL_ARRIVE_TIME_DWELL'] + df['median_dwell'],
    df['PLC_ACTUAL_ARRIVE_TIME'] + df['median_track_occ'],
    df['ITRAC_ACTUAL_ARRIVE_TIME'] + pd.Timedelta(seconds=30)
]

inner_result = np.select(inner_conditions, inner_values, default=pd.NaT)

# 外层判断组合结果
final_result = np.where(outer_condition,
                        inner_result,
                        df[['BEACON_ACTUAL_DEPART_TIME', 'PLC_ACTUAL_DEPART_TIME', 'ITRAC_ACTUAL_DEPART_TIME']].bfill(axis=1).iloc[:, 0])

df['CALCULATED_DEPART_TIME'] = final_result

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