PySpark使用Lag函数计算递归列C时无法识别字段问题如何解决

问题原因 #

你报错的根本原因是：Spark同一级withColumn操作中，无法引用当前正在定义的列。你在定义C列的逻辑里直接调用lag(F.col('C'))，此时C列还没有完成计算，Spark自然会提示找不到该字段。

另外你的窗口排序逻辑也有问题：你按A字段排序，但所有行的A值都是2，会导致窗口内的行顺序不稳定，最终计算结果可能不符合预期。

解决方法 #

你可以先把递归的计算逻辑做数学展开，完全不需要引用自身列，就能实现需求：
从逻辑可以推导得到：任意行的C值 = A字段值 - 分区内从第一行到当前行的B字段累加和。

完整实现代码如下：

from pyspark.sql import Row, functions as F
from pyspark.sql.window import Window
import pandas as pd

# 初始化数据
data = {'A': [2,2,2,2], 'B': [0.5, 0.5,1,1.5]}
df = pd.DataFrame(data)
ddf = spark.createDataFrame(df)

# 先加索引列保证窗口内行顺序和原始数据一致
ddf = ddf.withColumn('row_idx', F.monotonically_increasing_id())

# 定义窗口：按A分区，按索引排序，范围从分区首行到当前行
win = Window.partitionBy('A').orderBy('row_idx').rowsBetween(Window.unboundedPreceding, Window.currentRow)

# 计算C列
ddf = ddf.withColumn('C', F.col('A') - F.sum('B').over(win)).drop('row_idx')

# 输出结果
ddf.show()

运行后输出和你预期的结果完全一致：

+---+---+----+
|  A|  B|   C|
+---+---+----+
|  2|0.5| 1.5|
|  2|0.5| 1.0|
|  2|1.0| 0.0|
|  2|1.5|-1.5|
+---+---+----+

补充说明 #

如果你的实际场景逻辑更复杂无法做数学展开，可以使用窗口聚合UDF或者逐行迭代的方式处理，但当前场景下累加和的方案是性能最高、实现最简单的方案，支持大数据量的分布式计算。

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