Polars中计算两日期列间月末并展开：现有方法是否最优？

原代码是否可行？ #

你的代码是可行的，能够正确计算每一行valid_from到valid_to之间的所有月末日期，并通过explode展开结果。但需要注意：map_elements是逐行处理的操作，它无法利用Polars的向量化优化能力，当数据量较大时，执行效率会明显偏低。

更高效的无struct实现方法 #

Polars从0.18.0版本开始，pl.date_range支持直接传入列作为start和end参数（向量化输入），可以完全替代struct+map_elements的写法，代码更简洁且效率更高：

import polars as pl
from datetime import datetime

df = pl.DataFrame(
    {
        "id": ["A", "A", "A", "B", "B"],
        "value": ["1", "2", "3", "4", "5"],
        "valid_from": [
            datetime(2020, 1, 1),
            datetime(2021, 1, 1),
            datetime(2022, 1, 1),
            datetime(2020, 1, 1),
            datetime(2021, 1, 1),
        ],
        "valid_to": [
            datetime(2020, 12, 31),
            datetime(2021, 12, 31),
            datetime(2022, 12, 31),
            datetime(2020, 12, 31),
            datetime(2021, 12, 31),
        ],
    }
)

# 向量化实现，无需struct和map_elements
result = df.with_columns(
    pl.date_range(
        start=pl.col("valid_from"),
        end=pl.col("valid_to"),
        interval="1mo",
        eager=False
    ).dt.month_end().alias("test")
).explode("test")

print(result)

代码说明 #

1. pl.date_range直接接收valid_from和valid_to列作为起始/结束日期，为每行生成对应的月度日期序列；
2. dt.month_end()将序列中的每个日期转换为当月最后一天；
3. 最后用explode展开数组列，得到最终的行级结果。

这种写法完全依托Polars的向量化引擎，处理大数据集时的性能会比原代码提升数倍甚至数十倍。

兼容旧版Polars的方案 #

如果你使用的Polars版本低于0.18.0，可以通过计算月份偏移量的方式实现：

# 兼容旧版Polars的实现
result = df.with_columns(
    # 计算valid_from到valid_to的总月份数
    pl.when(pl.col("valid_to").dt.day() == 31)
      .then((pl.col("valid_to").dt.year() - pl.col("valid_from").dt.year()) * 12 + (pl.col("valid_to").dt.month() - pl.col("valid_from").dt.month()) + 1)
      .otherwise((pl.col("valid_to").dt.year() - pl.col("valid_from").dt.year()) * 12 + (pl.col("valid_to").dt.month() - pl.col("valid_from").dt.month()))
      .alias("total_months")
).with_columns(
    # 生成从0到total_months的偏移序列
    pl.int_ranges(0, pl.col("total_months")).alias("month_offsets")
).explode("month_offsets").with_columns(
    # 基于valid_from添加月份偏移，再取月末
    pl.col("valid_from").dt.offset_by(pl.col("month_offsets").cast(str) + "mo").dt.month_end().alias("test")
).drop("total_months", "month_offsets")

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