Polars group_by中if-else的正确实现及警告原因解析

警告产生的原因 #

Polars 的 group_by.agg() 要求传入的所有表达式必须是合法的聚合操作——即每个表达式都要能为每个分组生成唯一的单个值。你原代码中，pl.col("a").first().is_null() 作为 when 的判断条件时，没有被 Polars 识别为标准的聚合逻辑（尽管实际运行时它是按组执行的）。Polars 无法确保这个判断逻辑的输出行数和分组后的行数严格匹配，因此抛出警告，提示这种行为属于实验性特性，后续版本可能会调整。

更优实现方式 #

方式一：用聚合函数明确标记判断条件 #

直接在 when 的判断中使用 pl.first() 聚合函数，让 Polars 明确识别这是按组的聚合判断，消除警告：

df = pl.DataFrame({"a": [None, 1, 2, 3, 4], "b": [1, 1, 2, 2, 2]})

result = df.group_by("b").agg(
    pl.when(pl.first("a").is_null())
    .then(None)
    .otherwise(pl.mean("a"))
    .alias("a_mean")
)

print(result)

输出结果：

shape: (2, 2)
┌─────┬─────────┐
│ b   ┆ a_mean  │
│ --- ┆ ---     │
│ i64 ┆ f64     │
╞═════╪═════════╡
│ 1   ┆ null    │
│ 2   ┆ 3.0     │
└─────┴─────────┘

方式二：拆分聚合与判断逻辑（更清晰） #

先聚合出组内的首个值和均值，再通过 with_columns 完成条件判断，适合逻辑更复杂的场景：

df = pl.DataFrame({"a": [None, 1, 2, 3, 4], "b": [1, 1, 2, 2, 2]})

result = df.group_by("b").agg(
    first_a=pl.first("a"),
    mean_a=pl.mean("a")
).with_columns(
    a_mean=pl.when(pl.col("first_a").is_null()).then(None).otherwise(pl.col("mean_a"))
).drop("first_a", "mean_a")

print(result)

输出结果和方式一完全一致，且无警告。

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