PySpark如何将DataFrame单元格内CSV值拆分为列且避免硬编码表头

问题背景 #

当前Spark DataFrame的某一列单元格中存储了CSV格式的值，需要将其拆分扩展为新列，示例输入DataFrame如下：

a_id                                    features
1   2020     "a","b","c","d","constant1","1","0.1","aa"
2   2021     "a","b","c","d","constant2","1","0.2","ab"
3   2022     "a","b","c","d","constant3","1","0.3","ac","constant3","1.1","3.3","acx"
4   2023     "a","b","c","d","constant4","1","0.4","ad"
5   2024     "a","b","c","d","constant5","1","0.5","ae","constant5","1.2","6.3","xwy","a","b","c","d","constant5","2.2","8.3","bunr"
6   2025     "a","b","c","d","constant6","1","0.6","af"

features列包含多组CSV值，其中"a","b","c","d"作为表头会在部分单元格中重复出现（如第3行、第5行），需要提取每组表头对应的取值，预期输出如下：

a_id       a        d
1   2020   constant1   ["aa"]
2   2021   constant2   ["ab"]
3   2022   constant3   ["ac","acx"]
4   2023   constant4   ["ad"]
5   2024   constant5   ["ae","xwy","bunr"]
6   2025   constant6   ["af"]

无硬编码实现方案 #

你可以通过配置化表头列表的方式，完全避免手动统计表头数量、硬编码表头字符串的问题，后续新增列仅需要修改配置项即可，无需调整核心逻辑：

from pyspark.sql import functions as F

# ========== 仅需修改这里的配置项即可适配表头变更/新增场景 ==========
# 定义表头列表，后续新增表头直接往列表中追加元素即可
HEADER_LIST = ['"a"', '"b"', '"c"', '"d"']
# 定义需要提取的字段映射，key为输出列名，value为该字段在每组数据中的索引位置
FIELD_MAPPING = {
    "a": 0, # 每组第一个元素对应a列
    "d": 3  # 每组第四个元素对应d列
}
# ================================================================

# 自动计算拼接后的表头字符串、每组数据长度，无需手动统计
header_sep = ','.join(HEADER_LIST) + ','
group_len = len(HEADER_LIST)

df = df.withColumn("split_by_header", F.split("features", header_sep)) \
       # 过滤拆分后的空值
       .withColumn("split_by_header", F.expr("filter(split_by_header, x -> trim(x) != '')")) \
       # 将每组数据拆分为独立行
       .withColumn("group_item", F.explode("split_by_header")) \
       # 拆分单组内的所有字段
       .withColumn("group_fields", F.split("group_item", ","))

# 自动按配置提取所需字段，无需硬编码索引
for col_name, idx in FIELD_MAPPING.items():
    df = df.withColumn(col_name, F.col("group_fields")[idx])

# 分组聚合得到最终结果
result_df = df.groupBy("a_id") \
              .agg(
                  F.first("a").alias("a"),
                  F.collect_list("d").alias("d")
              )

result_df.show(truncate=False)

方案优势 #

1. 无硬编码逻辑：表头数量、表头字符串全部从配置的HEADER_LIST自动计算，无需手动统计修改
2. 易扩展：后续如果新增表头，或者需要提取更多列，仅需要修改HEADER_LIST和FIELD_MAPPING两个配置项即可，核心拆分逻辑完全不需要调整
3. 适配性强：即使后续每组数据的长度随表头增加而变化，也会自动适配group_len的取值，不需要手动修改分组规则

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