寻求Python & Polars中更高效的年度内能耗月度占比计算方法

我有一份包含多年月度数据的能耗CSV文件，需要计算每年中每个月的能耗占该年度总能耗的百分比（小数形式），例如2025年8月能耗占全年总能耗的12.3%。

目前我通过分组求和生成年度总能耗DataFrame，再用左连接关联原数据的方式实现需求，但这种写法不够贴合Polars的风格，希望得到更简洁、可读性更强的实现方案。

现有实现代码 #

df_util_compare = (
    df_util
        .select(
        [
            'Service End Date',
            'Metered Consumption',
            'Number of Days',
            'Cons Month',
            'Cons Year',
        ])  # select the needed columns
        # .sort("Service End Date")
)
with pl.Config(tbl_cols=15, tbl_rows=6, tbl_width_chars=180, tbl_formatting="UTF8_FULL_CONDENSED", fmt_str_lengths=120, tbl_hide_dataframe_shape=False):
    print(df_util_compare.__repr__())

df_yearly_total = (
    df_util_compare
        .select(['Metered Consumption', 'Cons Year'])
        .group_by("Cons Year", maintain_order=True).sum()
        .rename({"Metered Consumption": "Yearly Consumption"})
)
print(df_yearly_total)

df_util_compare = (
    df_util_compare
        .join(df_yearly_total, on='Cons Year', how="left")
        .with_columns((pl.col('Metered Consumption')/pl.col('Yearly Consumption')).alias('Monthly KW Portion'))
)
with pl.Config(tbl_cols=15, tbl_rows=6, tbl_width_chars=180, tbl_formatting="UTF8_FULL_CONDENSED", fmt_str_lengths=120, tbl_hide_dataframe_shape=False):
    print(df_util_compare)

优化后的Polars风格实现 #

利用Polars的**窗口函数（Window Functions）**可以直接在原数据集上计算分组聚合值，无需额外生成中间DataFrame再做关联，代码更简洁且符合Polars的向量化操作理念：

import polars as pl

df_util_compare = (
    df_util
    # 筛选所需列
    .select(
        'Service End Date',
        'Metered Consumption',
        'Number of Days',
        'Cons Month',
        'Cons Year'
    )
    # 窗口函数计算年度总能耗及月度占比
    .with_columns(
        # 按年度分组计算总能耗
        pl.col('Metered Consumption').sum().over('Cons Year').alias('Yearly Consumption'),
        # 直接计算月度能耗占年度的比例
        (pl.col('Metered Consumption') / pl.col('Metered Consumption').sum().over('Cons Year')).alias('Monthly KW Portion')
    )
)

# 配置并打印结果
with pl.Config(tbl_cols=15, tbl_rows=6, tbl_width_chars=180, tbl_formatting="UTF8_FULL_CONDENSED", fmt_str_lengths=120, tbl_hide_dataframe_shape=False):
    print(df_util_compare)

优化优势 #

• 代码更紧凑：无需拆分多步生成中间DataFrame，从数据筛选到结果计算一步完成
• 可读性更强：逻辑连贯，直接体现"按年度计算总能耗，再求月度占比"的业务逻辑
• 性能更优：避免了额外的join操作，完全利用Polars的向量化引擎处理数据

运行结果 #

shape: (48, 7)
┌──────────────────┬─────────────────────┬────────────────┬────────────┬───────────┬────────────────────┬────────────────────┐
│ Service End Date ┆ Metered Consumption ┆ Number of Days ┆ Cons Month ┆ Cons Year ┆ Yearly Consumption ┆ Monthly KW Portion │
│ ---              ┆ ---                 ┆ ---            ┆ ---        ┆ ---       ┆ ---                ┆ ---                │
│ date             ┆ i64                 ┆ i64            ┆ i64        ┆ i64       ┆ i64                ┆ f64                │
╞══════════════════╪═════════════════════╪════════════════╪════════════╪═══════════╪════════════════════╪════════════════════╡
│ 2025-09-03       ┆ 1387                ┆ 30             ┆ 9          ┆ 2025      ┆ 20296              ┆ 0.068339           │
│ 2025-08-04       ┆ 2496                ┆ 33             ┆ 8          ┆ 2025      ┆ 20296              ┆ 0.12298            │
│ 2025-07-02       ┆ 1684                ┆ 29             ┆ 7          ┆ 2025      ┆ 20296              ┆ 0.082972           │
│ …                ┆ …                   ┆ …              ┆ …          ┆ …         ┆ …                  ┆ …                  │
│ 2021-12-02       ┆ 2521                ┆ 30             ┆ 12         ┆ 2021      ┆ 3576               ┆ 0.704978           │
│ 2021-11-02       ┆ 777                 ┆ 29             ┆ 11         ┆ 2021      ┆ 3576               ┆ 0.217282           │
│ 2021-10-04       ┆ 278                 ┆ 13             ┆ 10         ┆ 2021      ┆ 3576               ┆ 0.07774            │
└──────────────────┴─────────────────────┴────────────────┴────────────┴───────────┴────────────────────┴────────────────────┘

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