如何优化大型10000×10000时间序列DataFrame的滑动窗口计算？

问题背景 #

需要处理10000×10000的时间序列DataFrame，针对每一行，基于最近255条非零行执行自定义函数myFunc计算，但当前Pandas和NumPy实现性能均不理想。

当前Pandas实现 #

df = self.df
df = df.replace(0, np.nan)
df = df.dropna(how='all', axis=0)
result = df.rolling(window=255).apply(myFunc, raw=True)
res = pd.DataFrame(index=self.df.index, columns=self.df.columns)
res.update(result)
return res[(res != 0).any(axis=1)]

当前NumPy实现 #

df_np = self.df.to_numpy()
result = np.zeros_like(df_np)
window_size = 255
indices = np.where(~np.all(df_np == 0, axis=1))
for i in indices[0]:
    if i >= window_size:
        valid_indices = indices[0][np.where(indices[0] <= i)[0]][-window_size:]
        if len(valid_indices) >= window_size:
            result[i] = np.apply_along_axis(myFunc, 0, df_np[valid_indices, :])
res = pd.DataFrame(result, index=self.df.index, columns=self.df.columns)
return res[(res != 0).any(axis=1)]

优化方案 #

1. 预提取有效行索引，避免重复查找 #

原NumPy实现中每次循环都重新查找valid_indices，效率极低。提前提取所有非零行的索引数组，直接基于该数组做滑动窗口：

df_np = self.df.to_numpy()
valid_rows = np.where(~np.all(df_np == 0, axis=1))[0]
window_size = 255
result = np.zeros_like(df_np)

# 遍历有效行中满足窗口大小的位置
for k in range(window_size-1, len(valid_rows)):
    # 直接取最近255条有效行
    window_rows = valid_rows[k-window_size+1 : k+1]
    result[valid_rows[k]] = myFunc(df_np[window_rows])

res = pd.DataFrame(result, index=self.df.index, columns=self.df.columns)
return res[(res != 0).any(axis=1)]

2. 向量化改造myFunc，替代apply_along_axis #

np.apply_along_axis本质是逐列循环，若myFunc能改造成支持二维输入的向量化函数，性能会大幅提升。例如：

• 若myFunc是计算窗口均值，直接用np.mean(window, axis=0)替代自定义循环
• 若为自定义逻辑，尽量用NumPy矩阵运算实现，避免逐元素/逐列处理

3. 用NumPy滑动窗口视图批量处理 #

利用np.lib.stride_tricks.as_strided创建滑动窗口视图（无数据复制），批量生成所有窗口后一次性计算：

df_np = self.df.to_numpy()
valid_rows = np.where(~np.all(df_np == 0, axis=1))[0]
window_size = 255
n_valid = len(valid_rows)
result = np.zeros_like(df_np)

if n_valid >= window_size:
    # 创建滑动窗口视图：形状为(窗口数, 窗口大小, 列数)
    window_view = np.lib.stride_tricks.as_strided(
        df_np[valid_rows],
        shape=(n_valid - window_size + 1, window_size, df_np.shape[1]),
        strides=(df_np.strides[0] * window_size, df_np.strides[0], df_np.strides[1])
    )
    # 假设myFunc支持三维输入，沿窗口维度(axis=1)计算
    batch_results = np.apply_along_axis(myFunc, 1, window_view)
    # 将结果映射回原数组对应行
    result[valid_rows[window_size-1:]] = batch_results

res = pd.DataFrame(result, index=self.df.index, columns=self.df.columns)
return res[(res != 0).any(axis=1)]

注意：使用as_strided需确保维度计算准确，避免内存越界。

4. 并行化处理窗口计算 #

若myFunc无法完全向量化，可通过多进程并行处理每个窗口，利用多核CPU资源：

from joblib import Parallel, delayed

df_np = self.df.to_numpy()
valid_rows = np.where(~np.all(df_np == 0, axis=1))[0]
window_size = 255
result = np.zeros_like(df_np)

def process_single_window(k):
    window = df_np[valid_rows[k-window_size+1 : k+1]]
    return myFunc(window)

if len(valid_rows) >= window_size:
    # 并行处理所有符合条件的窗口
    batch_results = Parallel(n_jobs=-1)(
        delayed(process_single_window)(k)
        for k in range(window_size-1, len(valid_rows))
    )
    result[valid_rows[window_size-1:]] = np.array(batch_results)

res = pd.DataFrame(result, index=self.df.index, columns=self.df.columns)
return res[(res != 0).any(axis=1)]

注意：并行化存在启动开销，适合myFunc计算成本较高的场景。

5. 减少不必要的数据复制 #

• 原Pandas代码中replace(0, np.nan)+dropna会生成新DataFrame，直接用NumPy判断非零行更高效
• 尽量在NumPy数组上完成核心计算，减少Pandas与NumPy之间的转换开销

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