PySpark实现：提取forenames列中排除first_name的剩余内容

我们有一个包含forenames和first_name列的PySpark DataFrame，其中first_name是forenames中第一个空格前的完整字符串（含连字符的情况会完整保留）。现在需要新增middle_name列，用来存储forenames中first_name之后的剩余内容。

原始数据 #

|forenames        | first_name |
+-----------------+------------+
|IVO KAI ROGERS   |  IVO       |
|DYLAN STUART JOHN|  DYLAN     |
|JOSH JACK        |  JOSH      |
|MONALISA ELIEN   |  MONALISA  |
|RACHEL-GREEN JOE |RACHEL-GREEN|

预期输出 #

|forenames        | first_name | middle_name |
+-----------------+------------+-------------+
|IVO KAI ROGERS   |  IVO       | KAI ROGERS  |
|DYLAN STUART JOHN|  DYLAN     | STUART JOHN |
|JOSH JACK        |  JOSH      | JACK        |
|MONALISA ELIEN   |  MONALISA  | ELIEN       |
|RACHEL-GREEN JOE |RACHEL-GREEN| JOE         |

解决方案 #

可以通过regexp_replace结合trim函数实现，同时处理forenames与first_name完全相同的边界情况：

from pyspark.sql import SparkSession
from pyspark.sql.functions import regexp_replace, trim, col, when

# 初始化SparkSession
spark = SparkSession.builder.appName("ExtractMiddleName").getOrCreate()

# 创建示例DataFrame
data = [
    ("IVO KAI ROGERS", "  IVO       "),
    ("DYLAN STUART JOHN", "  DYLAN     "),
    ("JOSH JACK", "  JOSH      "),
    ("MONALISA ELIEN", "  MONALISA  "),
    ("RACHEL-GREEN JOE", "RACHEL-GREEN")
]
df = spark.createDataFrame(data, ["forenames", "first_name"])

# 新增middle_name列
df = df.withColumn(
    "trimmed_first_name",
    trim(col("first_name"))
).withColumn(
    "middle_name",
    when(
        col("forenames") == col("trimmed_first_name"),
        ""
    ).otherwise(
        regexp_replace(col("forenames"), "^" + col("trimmed_first_name") + "\\s+", "")
    )
).drop("trimmed_first_name")

# 查看结果
df.show(truncate=False)

代码说明 #

1. trim(col("first_name"))：去除first_name前后的多余空格，确保匹配的是准确的名字主体，避免因空格导致匹配失败。
2. 正则表达式^...\\s+：
   • ^ 匹配字符串开头，确保只替换forenames最前面的first_name部分
   • \\s+ 匹配一个或多个空格，覆盖first_name后可能存在的多个空格场景
3. when条件判断：处理forenames与trimmed_first_name完全相同的情况（即没有中间名），此时将middle_name设为空字符串。

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