如何用Polars的select/filter/group_by/value_counts批量计算多类统计量

问题背景 #

我有一个包含Cost、State和StartTime三列的Polars DataFrame，数据定义如下：

import polars as pl
from datetime import datetime

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

DataFrame的结构与数据：

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

需要完成以下三个统计需求：

• 计算整个DataFrame以及当月数据的Cost均值与总成本，结果存入单个DataFrame；
• 计算每日数据频率（即统计相同日期的行数）；
• 统计State为"OK"和"KO"的占比。

如何在Polars中尽可能避免单独计算每个值来实现上述需求？

我的尝试代码：

this_year = datetime.now().year
this_month = datetime.now().month
filter_this_year_month_exp = (pl.col('StartTime').dt.year() == this_year) & (pl.col('StartTime').dt.month() == this_month)
df = df.with_columns(
    pl.col('StartTime').dt.date().alias('StartDate')
)
stats = df.select(
    pl.col("Cost").sum().alias("TotalCost"),
    pl.col("Cost").mean().alias("AverageCost"),
    pl.col("Cost").filter(filter_this_year_month_exp).sum().alias("TotalCostCurrentMonth"),
    pl.col("Cost").filter(filter_this_year_month_exp).mean().alias("AverageCostCurrentMonth"),
)
frequency_per_day_in_this_month = df.filter(
    filter_this_year_month_exp
).group_by('StartDate').len().select("len").mean().item()
# this returns just the counts, should be divided by len(df) to compute the percentage
df.select(
    pl.col("State").value_counts()
).unnest("State")

优化实现方案 #

通过Polars的表达式组合与聚合操作，可减少重复计算，让代码更简洁高效：

1. 统一计算全局与当月Cost统计 #

提前定义当月过滤条件，在单个select中完成所有聚合，避免重复遍历数据：

from datetime import datetime

# 复用当月过滤条件
current_month_filter = (
    pl.col("StartTime").dt.year() == datetime.now().year()
) & (pl.col("StartTime").dt.month() == datetime.now().month())

# 一次性生成所有Cost统计指标
cost_stats = df.select(
    # 全局统计
    pl.col("Cost").sum().alias("TotalCost"),
    pl.col("Cost").mean().alias("AverageCost"),
    # 当月统计
    pl.col("Cost").filter(current_month_filter).sum().alias("TotalCostCurrentMonth"),
    pl.col("Cost").filter(current_month_filter).mean().alias("AverageCostCurrentMonth")
)

2. 每日数据频率统计 #

按日期分组聚合即可得到每日行数，按需添加过滤条件：

# 所有日期的每日数据频率
daily_frequency = df.with_columns(pl.col("StartTime").dt.date().alias("StartDate"))\
                    .group_by("StartDate")\
                    .agg(pl.count().alias("DailyCount"))

# 当月的每日数据频率
daily_frequency_current_month = df.filter(current_month_filter)\
                                  .with_columns(pl.col("StartTime").dt.date().alias("StartDate"))\
                                  .group_by("StartDate")\
                                  .agg(pl.count().alias("DailyCount"))

# 当月每日平均频率
avg_daily_freq_current_month = daily_frequency_current_month.select(
    pl.col("DailyCount").mean().alias("AvgDailyCountCurrentMonth")
)

3. State占比统计 #

利用value_counts()结合内置聚合计算占比，无需手动除法：

state_ratio = df.select(pl.col("State").value_counts())\
               .unnest("State")\
               .with_columns(
                   (pl.col("count") / pl.col("count").sum()).alias("Ratio")
               )

整合所有统计（可选） #

若需将所有统计结果合并为一个DataFrame，可通过横向拼接实现：

# 将各类统计转为单行格式后合并
state_ratio_pivot = state_ratio.pivot(index=None, columns="State", values="Ratio").fill_null(0)
final_stats = pl.concat([cost_stats, avg_daily_freq_current_month, state_ratio_pivot], how="horizontal")

关键优化点 #

• 复用过滤条件：避免重复编写相同逻辑，提升代码可维护性；
• 批量聚合：在单个操作中完成多个统计指标，减少数据遍历次数；
• 依赖内置函数：利用Polars原生方法实现统计逻辑，避免手动计算；
• 保留DataFrame格式：尽量不提取单个标量值，方便后续数据处理。

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