使用PySpark展开嵌套JSON后如何去除父结构前缀获取原始列名

解决方案 #

方案1：未扁平化时一步生成目标结构（推荐）

直接通过结构体通配符选择子字段，无需生成带前缀的中间结果：

# 直接选择顶层字段 + foo结构体下的所有子字段
final_df = df.select("x", "y", "foo.*")

执行后打印schema即可验证结果符合预期：

final_df.printSchema()

方案2：已生成带前缀的DataFrame时批量重命名

如果已经完成扁平化得到了带foo_前缀的列，可以批量替换前缀：

from pyspark.sql.functions import col

prefix = "foo_"
# 遍历列名，匹配前缀的去除前缀，其他列保留原名
new_cols = [col(c).alias(c.removeprefix(prefix)) if c.startswith(prefix) else c for c in df.columns]
final_df = df.select(new_cols)

注：Python 3.9以下版本没有removeprefix方法，可以替换为c[len(prefix):]实现相同效果。

通用自动展开方案（适配未知结构的多结构体场景）

如果嵌套结构不固定、或者存在多个需要展开的结构体，可以通过遍历schema自动处理：

from pyspark.sql.types import StructType

select_list = []
for field in df.schema.fields:
    if isinstance(field.dataType, StructType):
        # 结构体类型展开所有子字段
        select_list += [f"{field.name}.{sub.name}" for sub in field.dataType.fields]
    else:
        # 普通类型直接保留原列
        select_list.append(field.name)

final_df = df.select(select_list)

注意：如果不同父结构体下存在同名子字段，直接展开会出现列名冲突，这种场景下需要保留父前缀做区分，不适用上述方法。

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