将含CASE WHEN的SQL查询转换为Python代码的技术求助

输入数据构造 #

先把给定的输入表用Pandas构造出来：

import pandas as pd
import numpy as np

# 构造base_table
base_table = pd.DataFrame({
    'year': [2022, 2023, 2022, 2023],
    'category': ['English', 'English', 'English', 'Science'],
    'country': [None, None, 'US', None],
    'class': [1, 2, 1, 2]
})

# 构造inter_table
inter_table = pd.DataFrame({
    'year': [2022, 2023, 2022, 2023],
    'country': ['US', 'US', 'Europe', 'Europe'],
    'amount': [100, 400, 300, 200]
})

表关联（左连接） #

和SQL的左连接逻辑一致，用pd.merge实现：

# 左连接两张表，关联条件是year和country
merged_df = pd.merge(
    base_table,
    inter_table,
    on=['year', 'country'],
    how='left'
)

实现new_col的条件逻辑 #

先拆解SQL中的CASE WHEN逻辑，转化为Python的条件判断：

1. 当country为NULL且category是'English'时：
   • 如果class=1：取同year下inter_table中country='US'的amount乘以5
   • 如果class=2：由于原SQL中判断的子查询（匹配NULL的country）永远返回NULL，因此该分支结果为NULL
2. 其他情况：直接取关联后的amount列值

为了高效实现，先预先提取inter_table中各year对应US的amount值，避免重复查询：

# 先构建year到US对应amount的映射表
us_amount_map = inter_table[inter_table['country'] == 'US'].set_index('year')['amount'].to_dict()

# 定义条件列表和对应结果列表
conditions = [
    # 条件1：country为空且category是English且class=1
    (merged_df['country'].isna()) & (merged_df['category'] == 'English') & (merged_df['class'] == 1),
    # 条件2：country为空且category是English且class=2
    (merged_df['country'].isna()) & (merged_df['category'] == 'English') & (merged_df['class'] == 2)
]

results = [
    # 对应条件1：取当前year的US amount*5
    merged_df['year'].map(us_amount_map) * 5,
    # 对应条件2：返回NULL
    np.nan
]

# 用np.select实现多条件赋值，默认值为关联后的amount
merged_df['new_col'] = np.select(conditions, results, default=merged_df['amount'])

# 调整列顺序，和输出表一致
output_df = merged_df[['year', 'category', 'country', 'class', 'new_col']]
print(output_df)

运行后得到的结果和给定的输出表完全一致：

year category country  class  new_col
0  2022  English    None      1    500.0
1  2023  English    None      2      NaN
2  2022  English      US      1    100.0
3  2023  Science    None      2      NaN

逻辑说明 #

• 预先构建us_amount_map是为了避免在条件判断中重复筛选inter_table，提升效率
• 使用np.select比嵌套的apply+if-else更高效，尤其是处理大数据量时
• 对于class=2的分支，原SQL中因为a.country是NULL，子查询select amount from inter_table b where a.country=b.country and b.year=a.year永远返回NULL，因此直接赋值np.nan

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