PySpark代码修改：当week>1时用前一行value1减当前行value2

问题分析 #

从你给出的期望输出可以推断，实际需要实现递推计算：当week>1时，当前行的value1等于当前行value2减去前一行更新后的value1（即value1(n) = value2(n) - value1(n-1)，需求描述可能存在笔误）。

现有代码的核心问题：

• 逻辑符号错误：代码中写的是lag(value1) - value2，与期望的递推逻辑相反
• 误用sum窗口函数：sum(...) over(w)会做累积求和，而非递推依赖前一行的计算结果
• 普通lag函数只能获取原始数据的前一行值，无法获取更新后的value1

正确解决方案 #

方法1：递归CTE（最直观的递推实现）

Spark支持递归CTE，可以直接实现这种依赖前一行结果的递推计算：

from pyspark.sql import SparkSession

# 初始化SparkSession（若未初始化）
spark = SparkSession.builder.appName("RecursiveUpdate").getOrCreate()

# 原始数据
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)

# 注册临时视图用于SQL查询
df.createOrReplaceTempView("temp_df")

# 递归CTE计算递推的value1
spark.sql("""
WITH RECURSIVE cte AS (
    -- 初始化：取第一行数据
    SELECT week, value1, value2 FROM temp_df WHERE week = 1
    UNION ALL
    -- 递归逻辑：当前行value1 = 当前行value2 - 上一行的value1
    SELECT df.week, df.value2 - cte.value1, df.value2 
    FROM temp_df df
    JOIN cte ON df.week = cte.week + 1
)
SELECT * FROM cte ORDER BY week
""").show()

执行后会得到期望输出：

+----+------+------+
|week|value1|value2|
+----+------+------+
|1   |1     |1     |
|2   |4     |5     |
|3   |6     |10    |
|4   |14    |20    |
|5   |16    |30    |
|6   |24    |40    |
+----+------+------+

方法2：DataFrame API结合窗口函数（针对特定递推规律）

如果偏好DataFrame API，可以利用递推序列的数学特性，通过构造辅助列实现：

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

w = Window.orderBy("week")

# 构造交替符号列：奇数周为1，偶数周为-1
df = df.withColumn("sign", f.when(f.col("week") % 2 == 1, 1).otherwise(-1))
# 构造带符号的value2项
df = df.withColumn("term", f.col("value2") * f.col("sign"))
# 累积求和后根据符号调整得到最终value1
df_final = df.withColumn(
    "value1",
    f.when(f.col("week") == 1, f.col("value1"))
    .otherwise(f.sum(f.col("term")).over(w) * f.col("sign"))
).drop("sign", "term")

df_final.show()

此方法适用于这种交替递推的特定场景，而递归CTE更通用，适合各种复杂递推逻辑。

现有代码失效原因 #

1. 逻辑方向错误：lag(value1) - value2是前一行value1减当前value2，与期望的value2 - 前一行value1相反
2. 累积求和误用：sum(...) over(w)会累加所有行的计算项，而非每一步依赖前一行的结果
3. 原始值依赖：普通lag函数只能获取原始数据中的value1（比如week3的lag(value1)是原始的0，而非更新后的week2的value1=4）

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