PySpark/Pandas按自定义逻辑重塑DataFrame并调整列顺序

核心问题出在长宽表转换后没有按「层级优先、同层级按字段顺序」的规则自定义列排序，下面提供优先推荐的PySpark实现，以及修正列顺序问题的Pandas实现：

PySpark 实现（适配大数据量场景） #

实现逻辑：

• 先过滤LVL_NUM <= 0的无效层级数据
• 提前枚举有效层级值提升透视性能，按BTCH_NUM分组透视所有指标字段
• 按要求重命名列为原字段名.LVL_NUM格式，严格按「LVL_NUM从小到大、同层级按INP_A→INP_B→OUTP_A→OUTP_B」的规则重排列顺序
• 按BTCH_NUM分区写出，每个批次单独生成对应parquet文件

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, first

# 初始化Spark会话，可根据自身集群配置调整参数
spark = SparkSession.builder.appName("batch_l2w_reshape").getOrCreate()

# 替换为你的原始数据读取路径
raw_df = spark.read.parquet("hdfs:///path/to/your/raw_data.parquet")

# 过滤有效层级
valid_df = raw_df.filter(col("LVL_NUM") > 0)

# 提前获取升序排列的所有有效层级值，避免pivot全量扫描枚举
level_list = [i.LVL_NUM for i in valid_df.select("LVL_NUM").distinct().orderBy("LVL_NUM").collect()]
metric_cols = ["INP_A", "INP_B", "OUTP_A", "OUTP_B"]

# 透视转宽表
pivot_df = valid_df.groupBy("BTCH_NUM").pivot("LVL_NUM", level_list).agg(
    *[first(c).alias(c) for c in metric_cols]
)

# 重命名列为指定格式
for lvl in level_list:
    for metric in metric_cols:
        pivot_df = pivot_df.withColumnRenamed(f"{lvl}_{metric}", f"{metric}.{lvl}")

# 按规则重排列顺序
sorted_columns = ["BTCH_NUM"]
for lvl in level_list:
    for metric in metric_cols:
        sorted_columns.append(f"{metric}.{lvl}")
final_df = pivot_df.select(*sorted_columns)

# 按BTCH_NUM分区写出parquet，每个分区对应一个批次的文件
final_df.write.partitionBy("BTCH_NUM").mode("overwrite").parquet("hdfs:///path/to/output/")

Pandas 实现（适配小数据量场景） #

针对现有Pandas逻辑的问题，只需要在透视完成后增加自定义列排序步骤即可，不需要改动核心转换逻辑：

import pandas as pd

# 替换为你的原始数据读取逻辑
raw_df = pd.read_parquet("./raw_data.parquet")

# 过滤有效层级
valid_df = raw_df[raw_df["LVL_NUM"] > 0].copy()

# 长宽转换
pivot_df = valid_df.pivot(
    index="BTCH_NUM",
    columns="LVL_NUM",
    values=["INP_A", "INP_B", "OUTP_A", "OUTP_B"]
).reset_index()

# 重命名列为指定格式
pivot_df.columns = [
    f"{col[0]}.{col[1]}" if col[1] != "" else col[0] 
    for col in pivot_df.columns
]

# 按规则重排列顺序
level_list = sorted(valid_df["LVL_NUM"].unique())
metric_order = ["INP_A", "INP_B", "OUTP_A", "OUTP_B"]
sorted_columns = ["BTCH_NUM"]
for lvl in level_list:
    for metric in metric_order:
        sorted_columns.append(f"{metric}.{lvl}")
final_df = pivot_df[sorted_columns]

# 按批次拆分保存为单独parquet文件
for btch_id, batch_data in final_df.groupby("BTCH_NUM"):
    batch_data.to_parquet(f"./output/reshape_BTCH_{btch_id}.parquet", index=False)

注意：如果同一BTCH_NUM + LVL_NUM组合存在多条重复记录，请将上述代码中的first（PySpark）聚合逻辑替换为符合业务要求的聚合规则（如max、sum、collect_list等），避免数据遗漏。

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