如何加速基于Datetime Index的时间边界筛选迭代代码？

我有一个带Datetime Index的DataFrame df：

col1     col2    col3
2020-01-02 08:50:00    360.0    -131.0  -943.0
2020-01-02 08:52:01    342.0    -130.0  -1006.0
2020-01-02 08:55:04    321.0    -130.0  -997.0
... ... ... ...
2022-01-03 14:44:56    1375.0   -91.0   -728.0
2022-01-03 14:50:57    1381.0   -118.0  -692.0
2022-01-03 14:50:58    1382.0   -115.0  -697.0
2022-01-03 14:50:59    1390.0   -111.0  -684.0
2022-01-03 14:55:58    1442.0   -106.0  -691.0

我需要实现一个函数，筛选出与已选前一个索引间隔不小于指定时长（如5分钟）的索引，示例调用为masked_df = time_mask(df.index, pd.Timedelta(minutes=5))，要求选取第一个符合间隔要求的索引。以下是我的迭代实现代码：

def get_clean_ix_from_rolling(idx, time_delt):
    
    clean_ix = []
    prev_ix = idx[0]
    clean_ix.append(prev_ix)
    for i, x in enumerate(idx):
        if((x-prev_ix) >= time_delt):
            clean_ix.append(x)
            prev_ix = x

    ix = pd.to_datetime(clean_ix)
    return ix

请问如何加速这段代码？

1. 用Numba编译循环 #

Python原生循环速度慢，用Numba的JIT编译可以直接加速现有逻辑，不需要大幅修改代码：

import numba
import pandas as pd

@numba.jit(nopython=True)
def numba_filter(idx_values, time_delt_ns):
    clean_ix = [idx_values[0]]
    prev_ix = idx_values[0]
    for x in idx_values[1:]:
        if (x - prev_ix) >= time_delt_ns:
            clean_ix.append(x)
            prev_ix = x
    return clean_ix

def get_clean_ix_from_rolling(idx, time_delt):
    idx_values = idx.astype('datetime64[ns]').values.astype('int64')
    time_delt_ns = time_delt.value
    clean_ix_ns = numba_filter(idx_values, time_delt_ns)
    return pd.to_datetime(clean_ix_ns)

Numba会把循环编译成机器码，处理大数据集时速度能提升几十倍。

2. 向量化实现（动态基准版） #

如果不想用第三方库，可以优化循环逻辑，减少不必要的操作，同时利用pandas的数组操作提升效率：

def get_clean_ix_from_rolling(idx, time_delt):
    idx_arr = idx.values
    keep = [True]
    last_keep = idx_arr[0]
    
    for x in idx_arr[1:]:
        if (x - last_keep) >= time_delt:
            keep.append(True)
            last_keep = x
        else:
            keep.append(False)
    
    return idx[keep]

这里直接操作底层numpy数组，比遍历idx的datetime对象更快，同时去掉了enumerate的冗余索引获取。

3. 固定窗口场景用resample（快速替代方案） #

如果你的需求可以简化为「按固定时间窗口取第一个点」（比如每5分钟取窗口内第一个索引），可以直接用pandas的resample，完全向量化，速度极快：

def get_clean_ix_from_rolling(idx, time_delt):
    freq = f"{time_delt.total_seconds()}S"
    resampled = pd.Series(index=idx).resample(freq).first().dropna()
    return resampled.index

注意：这个方法是按固定时间窗口筛选，和原逻辑的「以上一个选中点为动态基准」有区别，需要根据实际需求选择。

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