如何高效实现Pandas分组DataFrame的时间缺口补零重索引

问题描述 #

现有带时间戳的分组数据集，需为每个ID补全秒级时间缺口，缺失的ID和Col值填0。原采用循环+groupby的实现方式在20万行数据上耗时约1小时，面对5000万行（66列）的真实数据集效率严重不足，需寻求向量化替代方案。

原始数据集 #

Time                ID   Col
0   2023-02-08 06:00:00 1   10.01
1   2023-02-08 06:00:01 1   21.10
2   2023-02-08 06:00:05 1   21.02
3   2023-02-08 06:10:01 2   54.02
4   2023-02-08 06:10:03 2   56.02
5   2023-02-08 06:10:05 2   65.03

数据集生成代码 #

from datetime import datetime
import pandas as pd

data = {'Time': ["2023-02-08 06:00:00","2023-02-08 06:00:01","2023-02-08 06:00:05","2023-02-08 06:10:01","2023-02-08 06:10:03","2023-02-08 06:10:05"],
        'ID': ['1','1','1','2','2','2'],
        'Col': [10.01,21.1,21.02,54.02,56.02,65.03]}
df = pd.DataFrame(data)
df['Time'] = pd.to_datetime(df['Time'])

期望输出 #

Time            ID  Col
2023-02-08 06:00:00 1   10.01
2023-02-08 06:00:01 1   21.10
2023-02-08 06:00:02 0   0.00
2023-02-08 06:00:03 0   0.00
2023-02-08 06:00:04 0   0.00
2023-02-08 06:00:05 1   21.02
2023-02-08 06:10:01 2   54.02
2023-02-08 06:10:02 0   0.00
2023-02-08 06:10:03 2   56.02
2023-02-08 06:10:04 0   0.00
2023-02-08 06:10:05 2   65.03

向量化实现方案 #

以下方案利用pandas内置的向量化操作替代显式循环，大幅提升处理效率：

import pandas as pd

# 生成原始数据（同上）
data = {'Time': ["2023-02-08 06:00:00","2023-02-08 06:00:01","2023-02-08 06:00:05","2023-02-08 06:10:01","2023-02-08 06:10:03","2023-02-08 06:10:05"],
        'ID': ['1','1','1','2','2','2'],
        'Col': [10.01,21.1,21.02,54.02,56.02,65.03]}
df = pd.DataFrame(data)
df['Time'] = pd.to_datetime(df['Time'])

# 1. 为每个ID生成完整的秒级时间序列
def build_full_time_group(group):
    # 获取当前ID的时间起止范围
    time_start = group['Time'].min()
    time_end = group['Time'].max()
    # 生成连续秒级时间戳
    full_time_seq = pd.date_range(start=time_start, end=time_end, freq='S')
    # 返回包含时间和对应ID的框架
    return pd.DataFrame({'Time': full_time_seq, 'ID': group['ID'].iloc[0]})

# 按ID分组生成完整时间框架
full_time_frame = df.groupby('ID').apply(build_full_time_group).reset_index(drop=True)

# 2. 合并原始数据与完整时间框架，填充缺失值为0
final_result = full_time_frame.merge(df, on=['Time', 'ID'], how='left').fillna(0)

# 调整列顺序匹配期望输出
final_result = final_result[['Time', 'ID', 'Col']]
print(final_result)

效率提升说明 #

• 避免Python显式循环：groupby.apply结合pd.date_range利用pandas内部C优化逻辑，比纯Python循环快数十倍
• 合并与填充均为向量化操作：merge和fillna都是pandas高度优化的向量化方法，处理千万级数据的性能远高于逐行处理

大数据集额外优化建议 #

针对5000万行的超大数据集，可进一步优化：

• 数据类型优化：将ID设为category类型，减少内存占用并加速分组操作：df['ID'] = df['ID'].astype('category')
• 并行处理：使用dask.dataframe替代pandas，自动拆分数据集进行并行计算，适合内存无法容纳的超大数据
• 预计算时间范围：若所有ID的时间范围存在规律，可预先生成全局时间序列再按ID匹配，减少分组计算的开销

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