如何在Polars中使用GroupBy自定义函数评估分段预测效果？

对应你在Pandas中使用的groupby.apply逻辑，Polars提供了**map_groups**方法（推荐）和apply方法来实现相同功能，其中map_groups性能更优且类型更安全，是处理这类分组计算的首选。

基础实现（单指标） #

假设你的数据集df包含seg（分段标识）、true（实际值）、pred（预测值）三列，先定义你的评估函数，比如计算MAE：

import polars as pl

def eval_func(true: pl.Series, pred: pl.Series) -> float:
    # 示例：计算平均绝对误差
    return (true - pred).abs().mean()

使用map_groups按seg分组计算：

result = df.group_by("seg").map_groups(
    lambda group: pl.DataFrame({
        "seg": [group["seg"][0]],  # 保留分组标识
        "eval_result": [eval_func(group["true"], group["pred"])]
    })
)

多指标评估 #

如果需要同时计算多个评估指标（比如MAE、RMSE），只需修改评估函数返回字典，再展开到结果DataFrame中：

def eval_multi_func(true: pl.Series, pred: pl.Series) -> dict:
    mae = (true - pred).abs().mean()
    rmse = ((true - pred)**2).mean()**0.5
    return {"mae": mae, "rmse": rmse}

result = df.group_by("seg").map_groups(
    lambda group: pl.DataFrame({
        "seg": [group["seg"][0]],
        **eval_multi_func(group["true"], group["pred"])
    })
)

与Pandas代码的对应关系 #

• Pandas中groupby.apply接收每个分组的DataFrame，Polars的map_groups逻辑完全一致：每个分组的DataFrame会传入lambda函数。
• 如果你更习惯类似Pandas的apply写法，Polars也支持group_by(...).apply(...)，但map_groups在大型数据集上性能更优，且返回格式更可控。

注意事项 #

• 确保eval_func接收的是Polars Series而非Pandas Series（如果你的函数原本是为Pandas写的，只需把pd.Series换成pl.Series，大部分数值计算API是兼容的）。
• Polars默认开启并行计算，分组计算的速度会显著快于Pandas，尤其是数据量较大时。

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