如何基于掩码拆分DataFrame/Array并高效计算各分片统计量

现有两个数据结构：

• mat：尺寸为(N,T)，类型为布尔型或浮点型，支持空值
• masks：尺寸为(N,T)，类型为布尔型（无空值），可拆分为T个尺寸为N的掩码

目标是通过每个掩码将mat拆分为T个分片，快速高效地计算每个分片的多组统计量并存储。原遍历掩码的实现存在性能瓶颈，以下是针对核心疑问的解决方案：

1. 替代遍历的向量化优化方案 #

遍历的核心问题是逐次处理每个掩码，未利用numpy/pandas的向量化特性。我们可以通过广播扩展掩码，一次性生成所有分片的有效数据标识，再批量计算统计量，时间复杂度可从O(TNT)优化为O(N*T²)，且充分利用CPU向量指令加速。

以示例中的两个统计量为例，实现代码如下：

import pandas as pd
import numpy as np
import datetime

N = 1000
T = 100
t0 = datetime.date(2022, 1, 1)
t_last = t0 + datetime.timedelta(days=T-1)
dates = pd.date_range(t0, t_last, freq='D')

# 生成原始数据
mat_np = np.random.choice(a=[False, True, np.nan], size=(N, T), p=[0.5, 0.4, 0.1])
masks_np = (mat_np == False)

### 向量化计算stat_1（非空值计数）
# 扩展掩码为(N,T,T)：mask[i,j,k]表示第j个掩码选中第i行，且k >= j
mask_expanded = masks_np[:, :, np.newaxis] & np.tri(T, dtype=bool)[np.newaxis, :, :]
# 按N维度求和，得到(T,T)的统计结果
stat_1 = np.sum(~np.isnan(mat_np[:, np.newaxis, :]) & mask_expanded, axis=0)

### 向量化计算stat_2（累积最大值向前填充后求和）
# 先预处理每行的cummax+ffill
mat_filled = np.full_like(mat_np, np.nan)
for i in range(N):
    row = mat_np[i]
    # 计算累积最大值，处理空值
    cumax = np.maximum.accumulate(np.where(np.isnan(row), -np.inf, row))
    cumax[cumax == -np.inf] = np.nan
    # 向前填充空值
    valid_idx = np.where(~np.isnan(cumax), np.arange(T), 0)
    valid_idx = np.maximum.accumulate(valid_idx)
    mat_filled[i] = cumax[valid_idx]

# 结合掩码计算求和
stat_2 = np.sum(mat_filled[:, np.newaxis, :] * mask_expanded, axis=0)

# 转换为易读的DataFrame格式
results_df = pd.DataFrame(
    np.concatenate([stat_1[..., np.newaxis], stat_2[..., np.newaxis]], axis=2),
    index=dates,
    columns=pd.MultiIndex.from_product([dates, ['stat_1', 'stat_2']])
)

2. 并行计算的实现 #

并行计算适合T极大、单掩码处理逻辑复杂的场景，但要注意数据传输开销：如果单任务计算逻辑简单，并行的额外开销可能超过收益。

用joblib实现并行的代码示例：

from joblib import Parallel, delayed

def process_single_mask(date_idx, mat_df, masks_df):
    date = masks_df.columns[date_idx]
    mat_tmp = mat_df.loc[masks_df.loc[:, date]].loc[:, date:]
    return pd.DataFrame({
        'stat_1': mat_tmp.notna().sum(),
        'stat_2': mat_tmp.cummax(axis=1).ffill(axis=1).sum(),
    })

# 用所有CPU核心并行处理
results_list = Parallel(n_jobs=-1)(
    delayed(process_single_mask)(i, mat_df, masks_df)
    for i in range(len(masks_df.columns))
)

results_df = pd.concat(results_list, keys=masks_df.columns, axis=1)

若使用numpy数组，建议提前将数据转为共享内存（如multiprocessing.Array），避免进程间重复拷贝数据。

3. (N,T,T)张量的矩阵运算方案 #

将分片转为(N,T,T)张量是可行的，但需注意内存占用：当N=1000、T=100时，张量为1e7个元素，float64类型占80MB，完全可控；但T=1000时会达到8GB，内存压力骤增。

实现逻辑参考第一部分的向量化代码，核心是通过广播+布尔索引生成mask_expanded，再利用numpy广播运算批量计算。这种方式的优势是最大化利用CPU计算资源，比遍历快数倍到数十倍，但并非所有统计量都适合：涉及复杂行级逻辑的统计量，需要先预处理每行数据（如示例中的stat_2）。

4. DataFrame vs Array的选择 #

• numpy Array：纯数值计算速度快、内存占用低，向量化操作便捷，但需要手动管理索引和列名，多统计量存储需构造多维数组。
• pandas DataFrame：优势在于索引自动对齐、列名管理便捷、结果可读性高，适合直接输出业务化结果，但性能略低于numpy，大规模数据下差距更明显。

权衡建议：

• 核心计算阶段用numpy数组做向量化/并行计算，结果生成后再转为DataFrame存储展示。
• 若统计量涉及大量行列标签操作（如示例中的日期索引），直接用pandas更省心，但可结合df.to_numpy()转为数组完成核心计算，再转回DataFrame。

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