如何创建PySpark递归填充列？现有代码返回空值求解决方案

你原代码的问题在于：在withColumn中尝试引用**刚定义的new_column**的lag值，Spark在计算该列时，new_column还未生成，因此lag(F.col("new_column"))会返回null，导致后续行的计算结果为null或报错。

针对你的需求，有两种高效的实现方式：

方案一：利用累计和窗口函数（推荐，大数据量友好） #

观察你的期望结果可以发现：

• 第1行new_column=0
• 第n行（n>1）new_column = 第2行到第n行的other_column累加值

可以通过计算other_column的累计和，再减去第1行的other_column值来实现：

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

# 初始化示例数据
df = spark.createDataFrame([(1, 10), (2, 5), (3, 7), (4, 3)], ["row_num", "other_column"])

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

# 计算other_column的累计和
cumulative_sum = F.sum(F.col("other_column")).over(window_spec)
# 获取第一行的other_column值
first_other = F.first(F.col("other_column")).over(window_spec)

# 生成new_column
df_result = df.withColumn(
    "new_column",
    F.when(F.col("row_num") == 1, 0).otherwise(cumulative_sum - first_other)
)

df_result.show()

执行后输出：

+-------+-------------+----------+
|row_num|other_column|new_column|
+-------+-------------+----------+
|      1|           10|         0|
|      2|            5|         5|
|      3|            7|        12|
|      4|            3|        15|
+-------+-------------+----------+

方案二：递归CTE（适用于更复杂的依赖逻辑） #

如果你的依赖逻辑更复杂（比如初始值不是基于第一行的other_column），可以使用递归CTE实现：

# 注册DataFrame为临时表
df.createOrReplaceTempView("temp_df")

# 递归CTE计算
df_result = spark.sql("""
WITH RECURSIVE cte AS (
    -- 初始行：第一行new_column固定为0
    SELECT row_num, other_column, 0 AS new_column
    FROM temp_df WHERE row_num = 1
    -- 递归逻辑：后续行 = 前一行new_column + 当前行other_column
    UNION ALL
    SELECT t.row_num, t.other_column, c.new_column + t.other_column AS new_column
    FROM temp_df t
    JOIN cte c ON t.row_num = c.row_num + 1
)
SELECT * FROM cte ORDER BY row_num
""")

df_result.show()

该方案的输出与方案一一致，但递归CTE在数据量极大时性能可能不如窗口函数，因此优先推荐方案一。

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