如何在PySpark DataFrame中提取列中字符串生成列表新列？

问题描述 #

现有一个仅含单列的PySpark DataFrame，示例数据如下：

+------------------------------------------------------------------------------------------------------+
|data                                                                                                  |
+------------------------------------------------------------------------------------------------------+
{Job : "Job A", Condition: "Pass"},{Job : "Job B", Condition: "Pass"}
{Job : "Job C", Condition: "Pass"},{Job : "Job D", Condition: "Pass"},{Job : "Job E", Condition: "Pass"}
+------------------------------------------------------------------------------------------------------+

希望提取其中的Job名称，生成一个列表形式的新列，期望输出如下：

+-----------------------+
|data1                  |
+-----------------------+
["Job A","Job B"]
["Job C","Job D","Job E"]
+-----------------------+

简便实现方法 #

可以通过字符串正则提取+数组转换直接完成，核心利用Spark的regexp_extract_all函数（Spark 3.0及以上版本支持），一步生成目标数组列。

示例代码 #

from pyspark.sql import functions as F

# 假设原DataFrame名为df
result_df = df.withColumn(
    "data1",
    # 补全首尾大括号，确保正则能匹配所有Job项，然后提取所有Job名称
    F.regexp_extract_all(F.concat(F.lit("{"), F.col("data"), F.lit("}")), r'Job\s*:\s*"([^"]+)"', 1)
)

# 查看结果
result_df.select("data1").show(truncate=False)

关键说明 #

• concat(F.lit("{"), F.col("data"), F.lit("}"))：将每行的字符串补全为{...,...}的结构，让正则表达式可以统一匹配所有Job条目，避免首尾的条目漏匹配。
• regexp_extract_all：一次性提取所有符合正则规则的子串，直接返回数组类型，完美契合生成列表新列的需求。
• 若使用Spark 2.x版本（无regexp_extract_all），可先拆分片段再聚合，代码示例如下：

from pyspark.sql import functions as F

result_df = df.withColumn("job_item", F.split(F.col("data"), r"\s*},\s*")) \
    .withColumn("job_item", F.explode(F.col("job_item"))) \
    .withColumn("job_name", F.regexp_extract(F.col("job_item"), r'Job\s*:\s*"([^"]+)"', 1)) \
    .groupBy("data") \
    .agg(F.collect_list("job_name").alias("data1")) \
    .select("data1")

result_df.show(truncate=False)

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