Polars中when..then..otherwise批量赋值多列及性能优化咨询

问题分析 #

你当前的Polars代码重复执行了三次相同的过滤条件(pl.col("id") == col) & (pl.col("date") >= date)，每次when都会单独计算这个布尔掩码，对于百万级数据集来说，这种冗余计算会极大拖慢执行速度。

优化方案：复用过滤条件，单次判断批量赋值 #

方法1：预定义过滤条件，减少重复计算

先把过滤条件存为变量，后续when直接复用，Polars会只计算一次掩码，避免重复开销：

# 预定义过滤条件，仅计算一次
filter_mask = (pl.col("id") == col) & (pl.col("date") >= date)

elo_df = elo_df.with_columns([
    pl.when(filter_mask)
      .then(pl.lit(new_rating))
      .otherwise(pl.col("elo"))
      .alias("elo"),
    pl.when(filter_mask)
      .then(pl.lit(new_mu))
      .otherwise(pl.col("true_skill_mu"))
      .alias("true_skill_mu"),
    pl.when(filter_mask)
      .then(pl.lit(new_sigma))
      .otherwise(pl.col("true_skill_sigma"))
      .alias("true_skill_sigma")
])

方法2：Struct批量返回，单次判断完成多列赋值

更高效的方式是用struct将多个新值打包，通过一次when判断完成所有列的赋值，进一步降低判断开销：

filter_mask = (pl.col("id") == col) & (pl.col("date") >= date)

elo_df = elo_df.with_columns(
    pl.when(filter_mask)
      .then(
          pl.struct(
              pl.lit(new_rating).alias("elo"),
              pl.lit(new_mu).alias("true_skill_mu"),
              pl.lit(new_sigma).alias("true_skill_sigma")
          )
      )
      .otherwise(
          pl.struct(
              pl.col("elo"),
              pl.col("true_skill_mu"),
              pl.col("true_skill_sigma")
          )
      )
      .alias("updated_cols")
).unnest("updated_cols")

这种方式只执行一次条件判断，将符合条件的行的多列值一次性替换，性能提升更明显。

额外优化建议 #

1. 为过滤列建立索引/排序：
   如果你的id和date列常用于过滤，可通过排序或设置索引加速查询：

   # 若id列有序，设置为排序状态（Polars会优化过滤查询）
   elo_df = elo_df.set_sorted("id")
   # 或建立联合索引（Polars 0.19+版本支持）
   elo_df = elo_df.with_index(["id", "date"])
   

2. 使用Lazy API优化查询计划：
   将代码切换为Lazy模式，Polars会自动优化查询逻辑，合并重复计算：

   elo_df = (
       elo_df.lazy()
       .with_columns(
           elo=pl.when(filter_mask).then(new_rating).otherwise(pl.col("elo")),
           true_skill_mu=pl.when(filter_mask).then(new_mu).otherwise(pl.col("true_skill_mu")),
           true_skill_sigma=pl.when(filter_mask).then(new_sigma).otherwise(pl.col("true_skill_sigma"))
       )
       .collect()
   )
   

3. 确保新值为标量：
   确认new_rating、new_mu、new_sigma是单个标量值，而非Series，避免不必要的广播计算。

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