如何在Polars中实现分组后的Value Counts并生成指定格式输出？

需求说明 #

实现按Country分组后，统计Date列的年月格式（%Y-%m）数值分布，最终转置成以年月为行、国家为列的结构，并按年月排序，缺失值显示为null。

原始代码 #

Pandas可行实现

import polars as pl
import pandas as pd

pdf = pl.DataFrame({
    "Country": ["US", "UK", "US"],
    "Date": pd.to_datetime(["2023-01-01", "2023-02-01", "2023-03-01"])
})

df = pdf.to_pandas()

# Pandas逻辑：分组统计年月频次，转置排序
out = df.groupby('Country').apply(lambda row: row['Date'].dt.strftime('%Y-%m').value_counts()).unstack(0).sort_index()
print(out)

未达预期的Polars尝试

# 生成年月列
pdf = pdf.with_columns(
    Date_ym = pdf['Date'].dt.strftime("%Y-%m")
)

# 尝试分组统计，但未得到目标结构
out = (pdf.group_by("Country")
 .agg(pl.col('Date_ym')
      .value_counts())
 .explode('Date_ym')
)

print(out)

正确的Polars实现 #

import polars as pl

# 用Polars原生方法初始化数据，避免依赖Pandas
pdf = pl.DataFrame({
    "Country": ["US", "UK", "US"],
    "Date": pl.date_range(start="2023-01-01", end="2023-03-01", interval="1mo")
})

# 核心逻辑：生成年月列 → 分组统计频次 → 转置结构 → 排序结果
out = (pdf
       .with_columns(Date_ym = pl.col("Date").dt.strftime("%Y-%m"))
       .group_by(["Country", "Date_ym"])
       .agg(count=pl.count())
       .pivot(index="Date_ym", columns="Country", values="count")
       .sort("Date_ym")
       .rename({"Date_ym": "Date"})
)

print(out)

预期输出 #

shape: (3, 3)
┌─────────┬──────┬──────┐
│ Date    ┆ UK   ┆ US   │
│ ---     ┆ ---  ┆ ---  │
│ str     ┆ u32  ┆ u32  │
╞═════════╪══════╪══════╡
│ 2023-01 ┆ null ┆ 1    │
│ 2023-02 ┆ 1    ┆ null │
│ 2023-03 ┆ null ┆ 1    │
└─────────┴──────┴──────┘

关键步骤说明 #

• 原生日期处理：用Polars自带的date_range和dt.strftime生成年月列，无需依赖Pandas
• 联合分组统计：同时按Country和Date_ym分组，用pl.count()统计每组出现次数
• 转置结构：通过pivot方法直接将国家转为列、年月转为行，自动填充缺失值为null
• 排序重命名：按年月排序后，将Date_ym列名改为Date匹配目标输出

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