Pyspark中如何实现类似R reshape函数的宽表列转置功能

核心机制对齐说明 #

• 关于R reshape 的空值处理逻辑：转宽表过程不会主动丢弃空值，所有属于idvar分组、匹配timevar取值的非idvar列值（包括NA空值），会被直接映射到对应后缀的宽表列位，不存在自动过滤空值关联分组的行为。
• PySpark生态无原生完全等价R reshape 的开箱函数，但通过内置pivot算子搭配自定义聚合规则，可100%复刻R reshape 的转置逻辑，无需依赖第三方扩展包。

分步骤实现流程 #

1. 预处理对齐R脚本前置逻辑
   • 删除源数据中不需要的Planned.Treatment.ID列
   • 逐列匹配R脚本的重命名规则：将Treatment.Administration.days.x重命名为Treatment.Administration.days.plan，Treatment.Administration.days.y重命名为Treatment.Administration.days.act，使用PySpark原生withColumnRenamed算子即可实现。
2. 明确转置核心参数
   • 分组列（对应R的idvar）：完全沿用R代码中指定的固定维度列清单，作为宽表的行级主键
   • 转置键（对应R的timevar）：使用Treatment.ID，所有该字段的去重值会自动作为宽表列名的后缀来源
   • 待转置值列：除分组列、转置键外的所有列，全部纳入转置范围，和R reshape 自动识别非idvar、非timevar列的逻辑保持一致
3. 执行转置操作
   • 提前拉取Treatment.ID的全部去重值，避免PySpark pivot的全表扫描开销，同时完全匹配R自动遍历所有timevar值展开列的行为
   • 聚合规则选用F.first(ignorenulls=False)，保证取值行为和R无聚合直取的逻辑完全一致：同分组、同timevar下的任意值（包括空值）都会被直接映射到对应列位，不会跳过空值
   • 调整生成列的命名规则：将PySpark默认生成的{timevar值}_{值列名}格式，调整为{值列名}_{timevar值}，和R生成的列名规则（值列名.timevar值）完全匹配

可直接复用的参考代码 #

from pyspark.sql import functions as F

# 预处理阶段
trt_clean = (
    trt_stacked
    .drop("Planned.Treatment.ID")
    .withColumnRenamed("Treatment.Administration.days.x", "Treatment.Administration.days.plan")
    .withColumnRenamed("Treatment.Administration.days.y", "Treatment.Administration.days.act")
)

# 定义转置核心参数
idvar_cols = [
    "Trial.ID", "Arm.ID", "Re.randomized.arm.id", "Phase.ID", "Period.ID",
    "Phase", "Phase.Duration", "Phase.Duration.unit", "Phase.Description",
    "Period", "Period.Duration", "Period.Dur.Unit", "Period.Description"
]
timevar_col = "Treatment.ID"
value_cols = [c for c in trt_clean.columns if c not in idvar_cols + [timevar_col]]

# 提前获取所有timevar去重值
pivot_values = [row[timevar_col] for row in trt_clean.select(timevar_col).distinct().collect()]

# 构造聚合表达式
agg_exprs = [F.first(F.col(c), ignorenulls=False) for c in value_cols]

# 执行转置
trt_wide = (
    trt_clean
    .groupBy(*idvar_cols)
    .pivot(timevar_col, pivot_values)
    .agg(*agg_exprs)
)

# 调整列名格式匹配R输出
for old_col in trt_wide.columns:
    if old_col not in idvar_cols:
        time_val, val_col = old_col.split("_", 1)
        trt_wide = trt_wide.withColumnRenamed(old_col, f"{val_col}_{time_val}")

示例数据匹配说明 #

上述实现逻辑跑通提供的示例输入时，输出结果将完全匹配预期宽表结构：

• 所有分组维度列会作为固定列前置
• Treatment.ID=1对应的所有非维度列值，会映射到后缀为_1的对应列
• Treatment.ID=2对应的所有非维度列值，会映射到后缀为_2的对应列
• 所有空值会被原样保留在对应列位，不会被转置流程过滤

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