Pandas百万级患者数据聚合函数运行速度优化方案

核心瓶颈 #

当前耗时过高的本质是实现逻辑的时间复杂度达到了O(患者数*检验记录数)：逐行遍历100万患者时，每次都对几百万条的全量检验表做ID匹配+时间范围过滤，相当于重复扫描全表100万次，这类写法在百万级数据场景下必然出现数天级的耗时，不需要更换分布式集群，单机即可优化到分钟级。

按优先级落地的优化步骤 #

1. 消灭全表扫描与逐行循环（提速100~1000倍，覆盖90%以上耗时） #

• 先对检验结果表做预处理：将检验日期字段转为datetime64类型，设置patient_id为有序索引，pandas对有序索引的定位查询速度比普通列条件筛选快2个数量级以上，预处理代码：

  df['test_date'] = pd.to_datetime(df['test_date'])
  df = df.set_index('patient_id', sort=True)
  

• 完全抛弃DataFrame.apply(axis=1)逐行调用聚合函数的逻辑，改为一次性关联+向量化过滤：先把患者表的ID、首次日期firstAF、时间窗参数X/Y和检验表做内连接，再用向量化条件一次性筛出所有符合时间窗的记录，把100万次单独筛选替换为1次批量过滤：

  # 关联每个患者的时间窗参数到对应的检验记录
  merged = df.merge(
      customers[['patient_id', 'firstAF', 'X', 'Y']],
      left_index=True,
      right_on='patient_id',
      how='inner'
  )
  # 一次性过滤所有时间窗内的记录，无逐行python循环
  time_mask = (merged['test_date'] >= merged['firstAF'] - pd.to_timedelta(merged['X'], unit='d')) & \
              (merged['test_date'] <= merged['firstAF'] + pd.to_timedelta(merged['Y'], unit='d'))
  window_data = merged[time_mask]
  

2. 聚合计算逻辑优化（提速50~100倍） #

• 基础统计量（均值、标准差、最新值、最旧值）全部用pandas原生groupby内置聚合函数计算，这类函数是C层实现，比Python层自定义循环快几十倍。计算前先把window_data按test_date升序排序，分组后直接取首尾值即可得到最旧、最新结果，不需要分组后重复排序：

  window_data = window_data.sort_values('test_date')
  base_stats = window_data.groupby(['patient_id', 'test_id'])['test_value'].agg(
      mean_res='mean',
      std_res='std',
      oldest_res=lambda x: x.iloc[0],
      latest_res=lambda x: x.iloc[-1]
  )
  

• 替换逐组调用scipy.stats.linregress算斜率的逻辑：linregress每次调用都有大量Python层开销，可直接用最小二乘法斜率的代数公式转成向量化计算——对每个分组，先计算检验日期距firstAF的天数作为x轴值，检验结果作为y轴值，通过groupby批量统计sum(x)、sum(y)、sum(x*y)、sum(x**2)、count(n)五个统计量，再代入公式slope = (n*sum(xy) - sum(x)*sum(y)) / (n*sum(x2) - sum(x)**2)即可一次性算出所有分组的回归斜率，计算结果和linregress完全一致，速度提升50倍以上。

3. 引擎与存储优化（额外提速5~10倍） #

• 如果本地内存不足以加载全量数据，可直接替换pandas为polars计算引擎，语法和pandas高度兼容，默认支持多线程并行、懒执行优化、更低的内存占用，同一份逻辑无需大幅修改即可获得5~10倍的速度提升。
• 长期存储可将检验结果表按patient_id分区存为parquet格式，后续计算时可直接跳过无关患者的数据块，进一步减少IO和计算量。

预期收益 #

以上步骤全部落地后，在普通16核32G内存的单机上，100万患者、数百万条检验记录的全量计算耗时可从数天压缩到10分钟以内。

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