使用PySpark读取S3文件时如何获取文件名并新增PO ID列

Spark 原生提供了内置函数input_file_name()可以直接在读取阶段获取当前行对应的源文件完整路径，全程在计算节点内存中完成，无额外IO开销，是性能最优的实现方式，不需要事后单独拉取文件名做关联。

具体实现步骤 #

1. 导入依赖函数

from pyspark.sql.functions import input_file_name, regexp_extract

2. 读取文件的同时提取PO ID
   你可以直接通过正则匹配从文件路径中提取文件名开头的6位PO ID，适配任意层级的S3路径：

df1 = spark.read.csv("s3://bucket/fact/*.dat") \
    .withColumn("po_id", regexp_extract(input_file_name(), r'/(\d{6})_[^/]+\.dat$', 1))

正则规则r'/(\d{6})_[^/]+\.dat$'会自动匹配路径末尾符合「6位数字_xxx.dat」格式的文件名，直接捕获前6位数字作为po_id的值。

可选优化 #

• 如果存在不符合命名规则的文件，可通过过滤逻辑剔除无效数据：

df1 = df1.filter(df1.po_id != "")

• 不需要保留源文件路径的话无需额外存储中间列，上述代码不会产生冗余存储开销。

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