Polars中优雅聚合行子集：解决when...then聚合警告问题

针对你的需求，推荐使用聚合上下文内的filter操作来实现，这是Polars中更地道、不会触发警告的写法：

import polars as pl

df = pl.DataFrame(dict(
    grp = ['A', 'A', 'A', 'B', 'B', 'B'],
    subgroup = ['x', 'x', 'y', 'x', 'x', 'y'],
    value = [1, 2, 3, 4, 5, 6]
))

# 推荐写法
result = (
    df
    .group_by('grp')
    .agg(
        pl.col('value').mean().alias('mean_all'),
        # 分组内过滤subgroup为x的行，再计算均值
        pl.col('value').filter(pl.col('subgroup') == 'x').mean().alias('mean_x')
    )
)

print(result)

输出结果：

┌─────┬──────────┬────────┐
│ grp ┆ mean_all ┆ mean_x │
│ --- ┆ ---      ┆ ---    │
│ str ┆ f64      ┆ f64    │
╞═════╪══════════╪════════╡
│ A   ┆ 2.0      ┆ 1.5    │
│ B   ┆ 5.0      ┆ 4.5    │
└─────┴──────────┴────────┘

原写法触发警告的原因 #

你之前的pl.when(...).then(...).mean()写法，本质是先逐行生成一个仅保留subgroup=x对应value的Series（其余为null），再对这个Series求均值。虽然结果正确，但Polars的类型检查会将when/then识别为逐行转换操作而非聚合操作，后续的mean()虽为聚合，但整体结构不符合Polars聚合上下文的标准范式，因此触发“非有效聚合”警告。

另一种等价实现（可选） #

如果需要更明确地展示均值的计算逻辑，也可以通过“条件求和/条件计数”的方式实现：

result = (
    df
    .group_by('grp')
    .agg(
        pl.col('value').mean().alias('mean_all'),
        (
            pl.col('value').filter(pl.col('subgroup') == 'x').sum() 
            / pl.col('subgroup').filter(pl.col('subgroup') == 'x').count()
        ).alias('mean_x')
    )
)

这种写法逻辑更直观，但不如第一种简洁，按需选择即可。

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