Spark 2.3中替代sequence函数生成周度日期序列的方案

问题描述 #

原本借助Spark 2.4的sequence函数可以快速生成两个日期间的周度日期序列，但生产环境为Spark 2.3，无法使用该函数，需要实现相同功能。原代码示例如下：

data1 = [
    (1, "2022-09-01", "2023-01-01", 1),
    (1, "2022-09-01", "2023-02-01", 1),
    (1, "2022-09-11", "2023-01-01", 2),
    (1, "2022-09-01", "2023-01-01", 2),
    (1, "2022-09-21", "2023-01-01", 1),
]
df1 = spark.createDataFrame(
    data1, ["item", "start_d", "activation_d", "dept_id"]
)
df1 = df1.withColumn(
    "week_start",
    SF.explode(SF.expr("sequence(start_d, activation_d, interval 7 day)"))
)

解决方案 #

在Spark 2.3中，可以通过计算日期差生成索引序列，再基于起始日期推导周度日期的方式替代sequence函数，步骤如下：

1. 转换日期类型：确保start_d和activation_d为日期类型（原数据是字符串，需先转换）
2. 计算周数范围：计算两个日期之间的总天数差，除以7得到需要生成的周数（向上取整）
3. 生成索引序列：使用posexplode和range函数生成从0开始的索引序列
4. 推导周起始日期：基于起始日期加上7 * 索引天，得到每个周的起始日期
5. 过滤无效日期：移除超过activation_d的日期（避免最后一周超出结束范围）

完整代码实现 #

from pyspark.sql import functions as SF

data1 = [
    (1, "2022-09-01", "2023-01-01", 1),
    (1, "2022-09-01", "2023-02-01", 1),
    (1, "2022-09-11", "2023-01-01", 2),
    (1, "2022-09-01", "2023-01-01", 2),
    (1, "2022-09-21", "2023-01-01", 1),
]
df1 = spark.createDataFrame(
    data1, ["item", "start_d", "activation_d", "dept_id"]
)

# 转换为日期类型
df1 = df1.withColumn("start_d", SF.to_date("start_d")) \
         .withColumn("activation_d", SF.to_date("activation_d"))

# 计算需要生成的周数（向上取整，+1确保包含起始周）
df_with_week_count = df1.withColumn(
    "week_count",
    SF.ceil(SF.datediff("activation_d", "start_d") / 7) + 1
)

# 生成周起始日期序列并过滤超出范围的日期
df_result = df_with_week_count.select(
    "item", "start_d", "activation_d", "dept_id",
    SF.posexplode(SF.expr("range(0, week_count)")).alias("idx", "week_offset")
).withColumn(
    "week_start",
    SF.date_add("start_d", SF.col("week_offset") * 7)
).filter(
    SF.col("week_start") <= SF.col("activation_d")
).drop("week_count", "idx", "week_offset")

# 查看结果
df_result.show()

关键说明 #

• datediff("activation_d", "start_d")计算两个日期的天数差，除以7得到周数，ceil向上取整确保覆盖最后一周
• range(0, week_count)生成从0到周数-1的索引序列，posexplode将序列拆分为多行
• date_add("start_d", week_offset *7)基于起始日期加上对应周数的天数，得到每周起始日期
• 最后过滤掉超过activation_d的日期，保证结果和原sequence函数逻辑一致

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