PySpark DataFrame新增列 将YYYYMM格式日期转换为递增月份序号

实现思路 #

直接计算目标月份与基准月份201801的间隔月份数，结果加1即可得到要求的映射值，无需硬编码映射关系，天然适配后续逐月新增的月份数据。

完整代码示例 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import to_date, months_between, lit, col

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

# 示例测试数据，可替换为你自己的DataFrame
test_data = [("201801",), ("201802",), ("202305",), ("202410",)]
df = spark.createDataFrame(test_data, ["date_ym"])

# 新增month_index列作为映射结果
df_result = df.withColumn(
    "month_index",
    # 计算与基准日期2018-01-01的月份差，加1后转为整数
    (months_between(
        # 将YYYYMM字符串转为日期类型，如果date_ym是数字类型需要先转字符串：to_date(col("date_ym").cast("string"), "yyyyMM")
        to_date(col("date_ym"), "yyyyMM"),
        lit("2018-01-01")
    ) + 1).cast("int")
)

# 查看结果
df_result.show()

输出示例 #

+-------+-----------+
|date_ym|month_index|
+-------+-----------+
| 201801|          1|
| 201802|          2|
| 202305|         65|
| 202410|         82|
+-------+-----------+

注意事项 #

• 若你的date_ym列为整数存储的YYYYMM值，需要在to_date中先调用.cast("string")转为字符串类型，否则格式识别会报错
• 该方案无需修改即可适配后续新增的月份数据，会自动计算对应的序号
• 若存在早于201801的日期，会返回小于1的序号，可结合when函数自行定义这部分数据的处理逻辑

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