PySpark窗口函数未全量生效求助：lag递推计算异常

问题现象 #

需求是从第2周开始，每一行的value1等于前一行计算后的value1乘以当前行的value2。测试数据和初始代码如下：

测试数据：

from pyspark.sql import functions as f

data = [
    (1, 1, 1),
    (2, 0, 5),
    (3, 0, 10),
    (4, 0, 20),
    (5, 0, 30),
    (6, 0, 40)
]
columns = ["week", "value1", "value2"]
df = spark.createDataFrame(data, columns)

初始计算逻辑：

w=Window.orderBy("week")
                          
df2 = df.withColumn('value1',
                    f.when((f.col('week') > 1),
                           f.lag(df['value1']).over(w) * df['value2'] 
                          ).otherwise(
                        f.col('value1')
                    )
                   )

运行后仅第2周value1符合预期（15=5），但第3周及之后的value1仍为0（本该是510=50），未按递推逻辑计算。

问题根源 #

PySpark的窗口函数（如lag）是基于原始DataFrame的列值计算的，不会引用withColumn中新生成的列值。也就是说，lag(df['value1'])调用的是原始数据里的value1，不是刚计算出来的更新后的值。第3周时，lag取到的是原始数据第2周的value1（0），而非计算后的5，因此结果错误。

这种依赖前一行计算结果的递推场景，窗口函数无法处理，需要用递归CTE实现。

解决方案：递归CTE实现递推计算 #

递归CTE可以逐行迭代，用上一行的计算结果更新当前行的值，代码如下：

from pyspark.sql import functions as f
from pyspark.sql.window import Window

# 测试数据初始化
data = [
    (1, 1, 1),
    (2, 0, 5),
    (3, 0, 10),
    (4, 0, 20),
    (5, 0, 30),
    (6, 0, 40)
]
columns = ["week", "value1", "value2"]
df = spark.createDataFrame(data, columns)

# 添加连续行号，确保递归时按week顺序处理
window = Window.orderBy("week")
df = df.withColumn("row_num", f.row_number().over(window))

# 创建临时视图供递归CTE调用
df.createOrReplaceTempView("weekly_data")

# 递归CTE计算递推value1
result_df = spark.sql("""
    WITH recursive cte AS (
        -- 基础部分：取第一行作为初始值
        SELECT week, value1, value2, row_num
        FROM weekly_data
        WHERE row_num = 1
        UNION ALL
        -- 递归部分：用上一行的value1计算当前行的value1
        SELECT 
            curr.week,
            prev.value1 * curr.value2 AS value1,
            curr.value2,
            curr.row_num
        FROM cte prev
        JOIN weekly_data curr ON prev.row_num + 1 = curr.row_num
    )
    SELECT week, value1, value2 FROM cte ORDER BY week
""")

result_df.show()

运行结果 #

+----+-------+------+
|week| value1|value2|
+----+-------+------+
|   1|      1|     1|
|   2|      5|     5|
|   3|     50|    10|
|   4|   1000|    20|
|   5|  30000|    30|
|   6|1200000|    40|
+----+-------+------+

代码说明 #

1. 添加行号：通过row_number()生成连续行号，确保递归时能按week顺序依次处理每一行。
2. 递归CTE结构：
   • 基础部分：取出第一行数据作为递推的初始值。
   • 递归部分：将上一轮CTE的结果（前一行数据）与原始数据的下一行关联，用前一行的value1乘以当前行的value2得到新的value1，再合并到CTE中。
3. 结果输出：从CTE中取出最终数据，按week排序后展示。

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