PySpark中Lag窗口函数无法逐行递推计算问题求助

你的问题根源在于PySpark窗口函数的lag是基于原始DataFrame的列值，无法引用同一列中刚计算出来的新值。你需要的是逐行递推的迭代计算，这种场景下递归CTE（公共表表达式）是更合适的方案。

解决方案代码 #

from pyspark.sql import functions as f
from pyspark.sql.types import IntegerType

# 原始数据初始化
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)
df = df.withColumn("week", f.col("week").cast(IntegerType()))

# 注册临时表供递归CTE使用
df.createOrReplaceTempView("weekly_data")

# 执行递归CTE计算
result_df = spark.sql("""
    WITH RECURSIVE recursive_calc AS (
        -- 初始行：取第1周的原始数据
        SELECT week, value1, value2 FROM weekly_data WHERE week = 1
        UNION ALL
        -- 递推逻辑：用上一周计算后的value1减去当前周的value2
        SELECT 
            w.week,
            rc.value1 - w.value2 AS value1,
            w.value2
        FROM recursive_calc rc
        JOIN weekly_data w ON rc.week + 1 = w.week
    )
    SELECT * FROM recursive_calc ORDER BY week
""")

result_df.show()

运行结果 #

+----+------+------+
|week|value1|value2|
+----+------+------+
|   1|     1|     1|
|   2|    -4|     5|
|   3|   -14|    10|
|   4|   -34|    20|
|   5|   -64|    30|
|   6|  -104|    40|
+----+------+------+

逻辑说明 #

1. 初始部分：获取第1周的原始数据作为递推的起点。
2. 递归部分：通过rc.week + 1 = w.week关联下一周的数据，用上一轮计算得到的value1减去当前周的value2，生成当前周的新value1。
3. 递归过程会逐行迭代，直到遍历完所有周的数据。

原代码无效原因 #

你之前的代码中，lag(df['value1']).over(w)取的是原始DataFrame中前一行的value1值，而非计算后的新值。比如第3周时，lag取的是原始数据中week=2的value1（0），而不是计算后的-4，因此结果不符合预期。

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