PySpark中如何对单键多值DataFrame执行透视操作？

需求说明 #

需将包含key列和多个score列的DataFrame进行透视，把key的不同取值与各score列组合为新列，将原多行数据按分组列合并为单行。

示例输入DataFrame #

期望输出DataFrame #

实现步骤与代码 #

1. 导入依赖并创建示例DataFrame #

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

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

# 构造示例数据
data = [
    ("ABC", 1, "complete", "q1", 1, 2, 3),
    ("ABC", 1, "complete", "q2", 4, 5, 6),
    ("ABC", 2, "complete", "q1", 1, 6, 7),
    ("ABC", 2, "complete", "q2", 5, 6, 7)
]
schema = ["id", "test_id", "test_status", "key", "score1", "score2", "score3"]
df = spark.createDataFrame(data, schema)

2. 将多Score列转为长格式 #

使用stack函数把多个score列拆分为（score列名, score值）的行结构，方便后续组合key列：

# 定义分组列（透视后保留的列）
group_cols = ["id", "test_id", "test_status"]

# 转换为长格式：将score1/score2/score3拆分为score_col和score_val
melted_df = df.select(
    *group_cols,
    "key",
    F.expr("stack(3, 'score1', score1, 'score2', score2, 'score3', score3) as (score_col, score_val)")
)

stack函数参数说明：第一个数字是要拆分的列数量，后续依次是列名和对应的列值。

3. 组合Key与Score列名 #

拼接key和score列名，生成透视后的目标列名（如q1_score1）：

pivot_ready_df = melted_df.withColumn(
    "pivot_col",
    F.concat(F.col("key"), F.lit("_"), F.col("score_col"))
)

4. 分组透视生成结果 #

按分组列聚合，透视拼接后的列名，并取对应score值：

result_df = pivot_ready_df.groupBy(group_cols).pivot("pivot_col").agg(F.first("score_val"))

# 查看结果
result_df.show()

执行后即可得到与期望一致的输出DataFrame。

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