如何递归填充PySpark DataFrame列？求可行实现方案

你原来的代码无效是因为在withColumn中引用刚创建的C列的lag值时，Spark的列计算逻辑基于原DataFrame，新列尚未完全生成，导致后续行无法获取上一行的C值，最终返回null。

以下是两种可行的解决方案：

方法一：构造辅助列后累加（推荐，高效简洁） #

核心思路是将需求转化为带初始值的累加求和：第一行用A作为初始值，后续行用B值，对这些值做累加即可得到目标列C。

代码实现：

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

# 定义按row_num排序的窗口
window_spec = Window.orderBy("row_num")

# 1. 构造辅助列temp：第一行取A，其余行取B
df = df.withColumn("temp", F.when(F.col("row_num") == 1, F.col("A")).otherwise(F.col("B")))
# 2. 对temp列做累加和得到C
df = df.withColumn("C", F.sum("temp").over(window_spec))
# 可选：删除辅助列
df = df.drop("temp")

df.show()

运行后输出的DataFrame符合预期：

方法二：递归CTE（适用于复杂依赖场景） #

如果你的row_num不连续或者有更复杂的递归逻辑，可使用递归CTE实现：

代码实现：

from pyspark.sql import functions as F

# 将DataFrame注册为临时视图
df.createOrReplaceTempView("temp_df")

# 递归CTE计算C列
result_df = df.sparkSession.sql("""
    WITH recursive cte AS (
        -- 基例：第一行C=A
        SELECT row_num, A, B, A AS C
        FROM temp_df
        WHERE row_num = 1
        UNION ALL
        -- 递归部分：当前行C = 上一行C + 当前行B
        SELECT t.row_num, t.A, t.B, c.C + t.B AS C
        FROM temp_df t
        JOIN cte c ON t.row_num = c.row_num + 1
    )
    SELECT * FROM cte
""")

result_df.show()

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