Pandas优化：百万级DataFrame搜索函数性能提升方案

问题场景 #

现有一个包含4,827,996行、55列的Pandas DataFrame，需要为每行新增rbt列，规则如下：

• 统计当前行IPorig对应的、State为RST、且datetime在当前行时间前300秒内的记录数
• 若统计数量>10，rbt值为1；否则为0

原实现使用df.apply结合自定义函数，因逐行全表扫描导致运行速度极慢，需优化。

原实现代码 #

def rbt(saddr, state, stime):
    sub_df = df.loc[(df['IPorig'] == saddr) & (df['State'] == 'RST') & (df['datetime'] < stime) & (df['datetime'] > (stime - 300))]
    count = len(sub_df)
    return 1 if count > 10 else 0

df['rbt'] = df.apply(lambda x: rbt(x['IPorig'], x['State'], x['datetime']), axis=1)

原方案低效原因 #

df.apply逐行处理，相当于对480w行数据做480w次全表过滤，时间复杂度为O(n²)，完全不适合百万级数据集。

优化方案 #

思路核心 #

先预处理所有RST记录，通过分组+滑动窗口计数替代逐行扫描，将时间复杂度降至O(n log n)（主要来自排序操作）。

步骤1：提取并预处理RST记录 #

先筛选出所有State为RST的记录，按IPorig分组后按datetime排序（确保时间有序，为滑动窗口做准备）：

import pandas as pd

# 提取RST记录并按IP+时间排序
rst_df = df[df['State'] == 'RST'].sort_values(['IPorig', 'datetime'])

步骤2：分组计算滑动窗口内的RST数量 #

利用Pandas的滚动窗口功能，对每个IPorig分组统计300秒内的RST记录数（需排除当前行本身）：

# 转换datetime为时间类型（若原始数据是时间戳，需指定unit='s'）
rst_df['datetime'] = pd.to_datetime(rst_df['datetime'], unit='s')
# 设置datetime为索引，方便使用时间滚动窗口
rst_df = rst_df.set_index('datetime')

# 按IP分组，计算300秒窗口内的记录数
rst_df['rst_count'] = rst_df.groupby('IPorig')['IPorig'].rolling('300s').count().reset_index(level=0, drop=True)
# 减去当前行，得到当前时间之前300秒内的RST数量
rst_df['rst_count'] = rst_df['rst_count'] - 1

# 恢复原索引，方便后续合并
rst_df = rst_df.reset_index()

步骤3：合并结果到原DataFrame #

将统计结果合并回原表，填充空值后生成rbt列：

# 合并统计结果到原DataFrame（按id匹配）
df = df.merge(rst_df[['id', 'rst_count']], on='id', how='left')
# 无匹配记录的行填充为0
df['rst_count'] = df['rst_count'].fillna(0)
# 生成rbt列：计数>10则为1，否则为0
df['rbt'] = (df['rst_count'] > 10).astype(int)

进阶加速：Numba优化 #

如果上述方案仍有性能瓶颈，可使用Numba对分组后的滑动计数进行JIT编译，进一步提升速度：

import numba
import numpy as np

@numba.jit(nopython=True)
def count_rst_in_window(times):
    n = len(times)
    counts = np.zeros(n, dtype=np.int64)
    for i in range(n):
        # 用二分查找快速定位窗口左边界
        left = np.searchsorted(times, times[i] - 300, side='left')
        counts[i] = i - left
    return counts

# 分组处理每个IP的时间序列
def process_group(group):
    times = group['datetime'].values
    group['rst_count'] = count_rst_in_window(times)
    return group

# 重新处理RST记录
rst_df = df[df['State'] == 'RST'].sort_values(['IPorig', 'datetime'])
rst_df = rst_df.groupby('IPorig', group_keys=False).apply(process_group)

# 后续合并步骤同前
df = df.merge(rst_df[['id', 'rst_count']], on='id', how='left')
df['rst_count'] = df['rst_count'].fillna(0)
df['rbt'] = (df['rst_count'] > 10).astype(int)

优化效果 #

原方案的O(n²)复杂度被降至O(n log n)，对于480w行数据，运行速度可提升数十至数百倍，完全满足大数据量处理需求。

内容的提问来源于stack exchange，提问作者Willen Borges Coelho