PySpark中如何保留前一行更新值并用于下一行计算？

核心思路 #

由于lag函数只能获取原始行值，无法引用前一行的计算结果，我们可以通过分组累积的方式解决：将数据按非零RANK值分割为独立分组，每组内从初始非零值开始，基于行位置和累积求和完成计算，全程无需循环。

实现步骤 #

• 确定排序依据：必须保证数据有明确的排序列（如id、时间戳），否则行顺序无法固定。
• 生成分组ID：通过累积计数非零RANK值，把连续需要依赖计算的行归为同一组。
• 提取组内初始值：获取每组第一个非零的RANK和RANKA作为计算起点。
• 计算更新后的RANK：组内非零值保持原样，零值用初始值加上组内行位置偏移量。
• 计算更新后的RANKA：组内非零值保持原样，零值用初始值加上后续行薪资的累积和。

完整代码 #

from pyspark.sql import SparkSession
from pyspark.sql.window import Window
from pyspark.sql.functions import col, sum, when, first, row_number

# 初始化Spark会话
spark = SparkSession.builder.appName("CumulativeColumnCalc").getOrCreate()

# 示例数据（需确保有排序列id）
sample_data = [
    (1, 5, 1000, 500),
    (2, 0, 0, 300),
    (3, 0, 0, 200),
    (4, 3, 800, 400),
    (5, 0, 0, 150),
    (6, 0, 0, 250)
]
df = spark.createDataFrame(sample_data, ["id", "RANK", "RANKA", "Salary"])

# 全局排序窗口
order_window = Window.orderBy("id")

# 生成分组ID：累积统计非零RANK的数量作为组标识
df = df.withColumn(
    "group_id",
    sum(when(col("RANK") != 0, 1).otherwise(0)).over(order_window.rowsBetween(Window.unboundedPreceding, 0))
)

# 分组内排序窗口
group_window = Window.partitionBy("group_id").orderBy("id")

# 获取每组的初始RANK和RANKA值
df = df.withColumn("start_rank", first("RANK").over(group_window))
df = df.withColumn("start_ranka", first("RANKA").over(group_window))

# 计算组内行位置
df = df.withColumn("pos", row_number().over(group_window))

# 计算更新后的RANK
df = df.withColumn(
    "updated_RANK",
    when(col("RANK") != 0, col("RANK")).otherwise(col("start_rank") + col("pos") - 1)
)

# 计算薪资贡献及累积和，生成更新后的RANKA
df = df.withColumn(
    "salary_contribution",
    when(col("pos") == 1, 0).otherwise(col("Salary"))
)
df = df.withColumn(
    "cumulative_contribution",
    sum("salary_contribution").over(group_window)
)
df = df.withColumn(
    "updated_RANKA",
    when(col("RANKA") != 0, col("RANKA")).otherwise(col("start_ranka") + col("cumulative_contribution"))
)

# 展示结果
df.select("id", "RANK", "RANKA", "Salary", "updated_RANK", "updated_RANKA").show()

输出结果 #

+---+----+-----+------+-----------+-----------+
| id|RANK|RANKA|Salary|updated_RANK|updated_RANKA|
+---+----+-----+------+-----------+-----------+
|  1|   5| 1000|   500|          5|        1000|
|  2|   0|    0|   300|          6|        1300|
|  3|   0|    0|   200|          7|        1500|
|  4|   3|  800|   400|          3|         800|
|  5|   0|    0|   150|          4|         950|
|  6|   0|    0|   250|          5|        1200|
+---+----+-----+------+-----------+-----------+

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