PySpark：日期列与字典匹配问题解决方案-Python教程-PHP中文网

PySpark：日期列与字典匹配问题解决方案

心靈之曲

发布： 2025-10-17 16:51:01

原创

675人浏览过

pyspark：日期列与字典匹配问题解决方案

本文旨在解决在PySpark中将日期列与字典进行匹配时遇到的问题。通过`create_map`函数创建映射表达式，并结合`withColumn`和`filter`函数，实现高效的日期匹配。本文提供详细的代码示例和解释，帮助读者理解并解决类似问题，并提供了一些注意事项。

在PySpark中，经常需要根据日期进行数据处理，例如判断某一天是否为节假日。一种常见的场景是将DataFrame中的日期列与包含节假日信息的字典进行匹配，从而添加新的列来标识是否为节假日。本文将介绍如何使用create_map函数实现这一功能，并解决可能遇到的问题。

问题描述

假设我们有一个包含日期信息的Spark DataFrame，以及一个包含节假日信息的Python字典。我们的目标是创建一个新的DataFrame列，该列指示DataFrame中的日期是否在节假日字典中。

DataFrame的Schema如下：

root
 |-- id: long (nullable = false)
 |-- date: timestamp (nullable = false)
 |-- year: integer (nullable = false)
 |-- month: integer (nullable = false)
 |-- day: string (nullable = false)
 |-- day_of_year: string (nullable = false)
 |-- hour: string (nullable = false)
 |-- minute: string (nullable = false)
 |-- is_weekend: boolean (nullable = false)
 |-- only_date: date (nullable = false)

登录后复制

节假日字典（例如，从holidays包获取）如下：

AI建筑知识问答

用人工智能ChatGPT帮你解答所有建筑问题

查看详情

{datetime.date(2018, 12, 5): 'Day of Mourning for President George H.W. Bush', datetime.date(2018, 1, 1): "New Year's Day", datetime.date(2018, 1, 15): 'Martin Luther King Jr. Day', datetime.date(2018, 2, 19): "Washington's Birthday", datetime.date(2018, 3, 30): 'Good Friday', datetime.date(2018, 5, 28): 'Memorial Day', datetime.date(2018, 7, 4): 'Independence Day', datetime.date(2018, 9, 3): 'Labor Day', datetime.date(2018, 11, 22): 'Thanksgiving Day', datetime.date(2018, 12, 25): 'Christmas Day'}

登录后复制

解决方案

关键在于正确地将DataFrame中的日期列传递给create_map函数生成的映射表达式。在create_map中，我们需要使用col("only_date")来引用DataFrame中的only_date列。

以下是完整的代码示例：

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

# 创建 SparkSession
spark = SparkSession.builder.appName("DateMatching").getOrCreate()

# 创建示例 DataFrame (为了示例，这里手动创建)
data = [(1, datetime(2018, 1, 1)), (2, datetime(2018, 1, 15)), (3, datetime(2018, 1, 20))]
df = spark.createDataFrame(data, ["id", "date"]).withColumn("only_date", col("date").cast("date"))

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

# 创建映射表达式
mapping_expr = create_map([lit(x) for x in chain(*nyse_holidays.items())])

# 添加新列，指示是否为节假日
df = df.withColumn("is_holiday", mapping_expr[col("only_date")])

# 显示结果
df.show()

# 停止 SparkSession
spark.stop()

登录后复制

代码解释：

导入必要的库： 导入pyspark.sql.functions中的col, create_map, lit，以及itertools中的chain。
创建节假日字典： 使用holidays包创建包含2018年纽约证券交易所节假日的字典。
创建映射表达式： 使用create_map函数将节假日字典转换为PySpark可以使用的映射表达式。chain(*nyse_holidays.items())将字典的键值对展开为扁平的列表，lit(x)将每个键值对转换为字面量。
添加新列： 使用withColumn函数添加名为is_holiday的新列。mapping_expr[col("only_date")]表示根据only_date列的值在映射表达式中查找对应的值。如果only_date列的值在节假日字典中存在，则is_holiday列的值为对应的节假日名称；否则，为null。
显示结果： 使用show函数显示包含新列的DataFrame。