按unique_id与cutoff聚合计算指标后，如何合并Polars DataFrame？

我有如下结构的Polars DataFrame，需要基于unique_id和cutoff列做分组聚合，计算真实值列y与其他模型列（排除分组列及ds列）之间的评估指标（比如MSE），最终生成包含各模型指标均值及最优模型的汇总表。

输入DataFrame示例 #

shape: (5, 10)
┌───────────┬─────┬────────┬───────┬───┬──────────────┬──────────────┬──────────────┬──────────────┐
│ unique_id ┆ ds  ┆ cutoff ┆ y     ┆ … ┆ CrostonClass ┆ SeasonalNaiv ┆ HistoricAver ┆ DynamicOptim │
│ ---       ┆ --- ┆ ---    ┆ ---   ┆   ┆ ic           ┆ e            ┆ age          ┆ izedTheta    │
│ str       ┆ i64 ┆ i64    ┆ f32   ┆   ┆ ---          ┆ ---          ┆ ---          ┆ ---          │
│           ┆     ┆        ┆       ┆   ┆ f32          ┆ f32          ┆ f32          ┆ f32          │
╞═══════════╪═════╪════════╪═══════╪═══╪══════════════╪══════════════╪══════════════╪══════════════╡
│ H1        ┆ 701 ┆ 700    ┆ 619.0 ┆ … ┆ 742.668762   ┆ 691.0        ┆ 661.674988   ┆ 612.767517   │
│ H1        ┆ 702 ┆ 700    ┆ 565.0 ┆ … ┆ 742.668762   ┆ 618.0        ┆ 661.674988   ┆ 536.846252   │
│ H1        ┆ 703 ┆ 700    ┆ 532.0 ┆ … ┆ 742.668762   ┆ 563.0        ┆ 661.674988   ┆ 497.82428    │
│ H1        ┆ 704 ┆ 700    ┆ 495.0 ┆ … ┆ 742.668762   ┆ 529.0        ┆ 661.674988   ┆ 464.723236   │
│ H1        ┆ 705 ┆ 700    ┆ 481.0 ┆ … ┆ 742.668762   ┆ 504.0        ┆ 661.674988   ┆ 440.972351   │
└───────────┴─────┴────────┴───────┴───┴──────────────┴──────────────┴──────────────┴──────────────┘

当前实现代码 #

我尝试通过循环逐个计算模型指标，但不知道如何正确合并结果得到预期汇总表：

from datasetsforecast.losses import mse, mae, rmse


def evaluate_cross_validation(df, metric):
    models = df.drop(columns=['ds', 'cutoff', 'y', 'unique_id']).columns
    evals = []

    for model in models:
      eval_ = (
          df
          .groupby(['unique_id', 'cutoff'])
          .agg(
              pl.apply(
                  exprs=['y', model],
                  function=lambda args: metric(args[0], args[1]),
              )
          )
          .rename({'y': model})
          .sort(by=['unique_id', 'cutoff'])
      )
      evals.append(eval_)
    
    uid_cutoff = evals[0].select(['unique_id'])
    eval_dfs = pl.concat([df.drop(['unique_id', 'cutoff']) for df in evals], how='horizontal')
    evals = pl.concat([uid_cutoff, eval_dfs], how='horizontal')
    evals = evals.groupby(['unique_id']).mean() # 对每个unique_id和模型的所有cutoff指标求平均
    best_model = [min(row, key=row.get) for row in evals.drop('unique_id').rows(named=True)]
    evals = evals.with_columns(pl.lit(best_model).alias('best_model')).sort(by=['unique_id'])
    return evals

预期输出 #

