如何在PySpark DataFrame中基于同列已计算值逐行计算列值

你的问题核心是lag()函数只能读取原始数据的前一行值，无法引用当前批次计算出来的新值——而你要的是迭代计算：每一行的Inv都依赖上一行刚算出的Inv结果，这时候得用Spark的递归CTE（公共表表达式）来实现，lag()解决不了这种依赖计算后值的场景。

实现思路 #

• 基础查询：先给每个id分组内的行按year、month排序并加行号，然后取出每组第一行，用原始的inv值作为初始计算值
• 递归查询：关联上一步的结果，依次遍历每组的后续行，用上一行计算好的current_inv代入公式Inv = 上一行Inv - 当前行sup + 当前行unit，算出当前行的Inv值

代码示例 #

假设你的DataFrame包含id、year、month、sup、units、inv（初始值列），代码如下：

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

# 1. 给每个分组内的行添加排序后的行号
row_num_window = Window.partitionBy("id").orderBy("year", "month")
df_with_row_num = df.withColumn("row_num", F.row_number().over(row_num_window))

# 2. 定义递归CTE
recursive_cte = df_with_row_num.selectExpr(
    "id",
    "year",
    "month",
    "sup",
    "units",
    "row_num",
    "inv as current_inv"  # 第一行用原始inv作为初始值
).where("row_num = 1") \
.unionAll(
    df_with_row_num.join(
        recursive_cte,
        (df_with_row_num.id == recursive_cte.id) & (df_with_row_num.row_num == recursive_cte.row_num + 1)
    ).selectExpr(
        df_with_row_num.id,
        df_with_row_num.year,
        df_with_row_num.month,
        df_with_row_num.sup,
        df_with_row_num.units,
        df_with_row_num.row_num,
        "round(recursive_cte.current_inv - df_with_row_num.sup + df_with_row_num.units, 2) as current_inv"
    )
)

# 3. 生成最终结果，去掉行号列并排序
final_df = recursive_cte.drop("row_num").orderBy("id", "year", "month")

关键说明 #

• 如果你的数据没有初始inv值，在基础查询里把inv as current_inv改成初始化逻辑（比如0 as current_inv）即可
• 递归CTE的性能和分组大小相关，若单分组数据量极大，需要调整Spark的递归深度参数（spark.sql.recursiveCTE.maxIterations）
• 务必保证year、month的排序逻辑和业务要求一致，否则计算顺序会出错

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