如何在Polars DataFrame中计算多组统计量并整合结果？

现有一个包含Cost和StartTime两列的Polars DataFrame，数据及结构如下：

import polars as pl
from datetime import datetime

df = pl.DataFrame({
    "Cost": [1, 2, 3],
    "StartTime": [datetime(2023, 6, 23), datetime(2024, 6, 22), datetime(2024, 6, 21)],
})

DataFrame结构：

shape: (3, 2)
┌──────┬─────────────────────┐
│ Cost ┆ StartTime           │
│ ---  ┆ ---                 │
│ i64  ┆ datetime[μs]        │
╞══════╪═════════════════════╡
│ 1    ┆ 2023-06-23 00:00:00 │
│ 2    ┆ 2024-06-22 00:00:00 │
│ 3    ┆ 2024-06-21 00:00:00 │
└──────┴─────────────────────┘

需要计算以下统计量并将结果存入一个DataFrame中：

• 整个DataFrame的总成本和平均成本
• 仅2024年数据的总成本和平均成本

以下是分开计算的示例代码：

total_costs = df.select('Cost').sum()
average_costs = df.select('Cost').mean()
df_2024 = df.filter(pl.col('StartTime').dt.year() == 2024)
total_costs_2024 = df_2024.select('Cost').sum()
average_costs_2024 = df_2024.select('Cost').mean()

方法一：一次性聚合计算（推荐） #

直接在select中使用条件过滤聚合，无需创建中间DataFrame，代码简洁且性能更优：

result = df.select(
    # 全局统计
    pl.col("Cost").sum().alias("total_cost_all"),
    pl.col("Cost").mean().alias("avg_cost_all"),
    # 2024年数据统计
    pl.col("Cost").filter(pl.col("StartTime").dt.year() == 2024).sum().alias("total_cost_2024"),
    pl.col("Cost").filter(pl.col("StartTime").dt.year() == 2024).mean().alias("avg_cost_2024")
)

print(result)

输出结果：

shape: (1, 4)
┌──────────────┬─────────────┬────────────────┬───────────────┐
│ total_cost_all ┆ avg_cost_all ┆ total_cost_2024 ┆ avg_cost_2024 │
│ ---          ┆ ---         ┆ ---            ┆ ---           │
│ i64          ┆ f64         ┆ i64            ┆ f64           │
╞══════════════╪═════════════╪════════════════╪═══════════════╡
│ 6            ┆ 2.0         ┆ 5              ┆ 2.5           │
└──────────────┴─────────────┴────────────────┴───────────────┘

方法二：分组透视法 #

通过添加分组标签，分组聚合后再透视成宽表结构：

result = (
    df
    # 添加分组列，标记数据归属
    .with_columns(
        pl.when(pl.col("StartTime").dt.year() == 2024)
        .then("2024")
        .otherwise("all")
        .alias("group")
    )
    # 按分组计算统计量
    .group_by("group")
    .agg(
        pl.col("Cost").sum().alias("total_cost"),
        pl.col("Cost").mean().alias("avg_cost")
    )
    # 透视成单行多列的结果表
    .pivot(index=None, columns="group", values=["total_cost", "avg_cost"])
)

print(result)

输出结果与方法一完全一致。

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