优化Pandas DataFrame与分区SQL表关联，筛选指定时间交易记录

需求说明 #

现有：

• Pandas DataFrame df_A，包含ID和Added_Date列
• 大型分区SQL表transactions，包含ID、Transaction_Date、Year、Month、Day列

需要关联两者，筛选出transactions中每个ID对应的Transaction_Date在其Added_Date之后30天内的交易记录，生成新DataFrame，以下是优化实现方案。

示例数据 #

Pandas DataFrame 示例 #

import sqlite3
import pandas as pd 

data = {'ID': [1, 2, 3], 'Added_Date': ['2023-02-01', '2023-04-15', '2023-03-17']}
df_A = pd.DataFrame(data)
# 转换日期类型，避免后续隐式类型转换性能损耗
df_A['Added_Date'] = pd.to_datetime(df_A['Added_Date'])

内存SQL表示例（模拟大型分区表） #

# 创建内存SQLite数据库
conn = sqlite3.connect(':memory:')
c = conn.cursor()

# 创建交易表（模拟分区表结构）
c.execute('''CREATE TABLE transactions
             (ID INTEGER, transaction_date DATE, Year INTEGER, Month INTEGER, Day INTEGER)''')

# 插入示例数据，同时填充分区字段
sample_data = [
    (1, '2023-01-15', 2023, 1, 15), (1, '2023-02-10', 2023, 2, 10), (1, '2023-03-01', 2023, 3, 1),
    (2, '2023-04-01', 2023, 4, 1), (2, '2023-04-20', 2023, 4, 20), (2, '2023-05-05', 2023, 5, 5),
    (3, '2023-03-10', 2023, 3, 10), (3, '2023-03-25', 2023, 3, 25), (3, '2023-04-02', 2023, 4, 2)
]
c.executemany('INSERT INTO transactions VALUES (?, ?, ?, ?, ?)', sample_data)
conn.commit()

优化方案 #

方案1：SQL端完成全量筛选（最优） #

核心思路：将小表df_A传入SQL，在数据库端完成关联和筛选，仅返回符合条件的数据，避免把大型SQL表全量拉取到Python内存中，同时利用分区表的Year/Month键缩小扫描范围，大幅提升效率。

代码实现：

# 将df_A写入SQL临时表
df_A.to_sql('temp_added_dates', conn, index=False, if_exists='replace')

# 编写优化后的SQL查询：先通过分区键过滤，再精确匹配日期区间
query = '''
SELECT t.ID, t.transaction_date
FROM transactions t
JOIN temp_added_dates a ON t.ID = a.ID
WHERE 
    -- 利用分区键减少扫描的分区数量
    (t.Year = strftime('%Y', a.Added_Date) AND t.Month >= strftime('%m', a.Added_Date))
    OR (t.Year = strftime('%Y', date(a.Added_Date, '+30 days')) AND t.Month <= strftime('%m', date(a.Added_Date, '+30 days')))
    -- 精确筛选Transaction_Date在Added_Date后30天内的记录
    AND t.transaction_date BETWEEN a.Added_Date AND date(a.Added_Date, '+30 days')
'''

# 执行查询并读取结果
result_df = pd.read_sql(query, conn)
print(result_df)

方案2：批量分ID查询（适用于无法创建临时表的场景） #

如果无法在SQL端创建临时表，可按ID分组，为每个ID构造专属日期范围条件，批量查询后合并结果，避免全表关联的性能损耗。

代码实现：

result_list = []
for _, row in df_A.iterrows():
    id_val = row['ID']
    start_date = row['Added_Date'].strftime('%Y-%m-%d')
    end_date = (row['Added_Date'] + pd.Timedelta(days=30)).strftime('%Y-%m-%d')
    # 利用分区键缩小查询范围
    query = f'''
    SELECT ID, transaction_date
    FROM transactions
    WHERE 
        ID = {id_val}
        AND ((Year = {row['Added_Date'].year} AND Month >= {row['Added_Date'].month})
             OR (Year = {(row['Added_Date'] + pd.Timedelta(days=30)).year} AND Month <= {(row['Added_Date'] + pd.Timedelta(days=30)).month}))
        AND transaction_date BETWEEN '{start_date}' AND '{end_date}'
    '''
    temp_df = pd.read_sql(query, conn)
    result_list.append(temp_df)

result_df = pd.concat(result_list, ignore_index=True)
print(result_df)

方案3：数据类型预优化（基础性能保障） #

确保两端日期字段类型一致，避免隐式类型转换带来的性能损耗：

• Pandas端：将Added_Date转换为datetime64类型（如示例中df_A['Added_Date'] = pd.to_datetime(df_A['Added_Date'])）
• SQL端：确保transaction_date为日期类型（而非字符串），分区键Year/Month/Day为整数类型

输出结果 #

ID transaction_date
0   1        2023-02-10
1   1        2023-03-01
2   2        2023-04-20
3   2        2023-05-05
4   3        2023-03-25
5   3        2023-04-02

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