如何用Polars DataFrame计算序列分位数（含分组与非分组）

原始数据定义 #

import polars as pl

df = pl.from_repr("""
┌─────┬──────────┐
│ a   ┆ outcome  │
│ --- ┆ ---      │
│ i64 ┆ f64      │
╞═════╪══════════╡
│ 2   ┆ 0.17745  │
│ 2   ┆ 0.712477 │
│ 2   ┆ 0.038308 │
│ 1   ┆ 0.886266 │
│ 2   ┆ 0.578249 │
│ 1   ┆ 0.80318  │
└─────┴──────────┘
""")

方案1：不分组（全局计算分位数） #

先定义需要计算的分位数列表，通过Polars的quantile函数计算全局分位数，再整理成指定的DataFrame格式：

# 定义目标分位数
quantiles = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]

# 计算并整理结果
global_quantile_df = (
    df
    .select(pl.col("outcome").quantile(quantiles).alias("value"))
    .with_columns(pl.Series(quantiles).alias("quantile"))
    .select("quantile", "value")
)

# 输出结果
print(global_quantile_df)

输出示例：

┌──────────┬──────────┐
│ quantile ┆ value    │
│ ---      ┆ ---      │
│ f64      ┆ f64      │
╞══════════╪══════════╡
│ 0.1      ┆ 0.038308 │
│ 0.2      ┆ 0.17745  │
│ 0.3      ┆ 0.17745  │
│ 0.4      ┆ 0.578249 │
│ 0.5      ┆ 0.578249 │
│ 0.6      ┆ 0.712477 │
│ 0.7      ┆ 0.712477 │
│ 0.8      ┆ 0.80318  │
│ 0.9      ┆ 0.886266 │
│ 1.0      ┆ 0.886266 │
└──────────┴──────────┘

方案2：按变量分组计算分位数 #

按a列分组后，对每个分组单独计算分位数，再通过explode将列表型结果展开为多行：

# 按a列分组计算分位数并整理格式
grouped_quantile_df = (
    df
    .group_by("a")
    .agg(
        pl.col("outcome").quantile(quantiles).alias("value"),
        pl.Series(quantiles).alias("quantile")
    )
    .explode(["quantile", "value"])
    .select("a", "quantile", "value")
)

# 输出结果
print(grouped_quantile_df)

输出示例：

┌─────┬──────────┬──────────┐
│ a   ┆ quantile ┆ value    │
│ --- ┆ ---      ┆ ---      │
│ i64 ┆ f64      ┆ f64      │
╞═════╪══════════╪══════════╡
│ 1   ┆ 0.1      ┆ 0.80318  │
│ 1   ┆ 0.2      ┆ 0.80318  │
│ 1   ┆ 0.3      ┆ 0.80318  │
│ 1   ┆ 0.4      ┆ 0.80318  │
│ 1   ┆ 0.5      ┆ 0.80318  │
│ 1   ┆ 0.6      ┆ 0.886266 │
│ 1   ┆ 0.7      ┆ 0.886266 │
│ 1   ┆ 0.8      ┆ 0.886266 │
│ 1   ┆ 0.9      ┆ 0.886266 │
│ 1   ┆ 1.0      ┆ 0.886266 │
│ 2   ┆ 0.1      ┆ 0.038308 │
│ 2   ┆ 0.2      ┆ 0.038308 │
│ 2   ┆ 0.3      ┆ 0.17745  │
│ 2   ┆ 0.4      ┆ 0.17745  │
│ 2   ┆ 0.5      ┆ 0.578249 │
│ 2   ┆ 0.6      ┆ 0.578249 │
│ 2   ┆ 0.7      ┆ 0.712477 │
│ 2   ┆ 0.8      ┆ 0.712477 │
│ 2   ┆ 0.9      ┆ 0.712477 │
│ 2   ┆ 1.0      ┆ 0.712477 │
└─────┴──────────┴──────────┘

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