基于日期在PySpark中获取上月、下月及对应年份的方法

核心思路 #

利用PySpark内置的日期函数add_months直接对目标日期进行月份加减，再通过year和month函数提取对应的年份与月份，自动处理跨年度、月末日期等边界情况。

实现步骤 #

1. 导入依赖函数并初始化环境

from pyspark.sql import SparkSession
from pyspark.sql.functions import add_months, year, month, col, to_date

# 初始化SparkSession
spark = SparkSession.builder.appName("DateMonthCalc").getOrCreate()

2. 准备输入数据（示例）

假设你的输入数据包含字符串或日期类型的日期列，先确保列转换为标准date类型：

# 示例输入数据（模拟用户的输入结构）
data = [("2023-05-15",), ("2024-01-01",), ("2023-12-31",), ("2023-02-28",)]
df = spark.createDataFrame(data, ["current_date"])

# 若日期列是字符串类型，先转换为date类型
df = df.withColumn("current_date", to_date(col("current_date"), "yyyy-MM-dd"))

3. 计算上月、下月及对应年份

通过链式调用withColumn生成所需字段：

result_df = df.withColumn("prev_month", month(add_months(col("current_date"), -1))) \
              .withColumn("prev_year", year(add_months(col("current_date"), -1))) \
              .withColumn("next_month", month(add_months(col("current_date"), 1))) \
              .withColumn("next_year", year(add_months(col("current_date"), 1)))

4. 查看结果

result_df.show()

输出示例：

+------------+----------+---------+----------+---------+
|current_date|prev_month|prev_year|next_month|next_year|
+------------+----------+---------+----------+---------+
|  2023-05-15|         4|     2023|         6|     2023|
|  2024-01-01|        12|     2023|         2|     2024|
|  2023-12-31|        11|     2023|         1|     2024|
|  2023-02-28|         1|     2023|         3|     2023|
+------------+----------+---------+----------+---------+

关键说明 #

• add_months函数会自动处理跨年度、月末日期：比如2023-03-31减1个月得到2023-02-28（非闰年），加1个月得到2023-04-30，无需额外逻辑处理边界。
• 如果原始日期列格式非yyyy-MM-dd，需在to_date中指定对应格式（如to_date(col("current_date"), "yyyy/MM/dd")）。

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