如何加速基于Pandas的客户持仓暴露计算代码？

核心优化方向 #

• 消除逐客户循环，采用全量数据批量处理
• 用Pandas向量化操作替代低效的行级循环
• 优化时间轴生成逻辑，避免重复计算

具体实现步骤 #

1. 合并所有客户交易数据 #

先把字典里的所有DataFrame合并成一个大表，新增client字段标记归属：

import pandas as pd

# 假设你的原始字典名为 client_trade_data
all_trades = pd.concat(
    [df.assign(client=name) for name, df in client_trade_data.items()],
    ignore_index=True
)

2. 预处理交易数据 #

确保时间字段为datetime类型，计算持仓变动（买入为正，卖出为负，根据实际业务调整）：

# 转换时间格式（如果还没转）
all_trades['trade_time'] = pd.to_datetime(all_trades['trade_time'], errors='coerce')

# 计算持仓变动：假设买入时quantity为正，卖出为负；若不是则自行转换
all_trades['pos_change'] = all_trades['quantity']
# 示例：如果有side字段，买入是buy，卖出是sell
# all_trades['pos_change'] = all_trades['quantity'] * all_trades['side'].map({'buy':1, 'sell':-1})

3. 生成全局5分钟时间网格+客户标的组合 #

构建完整的时间轴和客户-标的组合，避免每个客户单独生成：

# 确定全局时间范围，对齐到5分钟刻度
start_time = all_trades['trade_time'].min().floor('5min')
end_time = all_trades['trade_time'].max().ceil('5min')
time_axis = pd.date_range(start=start_time, end=end_time, freq='5min')

# 获取所有唯一的客户-标的组合
client_symbol_pairs = all_trades[['client', 'symbol']].drop_duplicates()

# 生成全量笛卡尔积网格
full_grid = client_symbol_pairs.merge(
    pd.DataFrame({'time': time_axis}),
    how='cross'
)

4. 合并数据并计算累积持仓 #

将交易数据与全量网格合并，按客户-标的分组计算累积持仓：

# 合并交易数据，空值填充为0（对应无交易的时间点）
merged_data = full_grid.merge(
    all_trades[['client', 'symbol', 'trade_time', 'pos_change']],
    left_on=['client', 'symbol', 'time'],
    right_on=['client', 'symbol', 'trade_time'],
    how='left'
).fillna({'pos_change': 0})

# 按客户-标的分组，计算累积持仓
merged_data['position'] = merged_data.groupby(['client', 'symbol'])['pos_change'].cumsum()

# 整理成最终的宽表格式（可选，根据需求调整）
final_position = merged_data.pivot(
    index='time',
    columns=['client', 'symbol'],
    values='position'
).fillna(0)

额外加速技巧 #

• 启用Pandas的PyArrow后端，提升数据处理速度：

  pd.options.mode.dtype_backend = 'pyarrow'
  pd.set_option('mode.copy_on_write', True)
  

• 若数据量超1000万行，用Dask替代Pandas，自动并行处理多CPU核心
• 避免使用apply()、iterrows()等行级循环操作，全部用Pandas内置的向量化方法

为什么比原有代码快？ #

1. 去掉了逐客户循环，减少了Python层面的循环开销
2. 时间轴和客户-标的网格一次性生成，避免重复计算
3. 所有核心操作都是Pandas底层优化的C级向量化运算，远快于纯Python循环

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