如何基于Dataframe的range列动态聚合连续行的amount总和？

问题背景 #

现有初始DataFrame，每个id对应的range值一致，需完成以下聚合操作：

1. 按id分组，以range列值为动态连续行范围；
2. 对该范围内的amount列求和；
3. 按id、row_num升序排序。
   尝试过Window.rowsBetween但无法动态调用range列值，不愿使用case...when或硬编码方式，希望借助高级窗口/分区功能解决。

解决方案（以PySpark为例） #

核心思路 #

通过自连接+行号区间过滤实现动态窗口聚合，替代固定参数的rowsBetween：

• 先为每个id提取统一的range值，计算每行对应的聚合行号区间；
• 自连接匹配同id且行号在目标区间内的记录，最终聚合求和。

代码实现 #

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

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

# 模拟初始数据
sample_data = [
    ("A", 1, 2, 10), ("A", 2, 2, 20), ("A", 3, 2, 30), ("A", 4, 2, 40),
    ("B", 1, 1, 5), ("B", 2, 1, 15), ("B", 3, 1, 25)
]
df = spark.createDataFrame(sample_data, ["id", "row_num", "range", "amount"])

# 1. 提取每个id对应的统一range值
id_range_map = df.groupBy("id").agg(F.first("range").alias("range_val"))

# 2. 计算每行的聚合行号区间（确保起始行不小于1，结束行不超过当前id的最大行号）
window = Window.partitionBy("id")
df_with_bounds = df.join(id_range_map, on="id") \
    .withColumn("max_row", F.max("row_num").over(window)) \
    .withColumn("start_row", F.greatest(F.col("row_num") - F.col("range_val"), F.lit(1))) \
    .withColumn("end_row", F.least(F.col("row_num") + F.col("range_val"), F.col("max_row")))

# 3. 自连接匹配目标区间内的记录
joined_df = df_with_bounds.alias("main") \
    .join(df.alias("target"), 
          (F.col("main.id") == F.col("target.id")) &
          (F.col("target.row_num").between(F.col("main.start_row"), F.col("main.end_row"))),
          how="left")

# 4. 聚合求和并排序
result = joined_df.groupBy("main.id", "main.row_num") \
    .agg(F.sum("target.amount").alias("sum_amount")) \
    .orderBy("id", "row_num")

# 输出结果
result.show()

补充说明 #

• 代码中额外处理了结束行号超出当前id最大行号的场景，避免无效匹配；
• 方案完全基于数据动态计算窗口范围，无需硬编码或分支判断枚举场景，适配任意range值的变化；
• 若使用Pandas，可通过merge+布尔索引实现类似逻辑，核心思路一致。

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