如何基于日期列的年份对Polars LazyFrame进行透视？

输入数据 #

data = {
    'text': ['text A', 'text B', 'text A', 'text A', 'text B', 'text C'],
    'length': [7, 1, 4, 0, 2, 0],
    'date': [20210101, 20231003, 20220808, 20210616, 20210203, 20220419]
}

ldf = pl.LazyFrame(data)

错误尝试与问题 #

原代码手动指定年份且运行报错：

query = (
    ldf
        .group_by(pl.col('date').cast(pl.String).str.slice(0, 4).alias('year'), pl.col('text')).agg(
            pl.col('length').sum().filter(pl.col('year') == '2021').alias('2021'),
            pl.col('length').sum().filter(pl.col('year') == '2022').alias('2022'),
            pl.col('length').sum().filter(pl.col('year') == '2023').alias('2023'),
        )
)

报错信息：

ComputeError: column 'year' not available in schema Schema:
name: text, data type: Utf8
name: length, data type: Int64
name: date, data type: Int64

错误原因：year是group_by中临时生成的别名，并非原始数据列，在agg的表达式上下文里无法直接引用；同时分组逻辑错误，每组的year是固定值，过滤操作无意义。

期望输出 #

grp_texts  2021  2022  2023
    text A     7     4     0
    text B     2     0     1
    text C     0     0     0

正确实现方法 #

方法一：动态生成年份列（推荐，无需手动指定年份） #

通过提取年份、分组聚合、透视三步完成，自动识别所有存在的年份：

result = (
    ldf
    # 从整数日期中提取年份，转字符串用于后续列名
    .with_columns(year=pl.col("date") // 10000)
    # 按文本和年份分组，计算长度总和
    .group_by("text", "year")
    .agg(pl.col("length").sum().alias("total_length"))
    # 透视：将年份转为列，缺失值填充0
    .pivot(
        index="text",
        columns="year",
        values="total_length",
        fill_value=0
    )
    # 重命名列匹配期望输出
    .rename({"text": "grp_texts"})
)

print(result.collect())

方法二：手动指定年份（适用于固定年份场景） #

先将年份转为数据列，再按文本分组，针对每个年份过滤求和：

result = (
    ldf
    .with_columns(year=pl.col("date") // 10000)
    .group_by("text")
    .agg(
        pl.col("length").filter(pl.col("year") == 2021).sum().alias("2021"),
        pl.col("length").filter(pl.col("year") == 2022).sum().alias("2022"),
        pl.col("length").filter(pl.col("year") == 2023).sum().alias("2023"),
    )
    .rename({"text": "grp_texts"})
)

print(result.collect())

说明 #

• 整数除法//10000提取年份比字符串切片更高效，避免类型转换开销；
• 方法一的透视操作会自动处理所有存在的年份，无需手动维护年份列表，扩展性更强；
• 两种方法最终输出均匹配期望结果。

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