PySpark 2.4中按分组日期实现线性插值的方法

针对你提供的DataFrame，要按group分组后基于date对value列做线性插值，核心难点是获取连续空值对应的最近前非空值/日期和最近后非空值/日期，以下是具体实现方案：

完整实现代码 #

from pyspark.sql import SparkSession
from pyspark.sql.window import Window
from pyspark.sql.functions import col, last, first, datediff, when

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

# 准备数据
data = [
    ("A", "2022-01-01", 10.0),
    ("A", "2022-01-02", None),
    ("A", "2022-01-03", None),
    ("A", "2022-01-04", 40.0),
    ("B", "2022-01-01", 30.0),
    ("B", "2022-01-02", None),
    ("B", "2022-01-03", 60.0)
]
columns = ["group", "date", "value"]
df = spark.createDataFrame(data, columns)

# 转换日期类型
df = df.withColumn("date", col("date").cast("date"))

# 定义窗口1：获取当前行及之前的最近非空值/日期（前向填充最近非空值）
window_prev = Window.partitionBy("group").orderBy("date").rowsBetween(Window.unboundedPreceding, Window.currentRow)
df = df.withColumn("prev_value", last("value", ignoreNulls=True).over(window_prev))
df = df.withColumn("prev_date", last("date", ignoreNulls=True).over(window_prev))

# 定义窗口2：获取当前行及之后的最近非空值/日期（后向填充最近非空值）
window_next = Window.partitionBy("group").orderBy("date").rowsBetween(Window.currentRow, Window.unboundedFollowing)
df = df.withColumn("next_value", first("value", ignoreNulls=True).over(window_next))
df = df.withColumn("next_date", first("date", ignoreNulls=True).over(window_next))

# 计算线性插值：仅对value为空的行计算，非空行保留原值
df = df.withColumn(
    "interpolated_value",
    when(
        col("value").isNull(),
        col("prev_value") + (col("next_value") - col("prev_value")) * (datediff(col("date"), col("prev_date")) / datediff(col("next_date"), col("prev_date")))
    ).otherwise(col("value"))
)

# 展示结果
df.select("group", "date", "value", "prev_value", "next_value", "prev_date", "next_date", "interpolated_value").orderBy("group", "date").show()

关键逻辑说明 #

• 前向非空值获取：通过window_prev窗口范围从分组开头到当前行，结合last(..., ignoreNulls=True)跳过空值，拿到当前行最近的前一个非空value和对应的date，解决连续空值下lag只能取相邻行的问题。
• 后向非空值获取：通过window_next窗口范围从当前行到分组末尾，结合first(..., ignoreNulls=True)拿到当前行最近的后一个非空value和对应的date，替代lead无法跨空值取后续非空值的缺陷。
• 线性插值计算：使用公式 prev_value + (next_value - prev_value) * (当前日期与前非空日期的天数差 / 后非空日期与前非空日期的天数差)，仅对空值行计算插值，非空行保留原值。

输出结果 #

+-----+----------+-----+----------+----------+----------+----------+-----------------+
|group|      date|value|prev_value|next_value| prev_date| next_date|interpolated_value|
+-----+----------+-----+----------+----------+----------+----------+-----------------+
|    A|2022-01-01| 10.0|      10.0|      40.0|2022-01-01|2022-01-04|             10.0|
|    A|2022-01-02| null|      10.0|      40.0|2022-01-01|2022-01-04|             20.0|
|    A|2022-01-03| null|      10.0|      40.0|2022-01-01|2022-01-04|             30.0|
|    A|2022-01-04| 40.0|      40.0|      40.0|2022-01-04|2022-01-04|             40.0|
|    B|2022-01-01| 30.0|      30.0|      60.0|2022-01-01|2022-01-03|             30.0|
|    B|2022-01-02| null|      30.0|      60.0|2022-01-01|2022-01-03|             45.0|
|    B|2022-01-03| 60.0|      60.0|      60.0|2022-01-03|2022-01-03|             60.0|
+-----+----------+-----+----------+----------+----------+----------+-----------------+

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