如何在PySpark中从点分隔的字符串行提取指定文本？

输入数据 #

Submission ID
Dx.CS1.2023-21-01
DX.RE1.2023-12-01
DX.G1.2021-23-01
DX.G2.2022-10-12

期望输出 #

ID
CS1
RE1
G1
G2

解决方案 #

方法1：使用split函数分割字符串 #

利用字符串的分隔符特征（首行用空格分隔，其余行用.分隔），通过split函数结合when条件判断提取目标内容：

from pyspark.sql import SparkSession
from pyspark.sql.functions import split, when, col

spark = SparkSession.builder.appName("ExtractTargetText").getOrCreate()

# 构造输入DataFrame
data = [
    ("Submission ID",),
    ("Dx.CS1.2023-21-01",),
    ("DX.RE1.2023-12-01",),
    ("DX.G1.2021-23-01",),
    ("DX.G2.2022-10-12",)
]
df = spark.createDataFrame(data, ["input_str"])

# 提取目标文本
df = df.withColumn(
    "extracted_id",
    when(
        col("input_str").startswith("Submission"),
        split(col("input_str"), " ").getItem(1)
    ).otherwise(
        split(col("input_str"), "\\.").getItem(1)
    )
)

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

方法2：使用正则表达式提取 #

通过regexp_extract函数匹配通用规则，一次处理所有行：

from pyspark.sql import SparkSession
from pyspark.sql.functions import regexp_extract, col

spark = SparkSession.builder.appName("ExtractTargetText").getOrCreate()

# 构造输入DataFrame
data = [
    ("Submission ID",),
    ("Dx.CS1.2023-21-01",),
    ("DX.RE1.2023-12-01",),
    ("DX.G1.2021-23-01",),
    ("DX.G2.2022-10-12",)
]
df = spark.createDataFrame(data, ["input_str"])

# 正则表达式匹配：匹配首行空格后内容，或其余行两个点之间的内容
df = df.withColumn(
    "extracted_id",
    regexp_extract(col("input_str"), r"(?:Submission )?(\w+)(?:\.\d{4}-\d{2}-\d{2})?", 1)
)

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

正则说明：

• (?:Submission )?：非捕获组，匹配可选的前缀"Submission "
• (\w+)：捕获组，提取字母数字组合的目标内容
• (?:\.\d{4}-\d{2}-\d{2})?：非捕获组，匹配可选的日期后缀

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