如何用JSON字典中的可变筛选条件高效过滤股票价格DataFrame？

现有一个以日期为索引、股票代码为列的股票价格DataFrame：

import pandas as pd
import numpy as np

df = pd.DataFrame(np.arange(9).reshape((3,3)),
                  columns=['stock1','stock2','stock3'],
                  index=['2023-01-01','2023-01-02','2023-01-03'])

输出结构：

stock1  stock2  stock3
2023-01-01       0       1       2
2023-01-02       3       4       5
2023-01-03       6       7       8

同时有一份按日期存储对应股票名称列表的JSON数据：

filter_dict = {
    '2023-01-01': ['stock1','stock3','stock4', 'stock5'],
    '2023-01-02': ['stock2','stock4'],
    '2023-01-03': ['stock1','stock2','stock3','stock5']
}

需求是过滤DataFrame，让每个日期仅保留JSON列表中存在的股票对应值（DataFrame不含的股票直接忽略），最终得到如下结果：

stock1  stock2  stock3
2023-01-01       0     NaN       2
2023-01-02     NaN       4     NaN
2023-01-03       6       7       8

此前尝试pd.mask、pd.where等方法无效，仅逐行遍历可行，但大数据集下速度极慢，需要更高效的实现方式。

方法一：构建布尔掩码矩阵 #

通过构建与原DataFrame结构一致的布尔矩阵，用矢量化操作完成过滤，避免逐行遍历：

# 获取原DataFrame的列和索引
all_stocks = df.columns
all_dates = df.index

# 创建全False的布尔DataFrame，结构与原df对齐
mask = pd.DataFrame(False, index=all_dates, columns=all_stocks)

# 遍历过滤字典，标记需保留的位置
for date, stocks in filter_dict.items():
    # 只保留原DataFrame中存在的股票
    valid_stocks = [s for s in stocks if s in all_stocks]
    mask.loc[date, valid_stocks] = True

# 应用掩码过滤数据
result = df.where(mask)

方法二：利用Series展开与索引匹配 #

更简洁的矢量化实现，依赖Pandas的索引操作优化：

# 将过滤字典转为Series并展开为(date, stock)行数据
filter_series = pd.Series(filter_dict).explode()

# 生成布尔掩码：检查每个(date, stock)是否在过滤列表内
mask = df.stack().index.isin(filter_series.reset_index().values.tolist())

# 将掩码还原为原DataFrame结构并过滤
result = df.where(mask.unstack())

方案说明 #

两种方法均采用矢量化操作，底层基于C语言实现，比Python层面的逐行遍历效率高一个数量级以上，完全适配大数据场景。

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