如何在Python(pandas)中实现类SQL子查询的条件关联？

数据示例 #

先定义示例DataFrame：

import pandas as pd

# df1示例数据
df1 = pd.DataFrame({
    'id': [1, 2, 3, 4],
    'code': [150, 250, 350, 450]
})

# df2示例数据，s_date仅包含20210101和20220101
df2 = pd.DataFrame({
    's_date': [20210101, 20210101, 20220101, 20220101],
    'range1': [100, 300, 100, 300],
    'range2': [200, 400, 200, 400],
    'category': ['A', 'B', 'C', 'D']
})

参考SQL代码（用户原实现） #

假设你使用的SQL逻辑如下：

SELECT 
    df1.*,
    df2.category,
    df2.s_date
FROM df1
LEFT JOIN df2 
ON df1.code BETWEEN df2.range1 AND df2.range2
AND df2.s_date IN (20210101, 20220101);

Pandas实现方法 #

方法1：笛卡尔积合并后过滤（适合小数据量） #

通过临时键实现交叉合并，再筛选符合范围和日期条件的行，最后关联回原df1保留所有行：

# 添加临时键实现笛卡尔积
cross_merged = pd.merge(df1.assign(temp_key=1), df2.assign(temp_key=1), on='temp_key').drop('temp_key', axis=1)

# 过滤符合条件的记录
filtered = cross_merged[
    cross_merged['code'].between(cross_merged['range1'], cross_merged['range2']) &
    cross_merged['s_date'].isin([20210101, 20220101])
]

# 关联回df1，保留原表所有行
final_result = df1.merge(filtered[['id', 's_date', 'category']], on='id', how='left')

方法2：逐行匹配（灵活但大数据量效率低） #

使用apply逐行对df1的code匹配df2中的符合条件记录：

def match_category(row):
    # 筛选df2中满足范围和日期条件的行
    matched_rows = df2[
        (df2['range1'] <= row['code']) &
        (row['code'] <= df2['range2']) &
        (df2['s_date'].isin([20210101, 20220101]))
    ]
    # 返回所有匹配的category（用逗号分隔，可根据需求调整为列表或取第一个）
    return ','.join(matched_rows['category']) if not matched_rows.empty else None

# 新增category列
df1['category'] = df1.apply(match_category, axis=1)
# 如果需要保留s_date，可修改函数返回元组后拆分列

方法3：merge_asof高效匹配（适合有序范围） #

如果df2的range1是有序的，使用merge_asof可以大幅提升效率：

# 对df1按code排序，df2按range1排序并过滤日期
df1_sorted = df1.sort_values('code').reset_index(drop=True)
df2_filtered_sorted = df2[
    df2['s_date'].isin([20210101, 20220101])
].sort_values('range1').reset_index(drop=True)

# 用merge_asof匹配code >= range1的最近记录
merged = pd.merge_asof(
    df1_sorted,
    df2_filtered_sorted,
    left_on='code',
    right_on='range1',
    direction='backward'
)

# 再过滤code <= range2的记录
result = merged[merged['code'] <= merged['range2']]

# 关联回原df1保留所有行
final_result = df1.merge(result[['id', 's_date', 'category']], on='id', how='left')

方法选择建议 #

• 小数据集：方法1或方法2，实现简单直观
• 大数据集：优先方法3，merge_asof是基于排序的高效匹配，性能远优于前两种

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