PySpark DataFrame自定义Melt实现：按科目列表转换数据结构

场景还原 #

假设你的原始DataFrame结构如下（以English、History两个科目为例）：

需要将科目转为行记录，最终得到：

解决方案：动态生成Stack表达式 #

利用stack函数的多值返回特性，结合动态字符串拼接，实现自定义Melt。核心思路是将每个科目对应的所有指标列打包成stack的一个条目，动态适配传入的科目列表。

代码实现

from pyspark.sql import SparkSession

# 初始化SparkSession
spark = SparkSession.builder.appName("CustomMelt").getOrCreate()

# 构造示例数据
data = [
    ("Alice", 18, 85, 95, 78, 60, 79, 90, 72, 55),
    ("Bob", 17, 92, 95, 78, 60, 88, 90, 72, 55)
]
columns = [
    "Name", "Age",
    "English_Marks", "English_Highest", "English_Avg", "English_Lowest",
    "History_Marks", "History_Highest", "History_Avg", "History_Lowest"
]
df = spark.createDataFrame(data, columns)

def custom_melt(df, id_vars, subject_list, metric_suffixes):
    # 生成stack的每个条目：'科目名', 科目_指标1, 科目_指标2,...
    stack_entries = []
    for subject in subject_list:
        cols = [f"'{subject}'"] + [f"`{subject}_{suffix}`" for suffix in metric_suffixes]
        stack_entries.append(", ".join(cols))
    
    # 构造stack表达式
    stack_expr = f"stack({len(subject_list)}, {', '.join(stack_entries)}) as (Subject, {', '.join(metric_suffixes)})"
    
    # 执行selectExpr，保留id列+stack结果
    return df.selectExpr(*id_vars, stack_expr)

# 调用自定义melt函数
id_columns = ["Name", "Age"]
subjects = ["English", "History"]
metrics = ["Marks", "Highest", "Avg", "Lowest"]

melted_df = custom_melt(df, id_columns, subjects, metrics)
melted_df.show()

关键说明 #

• stack(N, ...)中的N是科目数量，每个科目对应一组'科目名', 指标列1, 指标列2,...的条目
• 用反引号包裹列名，避免列名包含特殊字符时出错
• 函数参数完全动态：id_vars是需要保留的非指标列，subject_list是传入的科目列表，metric_suffixes是每个科目对应的指标后缀，灵活适配不同的指标组合

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