在PySpark的数据处理工作中，对DataFrame进行多列多函数聚合，再将聚合结果进行重塑是高频需求。比如统计不同部门下多个数值列的最大值、最小值、平均值，再将聚合后的宽表转换为长表格式方便后续分析，这类操作可以通过合理的API组合高效实现。
多列多函数聚合的实现方法
使用agg函数配合字典传参
最常用的方式是先调用groupBy方法指定分组列，再调用agg方法，传入包含列名和聚合函数对应关系的字典，一次性完成多列多函数的聚合计算。
首先构造测试数据：
from pyspark.sql import SparkSession
from pyspark.sql.functions import sum, avg, max, min
# 初始化SparkSession
spark = SparkSession.builder.appName("multi_col_agg_demo").getOrCreate()
# 构造测试DataFrame
data = [
("部门A", "员工1", 100, 200),
("部门A", "员工2", 150, 250),
("部门B", "员工3", 200, 300),
("部门B", "员工4", 180, 280)
]
columns = ["dept", "employee", "score_1", "score_2"]
df = spark.createDataFrame(data, columns)
df.show()
执行多列多函数聚合，对score_1列计算总和与平均值，对score_2列计算最大值与最小值：
# 定义聚合规则字典，key为列名，value为聚合函数
agg_dict = {
"score_1": [sum, avg],
"score_2": [max, min]
}
# 执行分组聚合
agg_result = df.groupBy("dept").agg(*[func(col).alias(f"{col}_{func.__name__}") for col, funcs in agg_dict.items() for func in funcs])
agg_result.show()
使用F表达式简化聚合逻辑
如果聚合函数较多，也可以使用pyspark.sql.functions中的函数表达式来构建聚合逻辑，代码可读性会更好：
from pyspark.sql.functions import col
# 使用F表达式定义聚合列
agg_exprs = [
sum(col("score_1")).alias("score_1_sum"),
avg(col("score_1")).alias("score_1_avg"),
max(col("score_2")).alias("score_2_max"),
min(col("score_2")).alias("score_2_min")
]
agg_result2 = df.groupBy("dept").agg(*agg_exprs)
agg_result2.show()
聚合结果重塑的实现方法
使用stack函数将宽表转长表
聚合后的结果通常是宽表格式，每一列对应一个聚合指标，如果需要转换为长表，也就是将指标名称和对应的值放在两列中，可以使用stack函数实现。
以上面的agg_result2为例，将四个聚合指标重塑为长表：
# 定义需要重塑的列名和对应的别名
cols_to_stack = ["score_1_sum", "score_1_avg", "score_2_max", "score_2_min"]
# stack函数第一个参数是要转换的列数，后面交替写指标名和对应的值列
stack_expr = f"stack({len(cols_to_stack)}, " + ", ".join([f"'{col}', {col}" for col in cols_to_stack]) + ") as (metric, value)"
reshaped_df = agg_result2.select("dept", expr(stack_expr))
reshaped_df.show()
使用pivot与unpivot组合实现复杂重塑
如果聚合结果需要先按某个维度做透视，再转换为长表，可以结合pivot和unpivot操作完成。比如先按部门分组，再按指标类型做透视，最后将透视结果展开为长表：
from pyspark.sql.functions import expr
# 先给聚合结果添加指标类型列，方便后续pivot
agg_result2 = agg_result2.withColumn("metric_type", lit("score"))
# 执行pivot操作，按metric_type透视
pivot_df = agg_result2.groupBy("dept").pivot("metric_type").agg(
sum("score_1_sum").alias("sum_score_1"),
avg("score_1_avg").alias("avg_score_1"),
max("score_2_max").alias("max_score_2"),
min("score_2_min").alias("min_score_2")
)
pivot_df.show()
# 使用unpivot将透视后的宽表转回长表
unpivot_expr = "stack(4, 'sum_score_1', sum_score_1, 'avg_score_1', avg_score_1, 'max_score_2', max_score_2, 'min_score_2', min_score_2) as (metric, value)"
final_df = pivot_df.select("dept", expr(unpivot_expr))
final_df.show()
注意事项
- 使用agg函数传字典时，注意聚合函数要传入函数对象，不要直接调用，否则会提前执行计算。
- stack函数的参数顺序不能错，先写要转换的列数，再交替写指标名称和对应的列名。
- 如果聚合列较多，建议先定义好列名列表，再通过列表推导式生成聚合表达式，减少重复代码。
- 重塑操作会增加数据的行数，如果原始数据量较大，需要注意集群的资源使用情况。