如何在PySpark中以另一列变量为偏移量实现列的lead/lag偏移

报错原因 #

Spark SQL 的 LEAD/LAG 函数偏移量参数仅支持固定常量，不支持传入列作为动态偏移量，这是和 SQL Server 语法的核心差异，因此你之前的写法会抛出数据类型不匹配异常。

可行实现方案 #

下面提供两种可直接运行的实现方式：

方案1：自关联实现（兼容性好） #

通过排序生成行号后自关联匹配偏移行，逻辑直观，兼容所有 Spark 版本：

from pyspark.sql import SparkSession
from pyspark.sql import functions as F
from pyspark.sql.window import Window

# 初始化Spark会话
spark = SparkSession.builder.appName("dynamic_lead").getOrCreate()

# 构造样例数据
data = [
    ("2021/02/11", 50.12, 0),
    ("2021/02/12", 72.30, 4),
    ("2021/02/15", 81.87, 1),
    ("2021/02/16", 90.12, 2),
    ("2021/02/17", 91.31, 1),
    ("2021/02/18", 81.23, 2),
    ("2021/02/19", 73.45, 1),
    ("2021/02/22", 87.17, 0)
]
df = spark.createDataFrame(data, schema=["Date", "Value", "Shift_Index"])

# 按日期排序生成从0开始的连续行号
w_row = Window.orderBy("Date")
df = df.withColumn("row_id", F.row_number().over(w_row) - 1)

# 自关联匹配偏移行
df_result = df.alias("t1").join(
    df.alias("t2"),
    F.col("t1.row_id") + F.col("t1.Shift_Index") == F.col("t2.row_id"),
    how="left"
).select(
    "t1.Date",
    "t1.Value",
    "t1.Shift_Index",
    F.coalesce(F.col("t2.Value"), F.col("t1.Value")).alias("New_Value")
).orderBy("Date")

df_result.show()

如果偏好SQL写法，可使用如下语句：

-- 注册临时视图
df.createOrReplaceTempView("sample_data")

-- 执行查询
WITH numbered_data AS (
    SELECT 
        Date,
        Value,
        Shift_Index,
        ROW_NUMBER() OVER(ORDER BY Date ASC) - 1 AS row_id
    FROM sample_data
)
SELECT 
    t1.Date,
    t1.Value,
    t1.Shift_Index,
    COALESCE(t2.Value, t1.Value) AS New_Value
FROM numbered_data t1
LEFT JOIN numbered_data t2 ON t1.row_id + t1.Shift_Index = t2.row_id
ORDER BY t1.Date ASC

方案2：数组索引实现（性能更优） #

通过窗口函数将所有Value按排序顺序收集为数组，直接按索引取偏移值，避免自关联开销，适合中小数据集：

w_full = Window.orderBy("Date").rowsBetween(Window.unboundedPreceding, Window.unboundedFollowing)
w_row = Window.orderBy("Date")

df_result = df.withColumn("value_list", F.collect_list("Value").over(w_full))\
    .withColumn("curr_idx", F.row_number().over(w_row))\
    .withColumn("New_Value",
        F.when(
            F.col("curr_idx") + F.col("Shift_Index") <= F.size("value_list"),
            F.element_at("value_list", F.col("curr_idx") + F.col("Shift_Index"))
        ).otherwise(F.col("Value"))
    ).select("Date", "Value", "Shift_Index", "New_Value")

df_result.show()

两种方案运行后均可得到你需要的结果。

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