如何在Polars中按分组对多列计算均值与NA占比

原始数据 #

import polars as pl

polars_df = pl.DataFrame({
    "name": ["A","B","C"],
    "group": ["a","a","b"],
    "val1": [1, None, 3],
    "val2": [1, 5, None],
    "val3": [None, None, 3],
})

需求说明 #

按group列分组，对所有以val开头的列计算两个指标：

• 组内所有val列的整体均值（所有非NA值的平均）
• NA值占比（组内val列所有单元格中NA数量与总单元格数的比值）

Pandas低效实现（参考） #

df = polars_df.to_pandas()

pd.concat([
    df.groupby(["group"]).apply(lambda g: g.filter(like="val").mean().mean()).rename("mean"),
    df.groupby(["group"]).apply(lambda g: g.filter(like="val").isna().sum().sum() / (g.filter(like="val").shape[0] * g.filter(like="val").shape[1])).rename("percentage_na")
], axis=1)

Polars优化方案 #

方案一：转长格式后聚合（简洁高效） #

利用melt将所有val开头的列转为长格式，再分组计算，代码简洁且性能最优：

result = (
    polars_df
    .melt(id_vars="group", value_vars=pl.col("^val.*$"), variable_name="val_col")
    .group_by("group")
    .agg(
        mean=pl.col("value").mean(),
        percentage_na=pl.col("value").is_null().mean()
    )
)

print(result)

输出结果：

shape: (2, 3)
┌───────┬──────┬──────────────────┐
│ group ┆ mean ┆ percentage_na    │
│ ---   ┆ ---  ┆ ---              │
│ str   ┆ f64  ┆ f64              │
╞═══════╪══════╪══════════════════╡
│ a     ┆ 2.0  ┆ 0.5              │
│ b     ┆ 3.0  ┆ 0.333333333333333│
└───────┴──────┴──────────────────┘

说明：

• melt将多列val数据转为单列value，简化分组后的聚合逻辑
• pl.col("value").mean()直接计算所有非NA值的整体均值
• pl.col("value").is_null().mean()利用布尔值均值等于True占比的特性，直接得到NA值占比

方案二：不转格式直接聚合（适合保留宽表习惯场景） #

如果不想转换数据格式，可直接在分组聚合中完成计算：

result = (
    polars_df
    .group_by("group")
    .agg(
        mean=pl.concat_list(pl.col("^val.*$")).flatten().drop_nulls().mean(),
        percentage_na=(pl.col("^val.*$").is_null().sum().sum() / (pl.col("^val.*$").len() * pl.col("^val.*$").count()))
    )
)

print(result)

说明：

• pl.concat_list(pl.col("^val.*$")).flatten().drop_nulls().mean()：将组内所有val列的值合并为列表，展平后去掉空值再取均值
• 分母pl.col("^val.*$").len() * pl.col("^val.*$").count()计算组内val列的总单元格数，分子计算NA总数，两者相除得到占比

方案优势 #

Polars方案全程使用向量化操作，避免了Pandas中apply带来的Python循环开销，代码更简洁且性能更优，完全匹配需求中的期望结果。

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