Spark如何使用窗口函数获取前一行的所有列值

解决方案 #

完全可以实现，你可以先将需要的列打包为结构体（struct），再对整个结构体调用一次lag函数即可，无需多次调用窗口函数。

具体实现步骤 #

• 第一步：将需要获取的前一行字段封装为struct类型的列
• 第二步：对封装好的结构体列调用lag窗口函数，一次性拿到前一整行的对应字段
• 第三步：将当前行字段和前一行结构体中的字段传入自定义UDF完成计算

代码示例 #

from pyspark.sql import functions as F

# 1. 封装前一行需要的所有字段为struct，调用一次lag拿到前一行的结构体
df_with_last_row = df.withColumn(
    "last_row",
    F.lag(F.struct(F.col("column1"), F.col("column2"), F.col("column3")), 1).over(w)
)

# 2. 调用UDF，传入当前行字段和前一行结构体中的对应字段
df2 = df_with_last_row.withColumn(
    "calculation",
    my_udf(
        F.col("column1"), F.col("column2"), F.col("column3"),
        F.col("last_row.column1"), F.col("last_row.column2"), F.col("last_row.column3")
    )
)

# 计算完成后如果不需要last_row列可以删掉
df2 = df2.drop("last_row")

补充说明 #

这个方案只触发一次窗口计算，性能和调用一次lag的开销完全一致，远优于三次单独调用lag的实现。如果后续需要增加更多取前一行的字段，只需要在struct里新增对应列即可，扩展性也更好。

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