PySpark实现DataFrame两月份列区间内数值生成并新增列

需求说明 #

基于PySpark处理DataFrame，实现如下逻辑：

• 新增new_month列，取值为每行dvpt_month与lead_month两个字段值闭区间内的所有连续数值
• 区间内每个数值对应生成一行，新行除new_month外其余所有字段取值和原行完全一致

示例数据集 #

输入示例 #

输出示例 #

实现代码 #

直接用Spark内置函数实现，执行性能远高于自定义UDF，不需要额外编写复杂逻辑：

1. 导入需要的内置函数

from pyspark.sql.functions import sequence, col, explode

2. 核心处理逻辑

# df为加载完成的原始DataFrame
result_df = df.withColumn(
    "new_month",
    # 先生成dvpt_month到lead_month的连续整数序列，再将数组元素拆分为独立行
    explode(sequence(col("dvpt_month"), col("lead_month")))
)

注意事项 #

• 该写法要求dvpt_month、lead_month为整数类型，如果是字符串格式存储的月份值，需要先通过cast("int")转换为数值类型再处理
• 如果数据中存在dvpt_month > lead_month的异常行，sequence会返回空数组，这类行最终不会出现在结果中，需要保留的话可以提前做值校验、调整步长参数做兼容。

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