如何用Polars高效标记数据列中的节假日？

我需要在Polars DataFrame中添加一列，标记日期是否为智利节假日。目前找到的代码用了map_elements+lambda，但想改用更高效的Polars原生方法。尝试优化后，在Polars 0.20.31版本中所有日期的节假日标记都返回false，比如2024-06-20实际是节假日却被识别为false。

原低效代码及预期输出 #

以下是原本能用但效率较低的实现代码：

import polars as pl
import holidays

# 初始化智利节假日
cl_holidays = holidays.CL()

# 示例数据
data = {
    "Date": ["2024-06-20 00:00:00", "2024-06-21 00:00:00", "2024-06-22 00:00:00", "2024-06-23 00:00:00", "2024-06-24 00:00:00"],
    "Amount": [100, 200, 300, 400, 500],
    "User_Count" : [1, 2, 3, 4, 5]
}

# 创建DataFrame
df = pl.DataFrame(data)

# 添加'Is_Holiday'列
df = df.with_columns(
    (pl.col("Date").map_elements(lambda x: x.split(" ")[0] in cl_holidays, return_dtype=pl.Boolean)).alias("Is_Holiday")
).with_columns(pl.col("Date").str.strptime(pl.Datetime))

df

预期输出：

shape: (5, 4)
┌─────────────────────┬────────┬────────────┬────────────┐
│ Date                ┆ Amount ┆ User_Count ┆ Is_Holiday │
│ ---                 ┆ ---    ┆ ---        ┆ ---        │
│ datetime[μs]        ┆ i64    ┆ i64        ┆ bool       │
╞═════════════════════╪════════╪════════════╪════════════╡
│ 2024-06-20 00:00:00 ┆ 100    ┆ 1          ┆ true       │
│ 2024-06-21 00:00:00 ┆ 200    ┆ 2          ┆ false      │
│ 2024-06-22 00:00:00 ┆ 300    ┆ 3          ┆ false      │
│ 2024-06-23 00:00:00 ┆ 400    ┆ 4          ┆ false      │
│ 2024-06-24 00:00:00 ┆ 500    ┆ 5          ┆ false      │
└─────────────────────┴────────┴────────────┴────────────┘

尝试的优化方案及问题 #

我尝试了以下Polars原生方法，但所有Is_Holiday都返回false：

import polars as pl
import holidays

# 初始化智利节假日
cl_holidays = holidays.CL()

# 示例数据
data = {
    "Date": ["2024-06-20 00:00:00", "2024-06-21 00:00:00", "2024-06-22 00:00:00", "2024-06-23 00:00:00", "2024-06-24 00:00:00"],
    "Amount": [100, 200, 300, 400, 500],
    "User_Count" : [1, 2, 3, 4, 5]
}

# 创建DataFrame
df = pl.DataFrame(data)

# 添加'Is_Holiday'列
df.with_columns(
    Is_Holiday=pl.col('Date').str.to_datetime().dt.date().is_in(cl_holidays.keys())
)

实际错误输出：

shape: (5, 4)
┌─────────────────────┬────────┬────────────┬────────────┐
│ Date                ┆ Amount ┆ User_Count ┆ Is_Holiday │
│ ---                 ┆ ---    ┆ ---        ┆ ---        │
│ str                 ┆ i64    ┆ i64        ┆ bool       │
╞═════════════════════╪════════╪════════════╪════════════╡
│ 2024-06-20 00:00:00 ┆ 100    ┆ 1          ┆ false      │
│ 2024-06-21 00:00:00 ┆ 200    ┆ 2          ┆ false      │
│ 2024-06-22 00:00:00 ┆ 300    ┆ 3          ┆ false      │
│ 2024-06-23 00:00:00 ┆ 400    ┆ 4          ┆ false      │
│ 2024-06-24 00:00:00 ┆ 500    ┆ 5          ┆ false      │
└─────────────────────┴────────┴────────────┴────────────┘

解决方案 #

问题出在类型不匹配：cl_holidays.keys()返回的是Python原生datetime.date对象，而Polars的dt.date()返回的是Polars专属的Date类型，两者无法直接匹配。

解决方法是把节假日日期转换成Polars的Date类型，再进行匹配，这样既能保留Polars原生操作的高效性，又能正确识别节假日：

import polars as pl
import holidays

# 初始化智利节假日
cl_holidays = holidays.CL()
# 提取节假日日期并转为Polars Date类型
holiday_dates = pl.Series([pl.Date(d) for d in cl_holidays.keys()])

# 示例数据
data = {
    "Date": ["2024-06-20 00:00:00", "2024-06-21 00:00:00", "2024-06-22 00:00:00", "2024-06-23 00:00:00", "2024-06-24 00:00:00"],
    "Amount": [100, 200, 300, 400, 500],
    "User_Count": [1, 2, 3, 4, 5]
}

df = pl.DataFrame(data)

# 高效添加节假日标记列
df = df.with_columns(
    # 转换日期格式后判断是否在节假日列表中
    Is_Holiday=pl.col("Date").str.to_datetime().dt.date().is_in(holiday_dates)
).with_columns(
    # 把原Date列转为Datetime类型（按需保留）
    Date=pl.col("Date").str.to_datetime()
)

print(df)

运行后就能得到符合预期的输出，且完全使用Polars原生操作，效率远高于map_elements方法。

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