如何在PySpark中迭代基于前一行值填充col_b至0或非空行

问题描述 #

需要对col_b列做如下处理：

• 非空值作为起始值
• 后续空行依次减1赋值
• 直到数值变为0，之后的空行保持null

原始数据：

期望结果：

现有代码使用lag函数无法实现需求，因为lag只能读取原始数据的前一行值，无法复用已计算出的填充值；且数据集过大，不能收集到本地处理。

实现思路 #

核心通过分组标记+组内偏移计算实现，无需迭代或本地收集数据：

1. 标记分组：用非空col_b行作为分组起点，通过累积求和生成分组ID
2. 填充基准值：在分组内向前填充最近的非空col_b值，作为每个组的起始基准
3. 计算偏移量：统计当前行在分组内相对于起始行的位置差
4. 生成最终值：用基准值减去偏移量，结果≥0则保留，否则设为null

完整代码实现 #

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

# 初始化SparkSession
spark = SparkSession.builder.appName("FillDecreasingSequence").getOrCreate()

# 原始数据
data = [(1, None), (2, 3), (3, None), (4, None), (5, None), (6, None), (7, 6), (8, None), (9, None)]
df = spark.createDataFrame(data, ["col_a", "col_b"])

# 步骤1：生成分组ID，每个非空col_b作为新分组的起点
group_window = Window.orderBy("col_a")
df = df.withColumn("group_id", 
                   F.sum(F.when(F.col("col_b").isNotNull(), 1).otherwise(0)).over(group_window))

# 步骤2：为每个分组填充基准起始值（最近的非空col_b）
base_value_window = Window.partitionBy("group_id").orderBy("col_a").rowsBetween(Window.unboundedPreceding, Window.currentRow)
df = df.withColumn("base_value", F.last(F.col("col_b"), ignorenulls=True).over(base_value_window))

# 步骤3：计算当前行在分组内的偏移量（相对于起始行的位置差）
row_num_window = Window.partitionBy("group_id").orderBy("col_a")
df = df.withColumn("row_offset", F.row_number().over(row_num_window) - 1)

# 步骤4：计算最终的col_b值，小于0则设为null
df = df.withColumn("col_b", 
                   F.when(
                       F.col("base_value").isNull(), None
                   ).when(
                       (F.col("base_value") - F.col("row_offset")) >= 0, 
                       F.col("base_value") - F.col("row_offset")
                   ).otherwise(None))

# 清理临时列并展示结果
df = df.drop("group_id", "base_value", "row_offset")
df.show()

代码解释 #

1. 分组ID生成：通过sum累积非空值标记，让每个非空col_b后的行归属同一分组，为后续组内计算打基础
2. 基准值填充：last(..., ignorenulls=True)在分组内向前抓取最近的非空值，确保每行都能获取到当前分组的起始值
3. 偏移量计算：row_number()-1得到当前行相对于分组起始行的偏移数（起始行偏移为0，下一行偏移为1，以此类推）
4. 最终值计算：用基准值减去偏移量，判断结果是否≥0，满足则保留，否则设为null；同时处理无基准值的行（如第一行null）

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