使用Lag与Window函数动态更新Spark DataFrame列遇问题求助

你的代码无法正确计算递推式adstock，核心问题在于Spark的列计算是批量执行的：lag("adstock", 1)引用的是执行withColumn之前的旧adstock列值，而非刚计算出的新值。当旧列初始为0时，第一行计算后若遇到lag返回null（每个id的第一行无前置行），结果会变成null，后续行因引用null值持续传播，最终导致大部分行无结果。

Adstock的计算公式是典型的递推关系：adstock_t = col_lag_t + 0.9 * adstock_{t-1}，每个id的第一行（最早日期）的adstock直接等于col_lag。针对Spark的特性，推荐两种实现方式：

方式一：Pandas UDF（简洁高效，适合大分组场景） #

利用pandas_udf按id分组后逐行计算，性能优于递归CTE，适合你12.5万id的大数量场景：

from pyspark.sql.functions import pandas_udf, PandasUDFType
import pandas as pd

def calc_adstock(group: pd.DataFrame) -> pd.DataFrame:
    # 按日期排序保证递推顺序
    group = group.sort_values("dt")
    # 初始化adstock为第一行的col_lag
    group["adstock"] = group["col_lag"].copy()
    # 逐行递推计算
    for i in range(1, len(group)):
        group["adstock"].iloc[i] = group["col_lag"].iloc[i] + 0.9 * group["adstock"].iloc[i-1]
    return group

# 定义UDF，schema需包含原表所有字段+新增的adstock列
adstock_udf = pandas_udf(calc_adstock, schema=df.schema.add("adstock", "double"))

# 按id分组计算，得到最终结果
final_df = df.groupBy("id").apply(adstock_udf).orderBy("id", "dt")

方式二：递归CTE（纯Spark SQL逻辑，无依赖） #

如果无法使用Pandas UDF，可通过递归CTE逐行关联前置行的adstock值：

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

# 为每个id的行添加序号，用于递推关联
window_spec = Window.partitionBy("id").orderBy("dt")
df = df.withColumn("row_num", F.row_number().over(window_spec))

# 基础CTE：每个id的第一行，adstock直接等于col_lag
base_df = df.filter(F.col("row_num") == 1).withColumn("adstock", F.col("col_lag"))

# 递归计算后续行，直到覆盖所有数据
adstock_cte = base_df
total_rows = df.count()

while adstock_cte.count() < total_rows:
    # 获取尚未计算的行
    max_calculated_row = adstock_cte.select(F.max("row_num")).first()[0]
    next_rows = df.filter(F.col("row_num") == max_calculated_row + 1)
    
    # 关联上一行的adstock值进行计算
    new_rows = next_rows.join(
        adstock_cte,
        (next_rows["id"] == adstock_cte["id"]) & (next_rows["row_num"] == adstock_cte["row_num"] + 1),
        "left"
    ).select(
        next_rows["id"],
        next_rows["dt"],
        next_rows["col_lag"],
        next_rows["row_num"],
        (next_rows["col_lag"] + 0.9 * adstock_cte["adstock"]).alias("adstock")
    )
    
    adstock_cte = adstock_cte.unionByName(new_rows)

# 清理多余列，得到最终结果
final_df = adstock_cte.drop("row_num").orderBy("id", "dt")

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