PySpark技术问询：如何将后续指定行数的求和值赋值给当前行

嘿，作为PySpark新手遇到这种行关联计算很正常，咱们用窗口函数就能完美解决这个需求！核心思路是先给每个分组内的行确定顺序，再通过窗口范围选取当前行之后的指定行数求和，最后处理无后续行的情况。

步骤拆解 & 代码实现 #

咱们直接用你的示例数据来演示，假设要取后续2行的求和：

1. 导入依赖并初始化SparkSession

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

# 初始化SparkSession（如果你的环境已经有了可以跳过）
spark = SparkSession.builder.appName("NextRowsSum").getOrCreate()

2. 创建示例测试数据

# 原始数据
data = [
    (1, 2), (1, 3), (1, 1), (1, 9),
    (2, 1), (2, 6), (2, 8), (2, 1)
]
df = spark.createDataFrame(data, schema=["ID", "val"])

3. 给每个分组内的行添加顺序标记
   因为PySpark默认不保证数据的顺序，所以我们需要给每个ID分组内的行按原始顺序添加行号，确保后续能准确选取后续行：

# 定义分组排序的窗口，生成行号rn
row_num_window = Window.partitionBy("ID").orderBy(F.monotonically_increasing_id())
df_with_rn = df.withColumn("rn", F.row_number().over(row_num_window))

这里用monotonically_increasing_id()是因为你的原始数据没有自带排序字段，如果有时间戳/业务顺序字段，直接用那个字段排序更稳妥！

4. 定义后续行求和的窗口并计算sum_val
   我们需要的窗口是：按ID分组，按行号排序，选取当前行之后的1到2行（也就是后续2行）进行求和：

# 指定要取的后续行数
N = 2

# 定义求和窗口：当前行+1 到 当前行+N
sum_window = Window.partitionBy("ID").orderBy("rn").rowsBetween(F.currentRow() + 1, F.currentRow() + N)

# 计算求和，并用coalesce把无后续行的null转为0
result_df = df_with_rn.withColumn(
    "sum_val",
    F.coalesce(F.sum("val").over(sum_window), F.lit(0))
).select("ID", "sum_val")

5. 查看结果

result_df.show()

运行后就能得到和你示例完全一致的输出：

+---+-------+
| ID|sum_val|
+---+-------+
|  1|      4|
|  1|     10|
|  1|      9|
|  1|      0|
|  2|     14|
|  2|      9|
|  2|      1|
|  2|      0|
+---+-------+

关键知识点说明 #

• rowsBetween vs rangeBetween：这里必须用rowsBetween，因为我们是按行的数量来选取范围，而不是按字段值的范围。
• coalesce的作用：当当前行后面没有足够的行时，sum函数会返回null，用coalesce把这些null替换成0，满足你的需求。
• 灵活调整N值：只需要修改代码里的N变量，就能轻松切换成后续3行、5行等不同的求和需求。

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