如何用Polars内置函数实现季度字符串的加减及差值计算？

核心思路 #

先将季度字符串转换为Polars标准日期类型，利用Polars内置的offset_by()实现高效的季度加减操作，再通过日期的年/季度差值计算实现行与行的季度差，全程采用Polars向量化操作，充分发挥并行性能优势。

完整代码示例 #

import polars as pl

# 示例DataFrame
df = pl.DataFrame({
    "quarter": ["2020Q4", "2021Q1", "2021Q3", "1999Q1"]
})

# 1. 转换季度字符串为季度起始日期（Polars日期类型）
df = df.with_columns(
    quarter_date=pl.date(
        pl.col("quarter").str.extract(r"(\d{4})", 1).cast(pl.Int32),
        pl.col("quarter").str.extract(r"Q(\d)", 1).cast(pl.Int32)*3 - 2,
        1
    )
)

# 2. 实现任意季度加减（支持正负整数）
df = df.with_columns(
    # 加1季度
    quarter_plus_1=pl.col("quarter_date").offset_by("1q").dt.strftime("%YQ%q"),
    # 减2季度
    quarter_minus_2=pl.col("quarter_date").offset_by("-2q").dt.strftime("%YQ%q")
)

# 3. 计算当前行与上一行的季度差值
df = df.with_columns(
    quarter_diff=(
        (pl.col("quarter_date").dt.year() - pl.col("quarter_date").shift(1).dt.year()) * 4
        + (pl.col("quarter_date").dt.quarter() - pl.col("quarter_date").shift(1).dt.quarter())
    )
)

# 查看最终结果
print(df)

关键细节解释 #

• 季度转日期：通过正则提取年份和季度数字，构造对应季度的起始日期（如2020Q4对应2020-10-01），这是使用Polars日期函数的基础。
• 季度加减：offset_by()支持"nq"格式参数，n为正整数时执行加法，负整数时执行减法，自动处理跨年度的进位/借位（如2020Q4+1q自动转为2021Q1），且全程为向量化并行操作，性能远优于Python循环。
• 差值计算：通过年份差乘以4加上季度差，得到精确的季度间隔数，避免了按天数近似计算的误差；shift(1)获取上一行的日期数据，配合向量化计算实现高效的行间差值统计。

输出结果 #

shape: (4, 5)
┌──────────┬──────────────┬──────────────┬───────────────┬─────────────┐
│ quarter  ┆ quarter_date ┆ quarter_plus ┆ quarter_minus ┆ quarter_diff│
│ ---      ┆ ---          ┆ _1           ┆ _2            ┆ ---         │
│ str      ┆ date         ┆ ---          ┆ ---           ┆ i32         │
│          ┆              ┆ str          ┆ str           ┆             │
╞══════════╪══════════════╪══════════════╪═══════════════╪═════════════╡
│ 2020Q4   ┆ 2020-10-01   ┆ 2021Q1       ┆ 2020Q2        ┆ null        │
│ 2021Q1   ┆ 2021-01-01   ┆ 2021Q2       ┆ 2020Q3        ┆ 1           │
│ 2021Q3   ┆ 2021-07-01   ┆ 2021Q4       ┆ 2020Q4        ┆ 2           │
│ 1999Q1   ┆ 1999-01-01   ┆ 1999Q2       ┆ 1998Q3        ┆ -90         │
└──────────┴──────────────┴──────────────┴───────────────┴─────────────┘

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