PySpark中基于特定日期条件复制数据行的实现求助

需求说明 #

现有Spark数据集：

需要补充最大日期（2022-11-01）到目标日期（2023-01-01）之间缺失的月份起始日期（2022-12-01、2023-01-01），补充规则为：复制上一年对应月份的Value，将日期替换为缺失的目标日期，最终合并原数据得到如下输出：

实现步骤 #

基于你已有的日期检测代码，补充以下逻辑即可完成需求：

1. 将原数据的日期转为datetime类型，避免字符串操作误差
2. 针对每个缺失日期，计算其上一年同月的日期，从原数据中匹配对应的Value
3. 生成补充数据的DataFrame，与原数据合并
4. 按日期降序排序，匹配示例输出格式

完整代码实现 #

from datetime import datetime
from dateutil.relativedelta import relativedelta
from pyspark.sql import SparkSession
from pyspark.sql.functions import col

# 初始化SparkSession（如果未初始化）
spark = SparkSession.builder.appName("FillMissingDates").getOrCreate()

# 替换为你的shares_union数据集
data = [
    ("2022-10-01", "x"),
    ("2022-11-01", "y"),
    ("2021-12-01", "z"),
    ("2022-01-01", "xy")
]
shares_union = spark.createDataFrame(data, ["DATE", "Value"])
# 将DATE列转为date类型
shares_union = shares_union.withColumn("DATE", col("DATE").cast("date"))

# 获取原数据的最大日期
max_dt = max(shares_union.select("DATE").distinct().rdd.flatMap(lambda x: x).collect())
# 目标日期（你指定的2023-01-01）
target_dt = datetime(2023, 1, 1).date()

# 生成缺失的月份日期列表
diff_months = relativedelta(target_dt, max_dt).months
missing_dates = []
for m in range(1, diff_months + 1):
    missing_dt = (max_dt + relativedelta(months=m)).replace(day=1)
    missing_dates.append(missing_dt)

# 生成补充数据行
supplement_rows = []
for dt in missing_dates:
    # 计算上一年同月的日期
    prev_year_dt = dt - relativedelta(years=1)
    # 从原数据中获取对应Value
    value = shares_union.filter(col("DATE") == prev_year_dt).select("Value").first()[0]
    supplement_rows.append((dt, value))

# 转为DataFrame并与原数据合并
supplement_df = spark.createDataFrame(supplement_rows, ["DATE", "Value"])
final_df = shares_union.union(supplement_df).orderBy(col("DATE").desc())

# 查看结果
final_df.show()

代码说明 #

• 日期类型转换：确保DATE列为标准date类型，简化日期计算逻辑
• 缺失日期生成：沿用你原有的逻辑，精准生成目标区间内的所有月份起始日
• 补充数据匹配：对每个缺失日期，通过年份偏移找到上一年同月的记录，提取对应Value生成新行
• 合并排序：使用union合并原数据与补充数据，orderBy按日期降序排列，完全匹配示例输出格式

运行代码后，输出结果将与你期望的格式完全一致。

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