PySpark中如何为多列批量应用窗口聚合函数？

完全可以实现类似groupBy.agg()的批量窗口聚合操作，以下是简洁高效的实现方案，适配列数量不固定的场景：

步骤1：定义窗口规范 #

首先根据需求创建窗口对象：

from pyspark.sql import Window
from pyspark.sql.functions import col, sum, mean

# 定义窗口：按col1分区，col2排序，范围从分区起始到当前行的前一行
window_spec = Window.partitionBy("col1").orderBy("col2").rangeBetween(Window.unboundedPreceding, -1)

步骤2：指定排除列并筛选目标列 #

明确不需要添加聚合列的字段（比如分区列、排序列或其他业务无关列），筛选出需要处理的列：

# 替换为你实际需要排除的列名列表
exclude_cols = ["col1", "col2"]
# 筛选出所有需要生成聚合列的目标列
target_cols = [col_name for col_name in df.columns if col_name not in exclude_cols]

步骤3：批量生成窗口聚合表达式 #

通过循环动态生成每个目标列的多种窗口聚合表达式，并为新列设置清晰的别名：

# 批量生成sum、mean的窗口聚合表达式
window_exprs = [
    func(col(col_name)).over(window_spec).alias(f"{col_name}_{func.__name__}")
    for col_name in target_cols
    for func in [sum, mean]
]

步骤4：执行查询生成结果 #

使用select保留原DataFrame的所有列，并添加生成的聚合列：

# 生成包含原列和所有窗口聚合列的结果DataFrame
result_df = df.select("*", *window_exprs)

关键说明 #

• 该方案完全适配列数量不固定的场景，即使有数百至数千列，也无需手动编写每一列的聚合逻辑
• 别名格式{原列名}_{函数名}可以避免列名冲突，也便于后续识别聚合类型
• 若需要扩展其他窗口函数（如max、min），只需在[sum, mean]列表中添加对应的函数即可
• 注意窗口函数的性能特性：合理选择partitionBy的列，避免数据倾斜；rangeBetween或rowsBetween的范围设置要符合业务需求

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