大数据集下Pandas时间区间关联内存优化及SQL方案修复

一、低内存Python实现方案 #

原代码中df1.merge(df2, how='outer', on="location")会产生全量笛卡尔积（同location的所有df1行和df2行两两匹配），大数据集下会瞬间占满内存。可以通过按location分组处理避免全量笛卡尔积，大幅降低内存占用：

import pandas as pd
import numpy as np

# 确保时间列转换为datetime类型
df1['timestamp'] = pd.to_datetime(df1['timestamp'])
df2['time_start'] = pd.to_datetime(df2['time_start'])
df2['time_end'] = pd.to_datetime(df2['time_end'])

# 按location分组处理，每个组单独匹配降雨事件
result_list = []
for loc, group_df1 in df1.groupby('location'):
    # 获取当前location对应的所有降雨事件
    group_df2 = df2[df2['location'] == loc].copy()
    
    if group_df2.empty:
        # 无对应降雨事件，直接标记为0
        group_df1['quantity_rain'] = 0
        result_list.append(group_df1)
        continue
    
    # 同location内匹配时间区间
    merged = group_df1.merge(group_df2, on='location')
    merged['quantity_rain'] = merged['event'].where(
        merged['timestamp'].between(merged['time_start'], merged['time_end']),
        0
    )
    
    # 处理一个timestamp匹配多个降雨事件的情况（这里取第一个匹配值，可按需调整）
    merged = merged.groupby(['timestamp', 'factor', 'location'])['quantity_rain'].first().reset_index()
    result_list.append(merged)

# 合并所有分组结果
final_df = pd.concat(result_list, ignore_index=True)

方案优势 #

• 仅在同location的小数据子集内做匹配，避免全量笛卡尔积
• 内存占用与单个location的最大数据量正相关，而非整个数据集的乘积
• 无需额外依赖库，纯pandas实现

二、SQL方案代码修复 #

原SQL代码存在表名重复覆盖和连接对象不统一两个核心错误，修复后的代码如下：

%load_ext sql
import sqlite3
import pandas as pd

# 创建磁盘数据库连接，同时让魔法命令复用该连接
conn = sqlite3.connect('test_database')
%sql sqlite:///test_database

# 将DataFrame写入独立的SQL表，避免覆盖
df1.to_sql('df1', conn, if_exists='replace', index=False)
df2.to_sql('df2', conn, if_exists='replace', index=False)

# 执行区间匹配查询，无匹配时标记为0
query = """
SELECT 
    df1.timestamp,
    df1.factor,
    df1.location,
    COALESCE(df2.event, 0) AS quantity_rain
FROM df1
LEFT JOIN df2 
    ON df1.location = df2.location 
    AND df1.timestamp BETWEEN df2.time_start AND df2.time_end
"""

# 执行查询并转换为DataFrame
final_df = %sql $query
final_df = final_df.DataFrame()

修复要点 #

1. 修正表名：将df1和df2分别写入df1和df2表，避免原代码中df2覆盖df1表的问题
2. 统一连接对象：让%sql魔法命令连接到磁盘上的test_database，与sqlite3.connect创建的连接保持一致，确保能找到表
3. 优化查询逻辑：使用LEFT JOIN保留df1的所有记录，用COALESCE将无匹配的NULL值替换为0，符合需求
4. 处理多匹配场景：如果一个timestamp对应多个降雨事件，可在SQL中添加GROUP BY或聚合函数（如MAX(df2.event)）调整结果

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