PySpark中使用lag函数与when条件生成Round字段时后续值为空的问题

问题原因分析 #

你当前的实现存在两个核心问题：

1. 列计算依赖顺序问题：Spark的withColumn基于当前DataFrame状态计算，第二次赋值Round时，lag("Round")引用的是第一次赋值后的Round列（仅第一行有值，其余为null），导致后续行无法获取正确的前一行新值。
2. 窗口函数误用：lag是固定偏移量的窗口函数，不需要指定rowsBetween这类窗口框架，强行添加会触发AnalysisException。

正确实现思路 #

治疗轮次的本质是基于间隔天数的累计计数：每次间隔≥10天时轮次+1，否则保持不变。可以通过以下两步实现：

• 先创建一个增量标识列，标记需要轮次递增的行；
• 对增量列做累计求和，再加上初始值1得到最终轮次。

完整代码实现 #

from pyspark.sql.functions import col, when, sum as spark_sum
from pyspark.sql.window import Window

data = sc.parallelize([
    ('A', 1, None),
    ('A', 2, 3),
    ('A', 3, 13),
    ('A', 4, 4),
    ('B', 1, None),
    ('B', 2, 22),
    ('B', 3, 3),
    ('B', 4, 14),
    ('B', 5, 11),
])

df_reprex = spark.createDataFrame(data, ['ID', 'Event', 'Gap_Day'])

# 1. 创建增量标识列：间隔≥10天则+1，否则0；第一行(Gap_Day为null)标记为0
df_reprex = df_reprex.withColumn(
    "round_increment",
    when(col("Gap_Day").isNull(), 0)
    .when(col("Gap_Day") >= 10, 1)
    .otherwise(0)
)

# 2. 按ID分区、Event排序，计算累计增量，再加1得到Round
window_def = Window.partitionBy("ID").orderBy("Event").rowsBetween(Window.unboundedPreceding, Window.currentRow)
df_reprex = df_reprex.withColumn(
    "Round",
    spark_sum(col("round_increment")).over(window_def) + 1
)

# 查看结果
df_reprex.select("ID", "Event", "Gap_Day", "Round").show()

输出结果验证 #

执行后会得到符合预期的结果：

+---+-----+-------+-----+
| ID|Event|Gap_Day|Round|
+---+-----+-------+-----+
|  A|    1|   null|    1|
|  A|    2|      3|    1|
|  A|    3|     13|    2|
|  A|    4|      4|    2|
|  B|    1|   null|    1|
|  B|    2|     22|    2|
|  B|    3|      3|    2|
|  B|    4|     14|    3|
|  B|    5|     11|    4|
+---+-----+-------+-----+

关键说明 #

• 累计求和窗口：rowsBetween(Window.unboundedPreceding, Window.currentRow)确保计算从当前分区的第一行到当前行的累计增量，完美匹配轮次的递增逻辑；
• 避免循环依赖：通过独立的增量列计算，无需引用正在生成的Round列，彻底解决null值传递问题。

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