如何在PySpark中从DataFrame的文件名列提取对应文件ID

针对大型PySpark DataFrame场景，优先使用原生内置函数实现，避免自定义UDF的额外性能开销，具体实现如下：

• 导入内置正则提取函数

from pyspark.sql.functions import regexp_extract

• 新增id列提取目标编码

# 正则匹配规则：捕获文件名末尾紧邻_PO.dat的6位数字，避免路径中数字干扰
df = df.withColumn("id", regexp_extract("filename", r"(\d{6})_PO\.dat$", 1))

• 验证输出结果

df.show(truncate=False)

性能说明 #

regexp_extract是Spark原生向量化函数，直接在JVM执行层处理数据，完全规避了Python UDF带来的跨进程序列化/反序列化损耗，处理TB级大型数据集时性能比自定义UDF高5~10倍。

如果你的文件名规则存在变化，可灵活调整正则表达式：例如当6位数字是文件名中唯一的连续6位数字时，可简化正则为r"(\d{6})"适配更多场景。

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