shape: (5, 8)
┌───────────┬───────────┬───────────┬────────────┬────────────┬────────────┬────────────┬──────────┐
│ unique_id ┆ AutoARIMA ┆ HoltWinte ┆ CrostonCla ┆ SeasonalNa ┆ HistoricAv ┆ DynamicOpt ┆ best_mod │
│ ---       ┆ ---       ┆ rs        ┆ ssic       ┆ ive        ┆ erage      ┆ imizedThet ┆ el       │
│ str       ┆ f64       ┆ ---       ┆ ---        ┆ ---        ┆ ---        ┆ a          ┆ ---      │
│           ┆           ┆ f64       ┆ f64        ┆ f64        ┆ f64        ┆ ---        ┆ str      │
│           ┆           ┆           ┆            ┆            ┆            ┆ f64        ┆          │
╞═══════════╪═══════════╪═══════════╪════════════╪════════════╪════════════╪════════════╪══════════╡
│ H1        ┆ 1979.3021 ┆ 44888.019 ┆ 28038.7363 ┆ 1422.66668 ┆ 20927.6640 ┆ 1296.33398 ┆ DynamicO │
│           ┆ 85        ┆ 531       ┆ 28         ┆ 7          ┆ 62         ┆ 4          ┆ ptimized │
│           ┆           ┆           ┆            ┆            ┆            ┆            ┆ Theta    │
│ H10       ┆ 458.89271 ┆ 2812.9166 ┆ 1483.48413 ┆ 96.895832  ┆ 1980.36749 ┆ 379.621124 ┆ Seasonal │
│           ┆ 5         ┆ 26        ┆ 1          ┆            ┆ 3          ┆            ┆ Naive    │
│ H100      ┆ 8629.9482 ┆ 121625.37 ┆ 91945.1406 ┆ 12019.0    ┆ 78491.1914 ┆ 21699.6479 ┆ AutoARIM │
│           ┆ 42        ┆ 5         ┆ 25         ┆            ┆ 06         ┆ 49         ┆ A        │
│ H101      ┆ 6818.3486 ┆ 28453.395 ┆ 16183.6347 ┆ 10944.4580 ┆ 18208.4042 ┆ 63698.0732 ┆ AutoARIM │
│           ┆ 33        ┆ 508       ┆ 66         ┆ 08         ┆ 97         ┆ 42         ┆ A        │
│ H102      ┆ 65489.965 ┆ 232924.85 ┆ 132655.300 ┆ 12699.8959 ┆ 309110.468 ┆ 31393.5214 ┆ Seasonal │
│           ┆ 82        ┆ 1562      ┆ 781        ┆ 96         ┆ 75         ┆ 84         ┆ Naive    │

解决方案：优化后的Polars实现 #

原来的循环方式不仅效率低，还容易在合并DataFrame时出错。可以利用Polars的批量表达式处理，一次性计算所有模型的指标，避免循环和手动合并：

from datasetsforecast.losses import mse, mae, rmse
import polars as pl


def evaluate_cross_validation(df, metric):
    # 筛选出所有模型列
    model_cols = [col for col in df.columns if col not in ['ds', 'cutoff', 'y', 'unique_id']]
    
    # 1. 按unique_id和cutoff分组，计算每个模型的指标
    grouped_metrics = (
        df
        .groupby(['unique_id', 'cutoff'])
        .agg([
            pl.apply([pl.col('y'), pl.col(model)], lambda args: metric(args[0], args[1])).alias(model)
            for model in model_cols
        ])
    )
    
    # 2. 按unique_id分组，对每个模型的指标求跨cutoff的平均值
    avg_metrics = grouped_metrics.groupby('unique_id').mean()
    
    # 3. 添加best_model列：取每个unique_id下指标最小的模型名称
    avg_metrics = avg_metrics.with_columns(
        pl.struct(model_cols).apply(lambda s: min(s, key=s.get)).alias('best_model')
    )
    
    return avg_metrics.sort('unique_id')

代码说明 #

1. 批量计算指标：通过列表推导式生成所有模型列的计算表达式，一次性在agg中完成所有模型的指标计算，无需循环生成单个DataFrame。
2. 简化合并逻辑：直接在分组聚合后得到包含所有模型指标的DataFrame，避免手动拼接多个小DataFrame的繁琐操作。
3. 高效计算最优模型：利用pl.struct将模型列打包成结构体，通过apply直接找出每个行中指标最小的模型名称，逻辑更清晰。

这个实现既符合Polars的高效编程风格，也能准确生成你需要的汇总表。

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