PySpark中日期列与节假日字典匹配失败问题求助

问题根源 #

你当前代码的核心错误是使用字符串索引map，而非列表达式：mapping_expr["only_date"]是尝试查找key为字符串"only_date"的条目，而非用only_date列的值作为key去匹配map中的节假日日期。

另外，虽然holidays返回的datetime.date对象通过lit()可以自动转为Spark的DateType，和你的only_date列类型一致，但索引方式错误导致完全无法匹配。

修正方案 #

方案1：修正map索引方式（推荐） #

直接使用col("only_date")作为map的索引，同时可以直接生成is_holiday标记列，无需先查名称再过滤：

from pyspark.sql.functions import col, create_map, lit, isnull
from itertools import chain
import holidays

# 获取节假日字典
nyse_holidays = holidays.financial.ny_stock_exchange.NewYorkStockExchange(years=2018)

# 构造map表达式：key为Spark DateType，value为节假日名称
mapping_expr = create_map(*[lit(x) for x in chain(*nyse_holidays.items())])

# 添加节假日名称列+是否为节假日的标记列
df_result = df.withColumn("holiday_name", mapping_expr[col("only_date")]) \
              .withColumn("is_holiday", ~isnull(col("holiday_name")))

# 查看匹配结果
display(df_result.filter(col("is_holiday")))

方案2：用数组包含判断（更简洁，无需构造map） #

如果只需要标记是否为节假日，不需要节假日名称，可以直接将节假日日期转为Spark DateType的数组，用array_contains判断：

from pyspark.sql.functions import col, array_contains, lit, array
import holidays

nyse_holidays = holidays.financial.ny_stock_exchange.NewYorkStockExchange(years=2018)
# 将节假日日期转为Spark DateType的数组
holiday_dates = [lit(date) for date in nyse_holidays.keys()]

df_result = df.withColumn("is_holiday", array_contains(array(*holiday_dates), col("only_date")))

display(df_result.filter(col("is_holiday")))

额外说明 #

如果你想直接使用date（TimestampType）列，无需单独创建only_date，可以用date_trunc("day", col("date")).cast("date")将Timestamp转为DateType后再进行匹配：

from pyspark.sql.functions import date_trunc
# 方案1中替换为
mapping_expr[date_trunc("day", col("date")).cast("date")]
# 方案2中替换为
array_contains(array(*holiday_dates), date_trunc("day", col("date")).cast("date"))

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