如何在Polars中用向量化方案替代map_groups实现波动率分桶？

直接用Polars的窗口函数（.over("symbol")）结合分位数计算和条件判断，全程用表达式实现，完全替代map_groups()，性能会好很多。

步骤分解 #

• 计算每个symbol的分位数阈值：用pl.col("realized_volatility").quantile()配合.over("symbol")，分别算出low_vol和mid_vol（分位数可自定义）。
• 生成波动率分桶：用pl.when().then().when().then().otherwise()的条件逻辑，基于每行的realized_volatility和对应symbol的阈值标记分桶。

完整代码实现 #

import polars as pl

def vol_buckets(df: pl.DataFrame, low_quantile: float = 1/3, mid_quantile: float = 2/3) -> pl.DataFrame:
    return df.with_columns(
        # 按symbol分组计算波动率分位数阈值
        low_vol=pl.col("realized_volatility").quantile(low_quantile).over("symbol"),
        mid_vol=pl.col("realized_volatility").quantile(mid_quantile).over("symbol"),
        # 生成分桶标记
        vol_bucket=pl.when(pl.col("realized_volatility") <= pl.col("low_vol"))
                     .then(pl.lit("low"))
                     .when(pl.col("realized_volatility") <= pl.col("mid_vol"))
                     .then(pl.lit("mid"))
                     .otherwise(pl.lit("high"))
    ).drop("low_vol", "mid_vol")  # 不需要临时阈值列的话可以加这行

关键说明 #

• 窗口函数优势：.over("symbol")确保分位数计算在每个symbol分组内独立执行，是Polars原生向量化操作，性能远高于map_groups()的Python循环逻辑。
• 分位数自定义：通过low_quantile和mid_quantile参数可以灵活调整分桶阈值，比如四分位就改成0.25和0.75。
• 临时列处理：如果不需要保留low_vol和mid_vol，直接链式调用.drop()删除即可。

示例用法 #

# 测试数据
df = pl.DataFrame({
    "symbol": ["A", "A", "A", "B", "B", "B"],
    "realized_volatility": [0.1, 0.2, 0.3, 0.05, 0.15, 0.25]
})

# 生成分桶
result = vol_buckets(df)
print(result)

输出：

shape: (6, 3)
┌────────┬────────────────────┬────────────┐
│ symbol ┆ realized_volatility ┆ vol_bucket │
│ ---    ┆ ---                ┆ ---        │
│ str    ┆ f64                ┆ str        │
╞════════╪════════════════════╪════════════╡
│ A      ┆ 0.1                ┆ low        │
│ A      ┆ 0.2                ┆ mid        │
│ A      ┆ 0.3                ┆ high       │
│ B      ┆ 0.05               ┆ low        │
│ B      ┆ 0.15               ┆ mid        │
│ B      ┆ 0.25               ┆ high       │
└────────┴────────────────────┴────────────┘

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