PySpark DataFrame日期展开与缺失行回填实现问询

原始DataFrame #

+---+----------+------+
| id|      date|amount|
+---+----------+------+
|123|2022-11-11|100.00|
|123|2022-11-12|100.00|
|123|2022-11-13|100.00|
|123|2022-11-14|200.00|
|456|2022-11-14|300.00|
|456|2022-11-15|300.00|
|456|2022-11-16|300.00|
|789|2022-11-11|400.00|
|789|2022-11-12|500.00|
+---+----------+------+

需求说明 #

为每个id生成截至current_date() - 2的所有连续日期记录，缺失日期对应的amount值使用该id最近的已有值填充。例如当date_sub(current_date(), 2) = "2022-11-16"时，目标DataFrame如下：

+------+----------+-------+
|id    |    date  | amount|
+------+----------+-------+
|   123|2022-11-11|100,00 |
|   123|2022-11-12|100,00 |
|   123|2022-11-13|100,00 |
|   123|2022-11-14|200,00 |
|   123|2022-11-15|200,00 |
|   123|2022-11-16|200,00 |
|   456|2022-11-14|300,00 |
|   456|2022-11-15|300,00 |
|   456|2022-11-16|300,00 |
|   789|2022-11-11|400,00 |
|   789|2022-11-12|500,00 |
|   789|2022-11-13|500,00 |
|   789|2022-11-14|500,00 |
|   789|2022-11-15|500,00 |
|   789|2022-11-16|500,00 |
+------+----------+-------+

初始化代码 #

import findspark
findspark.init()

import pyspark
from pyspark.sql import SparkSession

spark = SparkSession.builder.master("local[4]").appName("Complete Rows").getOrCreate()

from pyspark.sql.functions import *
from pyspark.sql.types import StructType,StructField, IntegerType, DateType, DecimalType
from datetime import datetime
from decimal import Decimal

vdata = [
    (123,datetime.strptime('2022-11-11','%Y-%m-%d'),Decimal(100)),
    (123,datetime.strptime('2022-11-12','%Y-%m-%d'),Decimal(100)),
    (123,datetime.strptime('2022-11-13','%Y-%m-%d'),Decimal(100)),
    (123,datetime.strptime('2022-11-14','%Y-%m-%d'),Decimal(200)),
    (456,datetime.strptime('2022-11-14','%Y-%m-%d'),Decimal(300)),
    (456,datetime.strptime('2022-11-15','%Y-%m-%d'),Decimal(300)),
    (456,datetime.strptime('2022-11-16','%Y-%m-%d'),Decimal(300)),
    (789,datetime.strptime('2022-11-11','%Y-%m-%d'),Decimal(400)),
    (789,datetime.strptime('2022-11-12','%Y-%m-%d'),Decimal(500))]

schema = StructType([
    StructField("id",IntegerType(),False),
    StructField("date",DateType(),False),
    StructField("amount",DecimalType(10,2),False)])

df = spark.createDataFrame(vdata,schema)

df.show()

解决方案 #

步骤1：定义目标结束日期并获取每个id的起始日期 #

先计算目标结束日期end_date，再分组提取每个id的最小日期作为起始日期：

# 定义目标结束日期：当前日期减2天
end_date = date_sub(current_date(), 2)

# 获取每个id的起始日期（该id最早的记录日期）
id_dates = df.groupBy("id").agg(min("date").alias("start_date"))

步骤2：为每个id生成完整日期序列 #

用sequence函数生成从起始日期到目标结束日的连续日期，再通过explode将序列展开为每行一个日期：

# 生成连续日期序列并展开
full_dates = id_dates.withColumn(
    "date",
    explode(sequence(col("start_date"), end_date, expr("interval 1 day")))
).drop("start_date")

步骤3：关联原始数据并填充缺失值 #

将完整日期表与原始数据左关联，再用窗口函数结合last函数向前填充缺失的amount值：

from pyspark.sql.window import Window

# 定义窗口：按id分区，按日期升序排序
window_spec = Window.partitionBy("id").orderBy("date").rowsBetween(Window.unboundedPreceding, Window.currentRow)

# 左关联后填充缺失的amount值
result_df = full_dates.join(df, on=["id", "date"], how="left")\
    .withColumn("amount", last("amount", ignorenulls=True).over(window_spec))\
    .select("id", "date", "amount")

result_df.show()

关键说明 #

• sequence函数确保每个id拥有从起始日到目标结束日的所有连续日期
• last(..., ignorenulls=True)会在窗口内取当前行及之前最近的非空amount值，实现向前填充逻辑
• 窗口的rowsBetween设置保证只引用历史数据，避免提前使用未来的amount值

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