PySpark实现对以‘AB-’开头的字符串截取处理

需求说明 #

现有包含字符串列Column A的PySpark DataFrame，需新增列Column B：

• 若Column A的字符串以AB-开头，移除该前缀保留剩余内容
• 否则保持原字符串不变

示例数据 #

预期输出 #

解决方案 #

方法1：条件判断+字符串截取

通过when函数判断前缀，结合substring截取从第4位开始的字符（AB-长度为3，PySpark字符串索引从1开始）：

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

# 创建示例DataFrame
spark = SparkSession.builder.appName("PrefixRemove").getOrCreate()
data = [
    ("AB-001-1-12345-A",),
    ("AB-001-1-12346-B",),
    ("ABC012345B",),
    ("ABC012346B",)
]
df = spark.createDataFrame(data, ["Column A"])

# 生成Column B
df = df.withColumn(
    "Column B",
    when(col("Column A").startswith("AB-"), substring(col("Column A"), 4, 100))
    .otherwise(col("Column A"))
)

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

方法2：正则表达式替换

用regexp_replace匹配开头的AB-并替换为空，写法更简洁：

from pyspark.sql.functions import regexp_replace

# 生成Column B
df = df.withColumn(
    "Column B",
    regexp_replace(col("Column A"), "^AB-", "")
)

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

两种方法均可得到预期结果：正则方法更适配复杂前缀规则，条件判断方法则更直观易读。

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